wip; switching to grad-buffer-centric design

megatron/arguments.py

@@ -130,9 +130,11 @@ def parse_args(extra_args_provider=None, defaults={},
                 args.global_batch_size), flush=True)
     assert args.global_batch_size > 0
     if args.num_layers_per_virtual_pipeline_stage is not None:
-        assert args.pipeline_model_parallel_size > 2, \
-            'pipeline-model-parallel size should be greater than 2 with ' \
-            'interleaved schedule'
+        # >>> [ temporarily turning off ]
+        # assert args.pipeline_model_parallel_size > 2, \
+        #     'pipeline-model-parallel size should be greater than 2 with ' \
+        #     'interleaved schedule'
+        # <<<
         assert args.num_layers % args.num_layers_per_virtual_pipeline_stage == 0, \
             'number of layers is not divisible by number of layers per virtual ' \
             'pipeline stage'

megatron/optimizer/__init__.py

@@ -97,11 +97,11 @@ def get_megatron_optimizer(model,
     # from lutil import pax
     # pax(0, {
     #     "model" : model,
-    #     "param_groups" : param_groups,
-    #     "param_groups / 0" : param_groups[0],
-    #     "param_groups / 0 / params" : param_groups[0]["params"],
-    #     "param_groups / 1" : param_groups[1],
-    #     "param_groups / 1 / params" : param_groups[1]["params"],
+    #     # "param_groups" : param_groups,
+    #     # "param_groups / 0" : param_groups[0],
+    #     # "param_groups / 0 / params" : param_groups[0]["params"],
+    #     # "param_groups / 1" : param_groups[1],
+    #     # "param_groups / 1 / params" : param_groups[1]["params"],
     # })
     # <<<
 
@@ -164,7 +164,8 @@ def get_megatron_optimizer(model,
                       params_have_main_grad,
                       args.use_contiguous_buffers_in_local_ddp,
                       args.bf16,
-                      grad_scaler)
+                      grad_scaler,
+                      model)
         # <<<
 
     # FP32.

megatron/training.py

@@ -365,8 +365,12 @@ def setup_model_and_optimizer(model_provider_func,
 
     unwrapped_model = unwrap_model(model,
                                    (torchDDP, LocalDDP, Float16Module))
-    optimizer = get_megatron_optimizer(unwrapped_model, no_wd_decay_cond,
+    # >>>
+    # optimizer = get_megatron_optimizer(unwrapped_model, no_wd_decay_cond,
+    #                                    scale_lr_cond, lr_mult)
+    optimizer = get_megatron_optimizer(model, no_wd_decay_cond,
                                        scale_lr_cond, lr_mult)
+    # <<<
 
     opt_param_scheduler = get_optimizer_param_scheduler(optimizer)