From 6fa43a279129ccd358e8b627ec9f744dc2128050 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 14 Nov 2023 16:48:12 -0800 Subject: [PATCH] add cpu init check Signed-off-by: Chen Cui --- .../nlp/models/language_modeling/megatron_gpt_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py index d791505f3fe4..0624e373bca3 100644 --- a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py +++ b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py @@ -245,7 +245,7 @@ def __init__(self, cfg: DictConfig, trainer: Trainer): if self.megatron_amp_O2: - if not self.with_distributed_adam: + if not self.with_distributed_adam and not self.cfg.get("use_cpu_initialization", False): # Pre-allocate the model on GPU to have master parameters allocated on the same device with matching data type if isinstance(self.model, list): for module in self.model: