Add option to merge distributed optimizer buckets (NVIDIA#9414)

* Add option to merge distopt buckets in GPT Signed-off-by: Tim Moon <[email protected]> * Move distopt bucket merge logic to base LLM class Signed-off-by: Tim Moon <[email protected]> * Apply isort and black reformatting Signed-off-by: timmoon10 <[email protected]> --------- Signed-off-by: Tim Moon <[email protected]> Signed-off-by: timmoon10 <[email protected]> Co-authored-by: timmoon10 <[email protected]> Co-authored-by: Sangkug Lym <[email protected]>
paarthneekhara · Jan 18, 2025 · 6288e76 · 6288e76
1 parent b40cb63
commit 6288e76
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_base_model.py b/nemo/collections/nlp/models/language_modeling/megatron_base_model.py
@@ -861,7 +861,15 @@ def configure_optimizers(self):
 
             # Initialize param buckets if explicitly provided
             if getattr(self, 'distributed_adam_buckets', None) is not None:
-                for bucket in self.distributed_adam_buckets:
+                buckets = self.distributed_adam_buckets
+                if self.cfg.get('distributed_adam_bucket_merge_size', 1) > 1:
+                    # Merge buckets if needed
+                    stride = self.cfg.get('distributed_adam_bucket_merge_size', 1)
+                    buckets = [
+                        list(itertools.chain.from_iterable(buckets[i : i + stride]))
+                        for i in range(0, len(buckets), stride)
+                    ]
+                for bucket in buckets:
                     self._optimizer.init_params_bucket(bucket)
                 self._optimizer.init_params_bucket(self.parameters())
             if hasattr(self, 'distributed_adam_buckets'):