add the mcore interface for optim arg; average_in_collective (#10010)

erhoo82 · web-flow · commit 6cf59fab6941 · 2024-08-07T10:11:50.000-07:00
Signed-off-by: Sangkug Lym &lt;slym@nvidia.com&gt;
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py
@@ -545,6 +545,7 @@ def setup_mcore_distributed_parallel(self):
                 # mcore bucket_size is based on num of parameters, therefore not
                 # using bucket_cap_mb to configure bucket_size here
                 bucket_size=self.cfg.optim.get('ddp_bucket_size', None),
+                average_in_collective=self.cfg.optim.get('average_in_collective', True),
             )
             self.model = [
                 McoreDDP(

Original file line number	Diff line number	Diff line change
`@@ -545,6 +545,7 @@ def setup_mcore_distributed_parallel(self):`
`545`	`545`	`# mcore bucket_size is based on num of parameters, therefore not`
`546`	`546`	`# using bucket_cap_mb to configure bucket_size here`
`547`	`547`	`bucket_size=self.cfg.optim.get('ddp_bucket_size', None),`
	`548`	`+ average_in_collective=self.cfg.optim.get('average_in_collective', True),`
`548`	`549`	`)`
`549`	`550`	`self.model = [`
`550`	`551`	`McoreDDP(`