1.5T with more CUDA memory for plan A

Author: feifeibear

tests/test_bert.py

@@ -203,7 +203,7 @@ def test_bert_model(is_ckp: bool = False,
 
     if is_ps:
         manager = HybridPSManager()
-        manager.init([1024 * 1024 * 1024 * 2] * 1,
+        manager.init([1024 * 1024 * 1024 * 8] * 1,
                      [1024 * 1024 * 1024 * 4 * 4])
         # chunk 512 MB, good for CPU-GPU bandwidth
         client = HybridPSClient(gpu_index=0,