cache num fw args in train and eval

Signed-off-by: Maanu Grover <[email protected]>

Author: maanug-nv

nemo/tron/api.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import inspect
 from typing import Callable
 
 from nemo.tron.checkpointing import save_checkpoint
@@ -40,15 +39,6 @@ def megatron_pretrain(
     test_data_iterator = setup_output.test_data_iterator
     ckpt_context = setup_output.checkpointing_context
 
-    # Check num args to forward_step_func
-    num_fw_args = len(inspect.signature(forward_step_func).parameters)
-    fail_msg = f"""
-    forward_step_func has {num_fw_args} arguments. Only the following signatures are supported: 
-        2 args: forward_step_func(data_iterator: Iterable, model: GPTModel)
-        3 args: forward_step_func(state: GlobalState, data_iterator: Iterable, model: GPTModel)
-    """
-    assert num_fw_args in (2, 3), fail_msg
-
     ## TRAINING ##
     if not config.train_config.skip_train:
         print_rank_0("training ...")

nemo/tron/eval.py

@@ -25,7 +25,7 @@
 from nemo.tron import fault_tolerance
 from nemo.tron.state import GlobalState
 from nemo.tron.utils.common_utils import is_last_rank, print_rank_0, print_rank_last
-from nemo.tron.utils.train_utils import maybe_inject_state
+from nemo.tron.utils.train_utils import check_forward_step_func_num_args, maybe_inject_state
 
 
 def evaluate(
@@ -39,8 +39,10 @@ def evaluate(
     non_loss_data_func=None,
 ):
     """Evaluation."""
-    timers = state.timers
+    # Check num args to forward_step_func
+    num_fw_args = check_forward_step_func_num_args(forward_step_func)
 
+    timers = state.timers
     timers("evaluate", log_level=0).start(barrier=True)
 
     # Turn on evaluation mode which disables dropout.
@@ -67,7 +69,7 @@ def evaluate(
             if verbose:
                 print_rank_0(f"Evaluating iter {iteration}/{state.cfg.train_config.eval_iters}")
 
-            wrapped_forward_step = maybe_inject_state(forward_step_func, state)
+            wrapped_forward_step = maybe_inject_state(forward_step_func, state, num_fw_args=num_fw_args)
             forward_backward_func = get_forward_backward_func()
             # Don't care about timing during evaluation
             config.timers = None

nemo/tron/train.py

@@ -43,6 +43,7 @@
 from nemo.tron.utils.common_utils import append_to_progress_log, barrier_and_log, get_world_size_safe, print_rank_0
 from nemo.tron.utils.train_utils import (
     calc_params_l2_norm,
+    check_forward_step_func_num_args,
     logical_and_across_model_parallel_group,
     maybe_inject_state,
     reduce_max_stat_across_model_parallel_group,
@@ -68,6 +69,9 @@ def train(
     timers = global_state.timers
     straggler_timer = global_state.straggler_timer
 
+    # Check num args to forward_step_func
+    num_fw_args = check_forward_step_func_num_args(forward_step_func)
+
     # Turn on training mode which enables dropout.
     for model_module in model:
         model_module.train()
@@ -231,7 +235,7 @@ def train(
         # Run training step.
         fault_tolerance.on_training_step_start(global_state)
         loss_dict, skipped_iter, should_checkpoint, should_exit, exit_code, grad_norm, num_zeros_in_grad = train_step(
-            forward_step_func, train_data_iterator, model, optimizer, scheduler, global_state
+            forward_step_func, num_fw_args, train_data_iterator, model, optimizer, scheduler, global_state
         )
         fault_tolerance.on_training_step_end(global_state)
         if should_checkpoint:
@@ -403,6 +407,7 @@ def train(
 
 def train_step(
     forward_step_func,
+    num_fw_args,
     data_iterator,
     model,
     optimizer,
@@ -424,7 +429,7 @@ def train_step(
         optimizer.zero_grad()
 
         # Optionally inject state into forward step
-        wrapped_forward_step = maybe_inject_state(forward_step_func, global_state)
+        wrapped_forward_step = maybe_inject_state(forward_step_func, global_state, num_fw_args=num_fw_args)
 
         # Forward pass.
         forward_backward_func = get_forward_backward_func()

nemo/tron/utils/train_utils.py

@@ -15,7 +15,7 @@
 import inspect
 from datetime import datetime
 from functools import partial
-from typing import Callable
+from typing import Callable, Optional
 
 import torch
 from megatron.core import parallel_state
@@ -514,10 +514,23 @@ def reduce_aux_losses_tracker_across_ranks():
             torch.distributed.all_reduce(values, group=tracker[name]["avg_group"], op=torch.distributed.ReduceOp.AVG)
 
 
-def maybe_inject_state(forward_step_func: Callable, state: GlobalState) -> Callable:
-    num_fw_args = len(inspect.signature(forward_step_func).parameters)
+def maybe_inject_state(forward_step_func: Callable, state: GlobalState, num_fw_args: Optional[int] = None) -> Callable:
+    if not num_fw_args:
+        num_fw_args = len(inspect.signature(forward_step_func).parameters)
     if num_fw_args == 3:
         # inject global_state
         return partial(forward_step_func, state)
     else:
         return forward_step_func
+
+
+def check_forward_step_func_num_args(forward_step_func: Callable) -> int:
+    num_fw_args = len(inspect.signature(forward_step_func).parameters)
+    fail_msg = f"""
+    forward_step_func has {num_fw_args} arguments. Only the following signatures are supported: 
+        2 args: forward_step_func(data_iterator: Iterable, model: GPTModel)
+        3 args: forward_step_func(state: GlobalState, data_iterator: Iterable, model: GPTModel)
+    """
+    assert num_fw_args in (2, 3), fail_msg
+
+    return num_fw_args