test cpu offload with pp

pytorch · Oct 28, 2024 · 2ca9882 · 2ca9882
1 parent ab1e258
commit 2ca9882
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/test_runner.py b/test_runner.py
@@ -354,12 +354,13 @@ def build_test_list():
         OverrideDefinitions(
             [
                 [
+                    "--experimental.pipeline_parallel_degree 2",
                     "--training.enable_cpu_offload True",
                 ],
             ],
-            "Enable CPU Offload",
-            "enable_cpu_offload",
-            ngpu=2,
+            "Enable CPU Offload with PP",
+            "enable_cpu_offload+PP",
+            ngpu=4,
         ),
     ]
     return integration_tests_flavors

diff --git a/train.py b/train.py
@@ -151,8 +151,10 @@ def loss_fn(pred, labels):
         for m in model_parts:
             # apply SPMD-style PT-D techniques
             models_parallelize_fns[model_name](m, world_mesh, parallel_dims, job_config)
-            m.to_empty(device="cuda")
-            m.init_weights()
+            init_device = "cpu" if job_config.training.enable_cpu_offload else "cuda"
+            m.to_empty(device=init_device)
+            buffer_device = "cuda" if job_config.training.enable_cpu_offload else None
+            m.init_weights(buffer_device=buffer_device)
             m.train()
     else:
         # apply PT-D Tensor Parallel, activation checkpointing, torch.compile, Data Parallel