Skip to content

GitLab

Commit 3d96729a authored by Mostofa Patwary's avatar Mostofa Patwary
Browse files

fixed an issue with ICT

parent 816fb890
Show whitespace changes
Inline Side-by-side
Showing with 7 additions and 8 deletions
megatron/training.py
View file @ 3d96729a
...@@ -317,12 +317,11 @@ def setup_model_and_optimizer(model_provider_func): ...@@ -317,12 +317,11 @@ def setup_model_and_optimizer(model_provider_func):
assert args.DDP_impl == 'local' assert args.DDP_impl == 'local'
# get model without FP16 and/or TorchDDP wrappers # get model without FP16 and/or TorchDDP wrappers
model = unwrap_model(model) unwrapped_model = unwrap_model(model)
for module in model: if args.iteration == 0 and hasattr(unwrapped_model,
if args.iteration == 0 and hasattr(module,
'init_state_dict_from_bert'): 'init_state_dict_from_bert'):
print("Initializing ICT from pretrained BERT model", flush=True) print_rank_0("Initializing ICT from pretrained BERT model")
module.init_state_dict_from_bert() unwrapped_model.init_state_dict_from_bert()
if args.fp16: if args.fp16:
optimizer.reload_model_params() optimizer.reload_model_params()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment