wip

megatron/checkpointing.py

@@ -96,7 +96,9 @@ def do_forward_pass(neox_args, model, inference=False):
         tokens, attention_mask, position_ids = get_batch(
             neox_args, context_tokens_tensor[:, : neox_args.seq_length]
         )
-        logits = model((tokens, position_ids, attention_mask))
+        output = model((tokens, position_ids, attention_mask))
+        logits = output[0] if isinstance(output, tuple) else output
+
 
     # reset to train mode, if model was in training before
     if model_was_in_train:

tests/common.py

@@ -512,7 +512,6 @@ def model_setup(yaml_list=None, param_dict=None, clear_data=True):
     args_loaded.build_tokenizer()
 
     initialize_megatron(neox_args=args_loaded)
-    print("YAP")
     model, optimizer, lr_scheduler, reference_model = setup_model_and_optimizer(
         neox_args=args_loaded, use_cache=True
     )

tests/model/test_model_checkpoint.py

@@ -93,6 +93,7 @@ def test_checkpoint(self, param_dict, tmpdir):
             reloaded_model,
             reloaded_optimizer,
             reloaded_lr_scheduler,
+            reloaded_reference_model,
             args_reloaded,
         ) = model_setup(yaml_list=None, param_dict=param_dict, clear_data=False)
         iteration = load_checkpoint(