From 330f035c6d496ecfadf9379eda95fb71be86ee15 Mon Sep 17 00:00:00 2001 From: Dmytro Pykhtar Date: Wed, 30 Sep 2026 15:53:29 +0300 Subject: [PATCH] reuse train state from MLM Signed-off-by: Dmytro Pykhtar --- docs/bridge-rl-integration.md | 4 +- .../nightly/pages/bridge-rl-integration.mdx | 4 +- .../nightly/pages/training/callbacks.mdx | 10 +- docs/training/callbacks.md | 10 +- examples/models/bagel/pretrain_bagel.py | 2 +- .../pretrain_qwen3_eval_cp.py | 4 +- skills/nemo-mbridge-perf-cuda-graphs/SKILL.md | 2 +- .../bridge/models/bagel/bagel_step.py | 2 +- src/megatron/bridge/training/checkpointing.py | 39 ++-- src/megatron/bridge/training/eval.py | 18 +- src/megatron/bridge/training/pretrain.py | 2 +- src/megatron/bridge/training/setup.py | 23 +- .../bridge/training/setup_megatron_mimo.py | 2 +- src/megatron/bridge/training/state.py | 64 +----- src/megatron/bridge/training/train.py | 96 ++++---- .../bridge/training/train_megatron_mimo.py | 26 ++- .../bridge/training/utils/checkpoint_utils.py | 72 +----- .../bridge/training/utils/train_utils.py | 16 +- .../test_groups/data/test_loaders.py | 2 +- .../test_pretrain_megatron_mimo.py | 10 +- .../quantization/test_qat_workflow.py | 2 +- .../training/test_local_checkpointing.py | 4 +- tests/functional_tests/utils.py | 2 +- tests/unit_tests/conftest.py | 4 +- tests/unit_tests/models/qwen_omni/conftest.py | 4 +- .../test_megatron_mimo_checkpointing.py | 26 +-- .../test_pretrain_megatron_mimo.py | 2 +- .../post_training/test_checkpointing.py | 2 +- .../test_checkpoint_resume_correctness.py | 2 +- .../unit_tests/training/test_checkpointing.py | 90 ++++---- tests/unit_tests/training/test_eval.py | 2 +- .../training/test_functor_support.py | 4 +- .../training/test_peft_checkpointing.py | 12 +- .../training/test_pg_collection_wiring.py | 18 +- .../training/test_runtime_callbacks.py | 2 +- tests/unit_tests/training/test_setup.py | 2 +- tests/unit_tests/training/test_state.py | 164 +------------- tests/unit_tests/training/test_train.py | 14 +- .../training/utils/test_checkpoint_utils.py | 205 +----------------- .../training/utils/test_train_utils.py | 66 +++--- 40 files changed, 262 insertions(+), 773 deletions(-) diff --git a/docs/bridge-rl-integration.md b/docs/bridge-rl-integration.md index 0461e7bb10..75dbba0666 100644 --- a/docs/bridge-rl-integration.md +++ b/docs/bridge-rl-integration.md @@ -322,7 +322,7 @@ save_checkpoint( model=[model], optimizer=optimizer, opt_param_scheduler=scheduler, - num_floating_point_operations_so_far=state.train_state.floating_point_operations_so_far, + num_floating_point_operations_so_far=state.train_state.num_floating_point_operations_so_far, checkpointing_context=ckpt_ctx, ) ``` @@ -503,7 +503,7 @@ class MegatronBridgeAdapter: def save_ckpt(self, path: str): save_checkpoint(self.state, [self.model], self.optimizer, self.scheduler, - num_floating_point_operations_so_far=self.state.train_state.floating_point_operations_so_far, + num_floating_point_operations_so_far=self.state.train_state.num_floating_point_operations_so_far, checkpointing_context=self.ckpt_ctx) def export_hf(self, out_dir: str, trust_remote_code: bool = False): diff --git a/docs/fern/versions/nightly/pages/bridge-rl-integration.mdx b/docs/fern/versions/nightly/pages/bridge-rl-integration.mdx index 8f3fe70591..f689b0f428 100644 --- a/docs/fern/versions/nightly/pages/bridge-rl-integration.mdx +++ b/docs/fern/versions/nightly/pages/bridge-rl-integration.mdx @@ -315,7 +315,7 @@ save_checkpoint( model=[model], optimizer=optimizer, opt_param_scheduler=scheduler, - num_floating_point_operations_so_far=state.train_state.floating_point_operations_so_far, + num_floating_point_operations_so_far=state.train_state.num_floating_point_operations_so_far, checkpointing_context=ckpt_ctx, ) ``` @@ -494,7 +494,7 @@ class MegatronBridgeAdapter: def save_ckpt(self, path: str): save_checkpoint(self.state, [self.model], self.optimizer, self.scheduler, - num_floating_point_operations_so_far=self.state.train_state.floating_point_operations_so_far, + num_floating_point_operations_so_far=self.state.train_state.num_floating_point_operations_so_far, checkpointing_context=self.ckpt_ctx) def export_hf(self, out_dir: str, trust_remote_code: bool = False): diff --git a/docs/fern/versions/nightly/pages/training/callbacks.mdx b/docs/fern/versions/nightly/pages/training/callbacks.mdx index 5232491ad3..9f1f720a66 100644 --- a/docs/fern/versions/nightly/pages/training/callbacks.mdx +++ b/docs/fern/versions/nightly/pages/training/callbacks.mdx @@ -19,11 +19,11 @@ from megatron.bridge.recipes.qwen import qwen25_500m_pretrain_config class MyCallback(Callback): def on_train_start(self, context): context.user_state['start_time'] = time.time() - print(f"Training started at step {context.state.train_state.step}") + print(f"Training started at step {context.state.train_state.iteration}") def on_train_step_end(self, context): if context.loss_dict: - print(f"Step {context.state.train_state.step}: loss={context.loss_dict}") + print(f"Step {context.state.train_state.iteration}: loss={context.loss_dict}") def on_train_end(self, context): elapsed = time.time() - context.user_state['start_time'] @@ -47,7 +47,7 @@ from megatron.bridge.training.pretrain import pretrain from megatron.bridge.recipes.qwen import qwen25_500m_pretrain_config def log_step(context): - step = context.state.train_state.step + step = context.state.train_state.iteration if context.loss_dict: print(f"Step {step}: {context.loss_dict}") @@ -163,7 +163,7 @@ import torch.distributed as dist class RankZeroCallback(Callback): def on_train_step_end(self, context): if dist.get_rank() == 0: - print(f"Step {context.state.train_state.step} complete") + print(f"Step {context.state.train_state.iteration} complete") ``` ## Exception Handling @@ -232,7 +232,7 @@ class ProprietaryMetricsCallback(Callback): def on_train_step_end(self, context): if context.loss_dict: self.client.send({ - "step": context.state.train_state.step, + "step": context.state.train_state.iteration, "loss": context.loss_dict.get("lm loss"), "grad_norm": context.grad_norm, "cluster_id": os.environ.get("CLUSTER_ID"), diff --git a/docs/training/callbacks.md b/docs/training/callbacks.md index 15d619247e..8e73a5d2b8 100644 --- a/docs/training/callbacks.md +++ b/docs/training/callbacks.md @@ -19,11 +19,11 @@ from megatron.bridge.recipes.qwen import qwen25_500m_pretrain_config class MyCallback(Callback): def on_train_start(self, context): context.user_state['start_time'] = time.time() - print(f"Training started at step {context.state.train_state.step}") + print(f"Training started at step {context.state.train_state.iteration}") def on_train_step_end(self, context): if context.loss_dict: - print(f"Step {context.state.train_state.step}: loss={context.loss_dict}") + print(f"Step {context.state.train_state.iteration}: loss={context.loss_dict}") def on_train_end(self, context): elapsed = time.time() - context.user_state['start_time'] @@ -47,7 +47,7 @@ from megatron.bridge.training.pretrain import pretrain from megatron.bridge.recipes.qwen import qwen25_500m_pretrain_config def log_step(context): - step = context.state.train_state.step + step = context.state.train_state.iteration if context.loss_dict: print(f"Step {step}: {context.loss_dict}") @@ -164,7 +164,7 @@ import torch.distributed as dist class RankZeroCallback(Callback): def on_train_step_end(self, context): if dist.get_rank() == 0: - print(f"Step {context.state.train_state.step} complete") + print(f"Step {context.state.train_state.iteration} complete") ``` ## Exception Handling @@ -233,7 +233,7 @@ class ProprietaryMetricsCallback(Callback): def on_train_step_end(self, context): if context.loss_dict: self.client.send({ - "step": context.state.train_state.step, + "step": context.state.train_state.iteration, "loss": context.loss_dict.get("lm loss"), "grad_norm": context.grad_norm, "cluster_id": os.environ.get("CLUSTER_ID"), diff --git a/examples/models/bagel/pretrain_bagel.py b/examples/models/bagel/pretrain_bagel.py index 0948feec2a..1eb5633cbc 100644 --- a/examples/models/bagel/pretrain_bagel.py +++ b/examples/models/bagel/pretrain_bagel.py @@ -65,7 +65,7 @@ def _record_loss(output: Path, context: CallbackContext) -> None: if learning_rate is None: raise RuntimeError("BAGEL optimizer has no canonical learning rate") row = { - "step": context.state.train_state.step, + "step": context.state.train_state.iteration, "grad_norm": context.grad_norm, "lr": float(learning_rate), "losses": {name: value.detach().float().cpu().tolist() for name, value in sorted(context.loss_dict.items())}, diff --git a/examples/training_features/decentralized_pg/pretrain_qwen3_eval_cp.py b/examples/training_features/decentralized_pg/pretrain_qwen3_eval_cp.py index 0813a5e288..d75811d228 100644 --- a/examples/training_features/decentralized_pg/pretrain_qwen3_eval_cp.py +++ b/examples/training_features/decentralized_pg/pretrain_qwen3_eval_cp.py @@ -738,9 +738,9 @@ def _stats(xs: list[float]) -> tuple[float, float, float]: forward_backward_func=forward_backward_func, p2p_communicator=p2p_train, ) - # Bridge's train_step bumps state.train_state.step internally; record it + # Bridge's train_step bumps state.train_state.iteration internally; record it # so the next train_step uses a fresh micro-batch counter. - state.train_state.step = iteration + state.train_state.iteration = iteration # --- Eval step inside eval_cp_context (cp_eval) --- with eval_cp_context(model, eval_pgs, train_pgs): diff --git a/skills/nemo-mbridge-perf-cuda-graphs/SKILL.md b/skills/nemo-mbridge-perf-cuda-graphs/SKILL.md index e872b4ddd8..fdd9b39bd0 100644 --- a/skills/nemo-mbridge-perf-cuda-graphs/SKILL.md +++ b/skills/nemo-mbridge-perf-cuda-graphs/SKILL.md @@ -180,7 +180,7 @@ scope is `attn`, not `mamba`. model_config.cuda_graph_impl == "transformer_engine" and cuda_graph_helper is not None and not cuda_graph_helper.graphs_created() - and global_state.train_state.step - start_iteration == model_config.cuda_graph_warmup_steps + and global_state.train_state.iteration - start_iteration == model_config.cuda_graph_warmup_steps ): if model_config.cuda_graph_warmup_steps > 0 and should_toggle_forward_pre_hook: disable_forward_pre_hook(model, param_sync=False) diff --git a/src/megatron/bridge/models/bagel/bagel_step.py b/src/megatron/bridge/models/bagel/bagel_step.py index 8027b3dc8c..1b6dc5389a 100644 --- a/src/megatron/bridge/models/bagel/bagel_step.py +++ b/src/megatron/bridge/models/bagel/bagel_step.py @@ -139,7 +139,7 @@ def __call__( ) _initialize_scheduler(self.scheduler) if config.reset_reference_training_rng and not self.reference_rng_seeded: - if state.train_state.step == 0: + if state.train_state.iteration == 0: world_size = torch.distributed.get_world_size() if config.reference_training_world_size != world_size: raise RuntimeError("BAGEL reference world size differs from the training world size") diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 9e0c6be751..7bd7a078f4 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -64,6 +64,7 @@ unwrap_model, ) from megatron.training.checkpointing import save_tokenizer_assets +from megatron.training.utils.checkpoint_utils import get_checkpoint_train_state_filename, read_train_state from modelopt.torch.opt.plugins import ( restore_modelopt_state, save_modelopt_state, @@ -86,12 +87,10 @@ get_checkpoint_name, get_checkpoint_run_config_filename, get_checkpoint_tracker_filename, - get_checkpoint_train_state_filename, is_checkpoint_iteration_directory, is_hf_checkpoint_dir, join_paths, read_run_config, - read_train_state, ) from megatron.bridge.training.utils.log_utils import append_to_progress_log from megatron.bridge.training.utils.pg_utils import get_pg_collection @@ -1286,7 +1285,7 @@ def save_checkpoint( # Determine checkpoint type and save directory save_dir = ckpt_cfg.save is_global_non_persistent_ckpt = non_persistent_ckpt and ckpt_cfg.non_persistent_ckpt_type == "global" - checkpoint_step = train_state.step + checkpoint_step = train_state.iteration global_non_persistent_keep_count = 0 if is_global_non_persistent_ckpt: # Global non-persistent saves must leave a durable resume point, even when retention is configured as zero. @@ -1311,7 +1310,7 @@ def save_checkpoint( ckpt_type = CheckpointType.GLOBAL ckpt_format = ckpt_cfg.ckpt_format if ckpt_type == CheckpointType.GLOBAL else "torch" # torch for local - print_rank_0(f"saving checkpoint at iteration {train_state.step:7d} to {save_dir} in {ckpt_format} format") + print_rank_0(f"saving checkpoint at iteration {train_state.iteration:7d} to {save_dir} in {ckpt_format} format") # Collect rng state across data parallel ranks. if pg_collection is None: @@ -1333,7 +1332,7 @@ def save_checkpoint( ) # Checkpoint name. - checkpoint_name = get_checkpoint_name(save_dir, train_state.step, release=False) + checkpoint_name = get_checkpoint_name(save_dir, train_state.iteration, release=False) # Save LayerWiseDistributedOptimizer state - only for 'torch' (local) checkpoints. # For torch_dist/fsdp_dtensor formats, optimizer state is already included in the @@ -1375,7 +1374,7 @@ def save_checkpoint( optimizer, opt_param_scheduler, rng_state, - iteration=train_state.step, + iteration=train_state.iteration, optim_sd_kwargs=dict(metadata=sharded_sd_metadata), model_sd_kwargs=dict(metadata=sharded_sd_metadata), rerun_state=rerun_state, @@ -1562,7 +1561,7 @@ def save_checkpoint( parallelization_group=_checkpoint_dp_cp_group(pg_collection), ) async_save_request = checkpointing_context["local_checkpoint_manager"].save( - state_dict_for_save, train_state.step, is_async=bool(ckpt_cfg.async_save) + state_dict_for_save, train_state.iteration, is_async=bool(ckpt_cfg.async_save) ) checkpointing_context["local_checkpoint_cache"] = cacheable_metadata @@ -1579,7 +1578,7 @@ def save_checkpoint( maybe_save_dataloader_state( model, train_data_iterator, - train_state.step, + train_state.iteration, dataloader_save_path, pg_collection=pg_collection, data_parallel_group=( @@ -1604,7 +1603,7 @@ def save_checkpoint( config_filename = get_checkpoint_run_config_filename(checkpoint_name) tracker_filename = get_checkpoint_tracker_filename(save_dir) - step = train_state.step + step = train_state.iteration if ckpt_type == CheckpointType.LOCAL: def train_state_finalize_fn(): @@ -1820,7 +1819,7 @@ def cleanup_old_checkpoints_finalize_fn() -> None: if ckpt_cfg.async_save: schedule_async_save(state, async_save_request) - print_rank_0(f" scheduled an async checkpoint save at iteration {train_state.step:7d} to {save_dir}") + print_rank_0(f" scheduled an async checkpoint save at iteration {train_state.iteration:7d} to {save_dir}") if pending_hf_save_dir is not None: _save_hf_weights(state, model, pending_hf_save_dir) @@ -2794,7 +2793,7 @@ def _load_hf_pretrained_checkpoint( bridge = _build_auto_bridge_for_save(cfg, hf_source=hf_dir) bridge.load_hf_weights(model, hf_path=hf_dir) - state.train_state.step = 0 + state.train_state.iteration = 0 # HF checkpoints only initialize model weights. Full training resume should # load the complete Megatron checkpoint through ``checkpoint.load``. @@ -3174,7 +3173,7 @@ def _load_checkpoint_from_path( state.train_state.load_state_dict(state_dict["train_state_metadata"]) else: print_rank_0("WARNING: train_state_metadata not found in local checkpoint, counters reset") - state.train_state = TrainState(step=state_dict.get("iteration", 0)) + state.train_state = TrainState(iteration=state_dict.get("iteration", 0)) else: train_state_filename = get_checkpoint_train_state_filename(checkpoint_name) if file_exists(train_state_filename): @@ -3184,7 +3183,7 @@ def _load_checkpoint_from_path( state.train_state = _get_train_state_from_state_dict(state_dict) if cfg.checkpoint.finetune or release: - state.train_state.step = 0 + state.train_state.iteration = 0 # For local checkpoints, checkpoint_name is a CkptID tuple. # Normalize to string for downstream logging / wandb / mlflow. @@ -3394,7 +3393,7 @@ def _load_checkpoint_from_path( f" successfully loaded checkpoint from {load_dir} " f"[ t {pg_collection.tp.rank()}/{pg_collection.tp.size()}, " f"p {pg_collection.pp.rank()}/{pg_collection.pp.size()} ] " - f"at iteration {state.train_state.step}" + f"at iteration {state.train_state.iteration}" ) if not torch.distributed.is_initialized() or is_last_rank(): @@ -3404,11 +3403,11 @@ def _load_checkpoint_from_path( torch.cuda.empty_cache() - if state.train_state.step > 0: + if state.train_state.iteration > 0: is_local_chkpt = ckpt_type == CheckpointType.LOCAL fault_tolerance.on_checkpoint_loaded(is_local_chkpt=is_local_chkpt, global_state=state) - return state.train_state.step, state.train_state.floating_point_operations_so_far + return state.train_state.iteration, state.train_state.num_floating_point_operations_so_far def init_checkpointing_context(checkpoint_config: CheckpointConfig) -> dict[str, Any]: @@ -3571,7 +3570,7 @@ def _resolve_checkpoint_iteration(load_dir: str | None, ckpt_step_override: int tracker_filename = get_checkpoint_train_state_filename(load_dir, prefix=TRACKER_PREFIX) if file_exists(tracker_filename): train_state = read_train_state(tracker_filename) - iteration = train_state.step + iteration = train_state.iteration else: # Fallback to legacy Megatron-LM format (latest_checkpointed_iteration.txt) legacy_tracker_filename = get_checkpoint_tracker_filename(load_dir) @@ -3594,7 +3593,7 @@ def _get_non_persistent_iteration( train_state_filename = get_checkpoint_train_state_filename(non_persistent_global_dir, prefix=TRACKER_PREFIX) if file_exists(train_state_filename): train_state = read_train_state(train_state_filename) - iteration = train_state.step + iteration = train_state.iteration # if train_state.release: # raise RuntimeError("Non-persistent checkpoint can't be a release checkpoint") else: @@ -4109,7 +4108,7 @@ def _checkpoint_dp_cp_group(pg_collection: ProcessGroupCollection) -> torch.dist def _get_train_state_from_state_dict(state_dict: dict[str, Any]) -> TrainState: """Create a TrainState from the state dict from a Megatron-LM checkpoint.""" legacy_train_state = TrainState() - legacy_train_state.step = state_dict.get("iteration", 0) + legacy_train_state.iteration = state_dict.get("iteration", 0) # Extract training progress from checkpoint args (like Megatron-LM does) checkpoint_args = state_dict.get("args", None) @@ -4124,7 +4123,7 @@ def _get_train_state_from_state_dict(state_dict: dict[str, Any]) -> TrainState: legacy_train_state.consumed_valid_samples = 0 # Extract floating point operations count from state_dict (like Megatron-LM does) - legacy_train_state.floating_point_operations_so_far = state_dict.get("num_floating_point_operations_so_far", 0) + legacy_train_state.num_floating_point_operations_so_far = state_dict.get("num_floating_point_operations_so_far", 0) legacy_train_state.do_train = False legacy_train_state.do_valid = False legacy_train_state.do_test = False diff --git a/src/megatron/bridge/training/eval.py b/src/megatron/bridge/training/eval.py index 3a60292c89..dda9084748 100644 --- a/src/megatron/bridge/training/eval.py +++ b/src/megatron/bridge/training/eval.py @@ -475,40 +475,40 @@ def evaluate_and_print_results( ppl = math.exp(min(20, total_loss_dict[key].item())) string += "{} PPL: {:.6E} | ".format(key, ppl) if writer: - writer.add_scalar("{} validation".format(key), total_loss_dict[key].item(), state.train_state.step) + writer.add_scalar("{} validation".format(key), total_loss_dict[key].item(), state.train_state.iteration) writer.add_scalar( "{} validation vs samples".format(key), total_loss_dict[key].item(), state.train_state.consumed_train_samples, ) if state.cfg.logger.log_validation_ppl_to_tensorboard: - writer.add_scalar("{} validation ppl".format(key), ppl, state.train_state.step) + writer.add_scalar("{} validation ppl".format(key), ppl, state.train_state.iteration) writer.add_scalar( "{} validation ppl vs samples".format(key), ppl, state.train_state.consumed_train_samples ) if wandb_writer and is_last_rank(): - wandb_writer.log({"{} validation".format(key): total_loss_dict[key].item()}, state.train_state.step) + wandb_writer.log({"{} validation".format(key): total_loss_dict[key].item()}, state.train_state.iteration) if state.cfg.logger.log_validation_ppl_to_tensorboard: - wandb_writer.log({"{} validation ppl".format(key): ppl}, state.train_state.step) + wandb_writer.log({"{} validation ppl".format(key): ppl}, state.train_state.iteration) if mlflow_writer and is_last_rank(): mlflow_writer.log_metrics( - _sanitize_mlflow_metrics({f"val/{key}": total_loss_dict[key].item()}), step=state.train_state.step + _sanitize_mlflow_metrics({f"val/{key}": total_loss_dict[key].item()}), step=state.train_state.iteration ) if state.cfg.logger.log_validation_ppl_to_tensorboard: mlflow_writer.log_metrics( - _sanitize_mlflow_metrics({f"val/{key} ppl": ppl}), step=state.train_state.step + _sanitize_mlflow_metrics({f"val/{key} ppl": ppl}), step=state.train_state.iteration ) if comet_logger and is_last_rank(): comet_logger.log_metrics( - {"{} validation".format(key): total_loss_dict[key].item()}, step=state.train_state.step + {"{} validation".format(key): total_loss_dict[key].item()}, step=state.train_state.iteration ) if state.cfg.logger.log_validation_ppl_to_tensorboard: - comet_logger.log_metrics({"{} validation ppl".format(key): ppl}, step=state.train_state.step) + comet_logger.log_metrics({"{} validation ppl".format(key): ppl}, step=state.train_state.iteration) if process_non_loss_data_func is not None and writer and is_last_rank(): - process_non_loss_data_func(collected_non_loss_data, state.train_state.step, writer) + process_non_loss_data_func(collected_non_loss_data, state.train_state.iteration, writer) length = len(string) + 1 print_rank_last("-" * length) diff --git a/src/megatron/bridge/training/pretrain.py b/src/megatron/bridge/training/pretrain.py index f5b6077f77..50cab95522 100644 --- a/src/megatron/bridge/training/pretrain.py +++ b/src/megatron/bridge/training/pretrain.py @@ -166,7 +166,7 @@ def _pretrain( else: print_rank_0("skipping training ...") - iteration = state.train_state.step + iteration = state.train_state.iteration # VALIDATION if state.train_state.do_valid: diff --git a/src/megatron/bridge/training/setup.py b/src/megatron/bridge/training/setup.py index 6566a88f61..f94f2b2d02 100644 --- a/src/megatron/bridge/training/setup.py +++ b/src/megatron/bridge/training/setup.py @@ -172,9 +172,7 @@ def _should_load_checkpoint(cfg: ConfigContainer, checkpoint_manager: Checkpoint "local_checkpoint_manager" in checkpointing_context and checkpointing_context["local_checkpoint_manager"].find_latest() != -1 ) - has_global_non_persistent_checkpoint = _has_global_non_persistent_checkpoint( - cfg.checkpoint.load, cfg.checkpoint - ) + has_global_non_persistent_checkpoint = _has_global_non_persistent_checkpoint(cfg.checkpoint.load, cfg.checkpoint) if cfg.peft is not None: load_checkpoint_exists = cfg.checkpoint.load is not None and ( @@ -401,9 +399,7 @@ def modelopt_pre_wrap_hook(model): checkpoint_path = cfg.checkpoint.pretrained_checkpoint ckpt_step = None else: - raise RuntimeError( - "No checkpoint source is available for ModelOpt state restoration" - ) + raise RuntimeError("No checkpoint source is available for ModelOpt state restoration") if not has_modelopt_state(checkpoint_path, ckpt_step=ckpt_step): raise RuntimeError(f"No modelopt_state found in selected checkpoint={checkpoint_path}") @@ -472,7 +468,7 @@ def modelopt_pre_wrap_hook(model): state.tensorboard_logger, state.wandb_logger, comet_logger=state.comet_logger, - current_training_step=state.train_state.step, + current_training_step=state.train_state.iteration, ) _update_model_config_funcs( @@ -523,18 +519,18 @@ def modelopt_pre_wrap_hook(model): # Resume the dataloader stream position so a resumed run continues over the same data (currently # only Megatron Energon). Runs after the iterator is built and the model checkpoint load restored - # state.train_state.step. The default source is resolved by load_checkpoint from the checkpoint + # state.train_state.iteration. The default source is resolved by load_checkpoint from the checkpoint # actually selected (recorded as "dataloader_state_dir"); an explicit dataset.dataloader_load # overrides. Gated on step > 0 so only a real resume restores -- fresh, finetune, and # pretrained-init runs (step reset to 0) start the data stream from the beginning. - if state.train_state.step > 0: + if state.train_state.iteration > 0: dataloader_load_path = getattr(cfg.dataset, "dataloader_load", None) if dataloader_load_path is None: ckpt_ctx = getattr(checkpoint_manager, "checkpointing_context", {}) dataloader_load_path = ckpt_ctx.get("dataloader_state_dir") maybe_load_dataloader_state( train_data_iterator, - state.train_state.step, + state.train_state.iteration, dataloader_load_path, pg_collection=pg_collection, data_parallel_group=get_data_distribution_group(pg_collection, cfg.model), @@ -589,7 +585,9 @@ def _register_setup_pre_wrap_hook( ] else: model_cfg._pre_wrap_hooks[:] = [ - registered_hook for registered_hook in model_cfg._pre_wrap_hooks if registered_hook is not previous_hook + registered_hook + for registered_hook in model_cfg._pre_wrap_hooks + if registered_hook is not previous_hook ] setup_hooks[setup_hook_name] = hook @@ -660,8 +658,7 @@ def _update_model_config_funcs( # every backward finalizes the DP-outer axis and only the last microbatch's gradient # reaches the optimizer. if isinstance(model[0], FullyShardedDataParallelV2) or ( - isinstance(model[0], (DistributedDataParallel, FullyShardedDataParallelV1)) - and ddp_config.overlap_grad_reduce + isinstance(model[0], (DistributedDataParallel, FullyShardedDataParallelV1)) and ddp_config.overlap_grad_reduce ): assert model_config.no_sync_func is None, ( "config.no_sync_func must be None when the wrapper supplies its own no_sync " diff --git a/src/megatron/bridge/training/setup_megatron_mimo.py b/src/megatron/bridge/training/setup_megatron_mimo.py index cc3cc0eb29..5254b41b28 100644 --- a/src/megatron/bridge/training/setup_megatron_mimo.py +++ b/src/megatron/bridge/training/setup_megatron_mimo.py @@ -337,7 +337,7 @@ def setup_megatron_mimo( if build_data_iterators_fn is not None: logger.info(f"Rank {dist.get_rank()}: Building data iterators") train_state = global_state.train_state - is_resuming = train_state.step > 0 + is_resuming = train_state.iteration > 0 if is_resuming: import inspect diff --git a/src/megatron/bridge/training/state.py b/src/megatron/bridge/training/state.py index 717db4a53d..74ba978d80 100644 --- a/src/megatron/bridge/training/state.py +++ b/src/megatron/bridge/training/state.py @@ -19,13 +19,12 @@ from dataclasses import dataclass from typing import Any, Optional -import torch from megatron.core.dist_checkpointing.strategies.nvrx import has_nvrx_async_support from megatron.core.energy_monitor import EnergyMonitor from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.timers import Timers from megatron.core.utils import StragglerDetector -from torch.distributed.checkpoint.stateful import Stateful +from megatron.training.state import TrainState from torch.utils.tensorboard.writer import SummaryWriter from megatron.bridge.training.config import ConfigContainer @@ -42,62 +41,6 @@ HAVE_NVRX = has_nvrx_async_support() -@dataclass -class TrainState(Stateful): - """Dataclass to hold the state of the training process. - - Inherits from Stateful for distributed checkpointing compatibility. - Tracks iteration count, consumed samples, flags for train/valid/test phases, - and floating-point operations. - """ - - step: int = 0 - consumed_train_samples: int = 0 - skipped_train_samples: int = 0 - consumed_valid_samples: int = 0 - floating_point_operations_so_far: int = 0 - do_train: bool = False - do_valid: bool = False - do_test: bool = False - - def state_dict(self) -> dict[str, torch.Tensor]: - """Serializes the training state into a dictionary of tensors. - - Conforms to the Stateful interface for distributed checkpointing. - - Returns: - A dictionary where keys are state variable names and values are - their corresponding tensor representations. - """ - return { - "step": torch.tensor(self.step, dtype=torch.int64), - "consumed_train_samples": torch.tensor(self.consumed_train_samples, dtype=torch.int64), - "skipped_train_samples": torch.tensor(self.skipped_train_samples, dtype=torch.int64), - "consumed_valid_samples": torch.tensor(self.consumed_valid_samples, dtype=torch.int64), - "floating_point_operations_so_far": torch.tensor( - self.floating_point_operations_so_far, dtype=torch.float64 - ), - "do_train": torch.tensor(self.do_train, dtype=torch.bool), - "do_valid": torch.tensor(self.do_valid, dtype=torch.bool), - "do_test": torch.tensor(self.do_test, dtype=torch.bool), - } - - def load_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None: - """Load the training state from a state dictionary. - - Args: - state_dict: A dictionary containing the state variables as tensors. - """ - self.step = state_dict["step"].item() - self.consumed_train_samples = state_dict["consumed_train_samples"].item() - self.skipped_train_samples = state_dict["skipped_train_samples"].item() - self.consumed_valid_samples = state_dict["consumed_valid_samples"].item() - self.floating_point_operations_so_far = state_dict["floating_point_operations_so_far"].item() - self.do_train = state_dict["do_train"].item() - self.do_valid = state_dict["do_valid"].item() - self.do_test = state_dict["do_test"].item() - - @dataclass class FaultToleranceState: """Dataclass to hold state specific to fault tolerance mechanisms.""" @@ -482,11 +425,6 @@ def reset_for_restart(self) -> None: This cleans up all stateful components that need to be reinitialized between restart iterations. The async calls queue for checkpointing is handled separately in aborting in order to clean up persistent workers. """ - from megatron.bridge.training.utils.checkpoint_utils import read_train_state - - # The top-level tracker is overwritten after every checkpoint. A restart must - # not reuse the iteration cached by an earlier invocation in this process. - read_train_state.cache_clear() if self.cfg is not None: model_config = getattr(self.cfg.model, "transformer", self.cfg.model) for callback_name in ( diff --git a/src/megatron/bridge/training/train.py b/src/megatron/bridge/training/train.py index 542d948ecb..5989552750 100644 --- a/src/megatron/bridge/training/train.py +++ b/src/megatron/bridge/training/train.py @@ -175,7 +175,7 @@ def train( # # Note on reference semantics: # - functools.partial stores a reference to global_state, not a copy - # - When global_state.train_state.step changes, the partial sees the updated value + # - When global_state.train_state.iteration changes, the partial sees the updated value # - This is safe because GlobalState is a mutable object passed by reference # # For functors (classes with __call__ defined): @@ -192,11 +192,11 @@ def train( # Make sure rerun_state_machine has the right iteration loaded from checkpoint. rerun_state_machine = get_rerun_state_machine() - if rerun_state_machine.current_iteration != global_state.train_state.step: - print_rank_0(f"Setting rerun_state_machine.current_iteration to {global_state.train_state.step}...") - rerun_state_machine.current_iteration = global_state.train_state.step + if rerun_state_machine.current_iteration != global_state.train_state.iteration: + print_rank_0(f"Setting rerun_state_machine.current_iteration to {global_state.train_state.iteration}...") + rerun_state_machine.current_iteration = global_state.train_state.iteration - num_floating_point_operations_so_far = global_state.train_state.floating_point_operations_so_far + num_floating_point_operations_so_far = global_state.train_state.num_floating_point_operations_so_far num_floating_point_operations_since_last_log_event = 0.0 if energy_monitor is not None: @@ -267,7 +267,7 @@ def train( "Parameter hashes not matching across DP replicas" ) torch.distributed.barrier() - print_rank_0(f">>> Weight hashes match after {global_state.train_state.step} iterations...") + print_rank_0(f">>> Weight hashes match after {global_state.train_state.iteration} iterations...") # Capture CUDA Graphs. cuda_graph_helper = None @@ -337,16 +337,16 @@ def train( model_config, copy_main_params, model, optimizer, forward_backward_func ) - start_iteration = global_state.train_state.step + start_iteration = global_state.train_state.iteration print_rank_0(f"Starting training loop at iteration {start_iteration}") p2p_communicator = P2PCommunicator(pp_group=pg_collection.pp, config=model_config) data_distribution_group = get_data_distribution_group(pg_collection, config.model) dp_size = data_distribution_group.size() # Anchor for interval-average throughput logging: training_log reports the FLOPS # performed over each logging interval as the delta of - # floating_point_operations_so_far. Seed it with the current cumulative (0 fresh, + # num_floating_point_operations_so_far. Seed it with the current cumulative (0 fresh, # or the checkpoint value on resume) so the first interval's delta is correct. - global_state._flops_at_last_log = global_state.train_state.floating_point_operations_so_far + global_state._flops_at_last_log = global_state.train_state.num_floating_point_operations_so_far if hasattr(config.model, "dist_train") and getattr(config.model.dist_train, "use_dist_train", False) is True: forward_backward_func = forward_backward_pipelining_without_interleaving p2p_communicator = config.model._p2p_communicator @@ -379,18 +379,18 @@ def train( pre_hook_enabled = True # Run training iterations till done. - while global_state.train_state.step < train_config.train_iters: + while global_state.train_state.iteration < train_config.train_iters: # Handle profiling for this step nvtx_ctx = handle_profiling_step( prof_config, - global_state.train_state.step, + global_state.train_state.iteration, torch.distributed.get_rank(), prof, ) if nvtx_ctx is not None: nsys_nvtx_context = nvtx_ctx - nvtx_step = global_state.train_state.step + nvtx_step = global_state.train_state.iteration nvtx_range_push(suffix=f"training_step_{nvtx_step}") fault_tolerance.on_checkpointing_start(global_state) @@ -400,7 +400,7 @@ def train( # Update the timeout for all process groups after initialization # We update the timeout after the first successful iteration, # which takes longer than others usually - if global_state.train_state.step == start_iteration + 1: + if global_state.train_state.iteration == start_iteration + 1: distributed_timeout_seconds_after_init = global_state.cfg.dist.distributed_timeout_seconds_after_init if distributed_timeout_seconds_after_init is not None: update_pg_timeout(timedelta(seconds=distributed_timeout_seconds_after_init)) @@ -410,7 +410,7 @@ def train( # from the previous iteration, save a checkpoint. Then run consistency check # to make sure training configuration is still valid. update_num_microbatches(global_state.train_state.consumed_train_samples, consistency_check=False, verbose=True) - if get_num_microbatches() != num_microbatches and global_state.train_state.step != 0: + if get_num_microbatches() != num_microbatches and global_state.train_state.iteration != 0: assert get_num_microbatches() > num_microbatches, ( f"Number of microbatches should be increasing due to batch size rampup; " f"instead going from {num_microbatches} to {get_num_microbatches()}" @@ -432,12 +432,12 @@ def train( # Completely skip iteration if needed. if _should_skip_and_handle_iteration(global_state, train_data_iterator, pg_collection): - if global_state.train_state.step == start_iteration + 1: - start_iteration = global_state.train_state.step + if global_state.train_state.iteration == start_iteration + 1: + start_iteration = global_state.train_state.iteration nvtx_range_pop(suffix=f"training_step_{nvtx_step}") handle_profiling_stop( config.profiling, - global_state.train_state.step, + global_state.train_state.iteration, torch.distributed.get_rank(), prof, nsys_nvtx_context, @@ -449,7 +449,7 @@ def train( model_config.cuda_graph_impl == "transformer_engine" and cuda_graph_helper is not None and not cuda_graph_helper.graphs_created() - and global_state.train_state.step - start_iteration == model_config.cuda_graph_warmup_steps + and global_state.train_state.iteration - start_iteration == model_config.cuda_graph_warmup_steps ): if model_config.cuda_graph_warmup_steps > 0 and should_toggle_forward_pre_hook: disable_forward_pre_hook(model, param_sync=False) @@ -461,7 +461,7 @@ def train( if ( vision_cuda_graph_helper is not None and not vision_cuda_graph_helper.graphs_created() - and global_state.train_state.step - start_iteration == model_config.cuda_graph_warmup_steps + and global_state.train_state.iteration - start_iteration == model_config.cuda_graph_warmup_steps ): vision_cuda_graph_helper.create_cudagraphs() vision_cuda_graph_helper.cuda_graph_set_manual_hooks() @@ -552,7 +552,7 @@ def train( nvtx_range_pop(suffix=f"training_step_{nvtx_step}") if ( prof_config is not None - and global_state.train_state.step < prof_config.profile_step_end + and global_state.train_state.iteration < prof_config.profile_step_end and (prof is not None or nsys_nvtx_context is not None) ): handle_profiling_stop( @@ -567,12 +567,12 @@ def train( # Enable forward pre-hooks after first set of forward and backward passes. # When running in fp16, skip all NaN iterations until steady-state loss scaling value # is reached. - if global_state.train_state.step == start_iteration: + if global_state.train_state.iteration == start_iteration: if skipped_iter: # Only enable forward pre-hook after a training step has successfully run. Relevant # for fp16 codepath where first XX iterations are skipped until steady-state loss # scale value is reached. - start_iteration = global_state.train_state.step + 1 + start_iteration = global_state.train_state.iteration + 1 else: # Enable forward pre-hook after training step has successfully run. All subsequent # forward passes will use the forward pre-hook / `param_sync_func` in @@ -595,10 +595,10 @@ def train( and vision_cuda_graph_helper.graphs_created() ): vision_cuda_graph_helper.cuda_graph_set_manual_hooks() - global_state.train_state.step += 1 + global_state.train_state.iteration += 1 # If fsdp_manual_registration is enabled, manually register FSDP communication buffers after one training step. - if global_state.train_state.step == start_iteration + 1 and config.ddp.use_megatron_fsdp: + if global_state.train_state.iteration == start_iteration + 1 and config.ddp.use_megatron_fsdp: _maybe_register_fsdp_buffers(config, model) batch_size = dp_size * train_config.micro_batch_size * get_num_microbatches() @@ -641,8 +641,8 @@ def train( patch_squared_sum=flops_stats.vision_patch_squared_sum, merged_token_sum=flops_stats.vision_merged_token_sum, ) - global_state.train_state.floating_point_operations_so_far += num_floating_point_operations_in_batch - num_floating_point_operations_so_far = global_state.train_state.floating_point_operations_so_far + global_state.train_state.num_floating_point_operations_so_far += num_floating_point_operations_in_batch + num_floating_point_operations_so_far = global_state.train_state.num_floating_point_operations_so_far num_floating_point_operations_since_last_log_event += num_floating_point_operations_in_batch # Logging. @@ -694,9 +694,10 @@ def train( global_state.train_state.do_valid and val_config.eval_interval and ( - val_config.start_eval_at_iter is None or global_state.train_state.step >= val_config.start_eval_at_iter + val_config.start_eval_at_iter is None + or global_state.train_state.iteration >= val_config.start_eval_at_iter ) - and global_state.train_state.step % val_config.eval_interval == 0 + and global_state.train_state.iteration % val_config.eval_interval == 0 ): if energy_monitor is not None: energy_monitor.pause() @@ -707,7 +708,7 @@ def train( if train_config.manual_gc and train_config.manual_gc_eval: # Collect all objects. gc.collect() - prefix = f"iteration {global_state.train_state.step}" + prefix = f"iteration {global_state.train_state.iteration}" timers("eval-time", log_level=0).start(barrier=True) evaluate_and_print_results( global_state, @@ -736,25 +737,25 @@ def train( # Miscellaneous post-training-step functions (e.g., FT heartbeats, GC). # Some of these only happen at specific iterations. - maybe_synchronize_training_step(config.train.train_sync_interval, global_state.train_state.step) + maybe_synchronize_training_step(config.train.train_sync_interval, global_state.train_state.iteration) num_floating_point_operations_since_last_log_event = maybe_report_stragglers( config.logger.log_interval, bool(getattr(config.straggler, "log_straggler", False)), straggler_timer, - global_state.train_state.step, + global_state.train_state.iteration, num_floating_point_operations_since_last_log_event, ) maybe_check_weight_hash_across_dp_replicas( model, optimizer, config.train.check_weight_hash_across_dp_replicas_interval, - global_state.train_state.step, + global_state.train_state.iteration, should_toggle_forward_pre_hook, ) nvtx_range_pop(suffix=f"training_step_{nvtx_step}") handle_profiling_stop( config.profiling, - global_state.train_state.step, + global_state.train_state.iteration, torch.distributed.get_rank(), prof, nsys_nvtx_context, @@ -762,7 +763,7 @@ def train( maybe_run_manual_gc( config.train.manual_gc, config.train.manual_gc_interval, - global_state.train_state.step, + global_state.train_state.iteration, ) # Checkpoint and decide whether to exit. @@ -787,9 +788,12 @@ def train( ckpt_config = config.checkpoint if ( ckpt_config.save - and global_state.train_state.step != 0 + and global_state.train_state.iteration != 0 and ckpt_config.save_interval != 0 - and (ckpt_config.save_interval is None or global_state.train_state.step % ckpt_config.save_interval != 0) + and ( + ckpt_config.save_interval is None + or global_state.train_state.iteration % ckpt_config.save_interval != 0 + ) ): save_checkpoint_and_time( global_state, @@ -1275,9 +1279,9 @@ def compute_throughputs_and_append_to_progress_log( # Compute job throughput. # num_floating_point_operations_so_far keeps track of floating-point operations # completed at the start of job. - job_throughput = (num_floating_point_operations_so_far - state.train_state.floating_point_operations_so_far) / ( - (time.time() - state.start_time) * 10**12 * get_world_size_safe() - ) + job_throughput = ( + num_floating_point_operations_so_far - state.train_state.num_floating_point_operations_so_far + ) / ((time.time() - state.start_time) * 10**12 * get_world_size_safe()) # Compute cumulative throughput since jobs of this world size were launched. # `get_start_time_from_progress_log` returns start time and number of floating-point @@ -1292,7 +1296,7 @@ def compute_throughputs_and_append_to_progress_log( saved_ckpt_prefix = "Saving async checkpoint" if state.cfg.checkpoint.async_save else "Saved checkpoint" append_to_progress_log( state.cfg.checkpoint.save, - f"{saved_ckpt_prefix}\tIteration: {state.train_state.step}\t" + f"{saved_ckpt_prefix}\tIteration: {state.train_state.iteration}\t" f"Job throughput: {job_throughput:.1f} MODEL_TFLOP/s/GPU\t" f"Cumulative throughput: {cumulative_throughput:.1f} MODEL_TFLOP/s/GPU\t" f"Floating-point operations: {num_floating_point_operations_so_far:.2e}\t" @@ -1453,7 +1457,7 @@ def checkpoint_and_decide_exit( if ( state.cfg.checkpoint.save and state.cfg.checkpoint.save_interval - and state.train_state.step % state.cfg.checkpoint.save_interval == 0 + and state.train_state.iteration % state.cfg.checkpoint.save_interval == 0 ): save_checkpoint_and_time( state, @@ -1472,7 +1476,7 @@ def checkpoint_and_decide_exit( elif ( state.cfg.checkpoint.save and state.cfg.checkpoint.non_persistent_save_interval - and state.train_state.step % state.cfg.checkpoint.non_persistent_save_interval == 0 + and state.train_state.iteration % state.cfg.checkpoint.non_persistent_save_interval == 0 ): save_checkpoint_and_time( state, @@ -1514,7 +1518,7 @@ def checkpoint_and_decide_exit( return True # Exit based on iterations. - if state.cfg.train.exit_interval and state.train_state.step % state.cfg.train.exit_interval == 0: + if state.cfg.train.exit_interval and state.train_state.iteration % state.cfg.train.exit_interval == 0: if state.cfg.checkpoint.save and not saved_checkpoint: save_checkpoint_and_time( state, @@ -1528,7 +1532,7 @@ def checkpoint_and_decide_exit( callback_manager=callback_manager, module_name=module_name, ) - barrier_and_log(f"exiting program at iteration {state.train_state.step}") + barrier_and_log(f"exiting program at iteration {state.train_state.iteration}") return True @@ -1599,14 +1603,14 @@ def _should_skip_and_handle_iteration( bool: True if the iteration was skipped, False otherwise """ cfg = global_state.cfg - if (global_state.train_state.step + 1) not in cfg.train.iterations_to_skip: + if (global_state.train_state.iteration + 1) not in cfg.train.iterations_to_skip: return False # Perform dummy train step to fast forward train_data_iterator _dummy_train_step(global_state, train_data_iterator, pg_collection) # Update step and sample counters - global_state.train_state.step += 1 + global_state.train_state.iteration += 1 dp_size = get_data_distribution_group(pg_collection, cfg.model).size() batch_size = dp_size * cfg.train.micro_batch_size * get_num_microbatches() global_state.train_state.consumed_train_samples += batch_size diff --git a/src/megatron/bridge/training/train_megatron_mimo.py b/src/megatron/bridge/training/train_megatron_mimo.py index 600741e5bd..bccafdd914 100644 --- a/src/megatron/bridge/training/train_megatron_mimo.py +++ b/src/megatron/bridge/training/train_megatron_mimo.py @@ -304,7 +304,7 @@ def train_megatron_mimo( # Main training loop timers("interval-time", log_level=0).start(barrier=True) - while train_state.step < train_config.train_iters: + while train_state.iteration < train_config.train_iters: # Finalize any pending async saves (non-blocking). Placed at the top # of the loop so async saves get a full iteration to complete. checkpoint_manager.finalize_async_saves( @@ -315,7 +315,7 @@ def train_megatron_mimo( # Handle profiling nsys_ctx = handle_profiling_step( prof_config, - train_state.step, + train_state.iteration, dist.get_rank(), prof, ) @@ -348,7 +348,7 @@ def train_megatron_mimo( history_wct.append(iteration_time) # Update training state - train_state.step += 1 + train_state.iteration += 1 train_state.consumed_train_samples += micro_batch_size * num_microbatches * cfg.data_parallel_size # Get learning rate from first scheduler @@ -389,15 +389,17 @@ def train_megatron_mimo( # Evaluation at specified intervals if ( eval_interval - and (cfg.validation.start_eval_at_iter is None or train_state.step >= cfg.validation.start_eval_at_iter) - and train_state.step % eval_interval == 0 + and ( + cfg.validation.start_eval_at_iter is None or train_state.iteration >= cfg.validation.start_eval_at_iter + ) + and train_state.iteration % eval_interval == 0 and valid_data_iterator is not None ): if train_config.manual_gc and train_config.manual_gc_eval: gc.collect() evaluate_and_print_results( state=global_state, - prefix=f"iteration {train_state.step}", + prefix=f"iteration {train_state.iteration}", forward_step_func=forward_step_func, data_iterator=valid_data_iterator, model=[model], @@ -414,7 +416,7 @@ def train_megatron_mimo( maybe_run_manual_gc( train_config.manual_gc, train_config.manual_gc_interval, - train_state.step, + train_state.iteration, ) # Checkpointing (interval, signal, duration, exit-interval) and exit decision. @@ -433,12 +435,12 @@ def train_megatron_mimo( if not profiling_stopped: handle_profiling_stop( prof_config, - train_state.step, + train_state.iteration, dist.get_rank(), prof, nsys_nvtx_context, ) - profiling_stopped = prof_config is not None and train_state.step == prof_config.profile_step_end + profiling_stopped = prof_config is not None and train_state.iteration == prof_config.profile_step_end if should_exit: break @@ -449,9 +451,9 @@ def train_megatron_mimo( ckpt_config = cfg.checkpoint if ( ckpt_config.save - and train_state.step != 0 + and train_state.iteration != 0 and ckpt_config.save_interval != 0 - and (ckpt_config.save_interval is None or train_state.step % ckpt_config.save_interval != 0) + and (ckpt_config.save_interval is None or train_state.iteration % ckpt_config.save_interval != 0) ): save_checkpoint_and_time( state=global_state, @@ -468,7 +470,7 @@ def train_megatron_mimo( if not profiling_stopped: handle_profiling_stop( prof_config, - train_state.step, + train_state.iteration, dist.get_rank(), prof, nsys_nvtx_context, diff --git a/src/megatron/bridge/training/utils/checkpoint_utils.py b/src/megatron/bridge/training/utils/checkpoint_utils.py index ecb3b934a7..be5ca2e522 100644 --- a/src/megatron/bridge/training/utils/checkpoint_utils.py +++ b/src/megatron/bridge/training/utils/checkpoint_utils.py @@ -23,13 +23,12 @@ import torch import yaml from megatron.core.msc_utils import MultiStorageClientFeature +from megatron.training.utils.checkpoint_utils import TRAIN_STATE_FILE -from megatron.bridge.training.state import TrainState from megatron.bridge.training.utils.config_utils import apply_run_config_backward_compat from megatron.bridge.utils.common_utils import get_rank_safe, get_world_size_safe, print_rank_0 -TRAIN_STATE_FILE = "train_state.pt" TRACKER_PREFIX = "latest" CONFIG_FILE = "run_config.yaml" @@ -110,24 +109,6 @@ def get_checkpoint_name(checkpoints_path: str, iteration: int, release: bool = F return common_path -def get_checkpoint_train_state_filename(checkpoints_path: str, prefix: Optional[str] = None) -> str: - """Get the filename for the train state tracker file. - - This file typically stores metadata about the latest checkpoint, like the iteration number. - - Args: - checkpoints_path: Base directory where checkpoints are stored. - prefix: Optional prefix (e.g., 'latest') to prepend to the filename. - - Returns: - The full path to the train state tracker file. - """ - if prefix is None: - return join_paths(checkpoints_path, TRAIN_STATE_FILE) - else: - return join_paths(checkpoints_path, f"{prefix}_{TRAIN_STATE_FILE}") - - def get_checkpoint_run_config_filename(checkpoints_path: str) -> str: """Get the filename for the run configuration file within a checkpoint directory. @@ -446,57 +427,6 @@ def read_run_config(run_config_filename: str) -> dict[str, Any]: return config_dict -@lru_cache() -def read_train_state(train_state_filename: str) -> TrainState: - """Read the train state metadata from a YAML file (rank 0 only). - - Reads the file on rank 0 and broadcasts the result to other ranks if - torch.distributed is initialized. Otherwise, loads the file locally. - - Args: - train_state_filename: Path to the train state YAML file. - - Returns: - An initialized TrainState object. - """ - if torch.distributed.is_initialized(): - state_obj = [None] - if get_rank_safe() == 0: - try: - if MultiStorageClientFeature.is_enabled(): - msc = MultiStorageClientFeature.import_package() - state_dict = msc.torch.load(train_state_filename, map_location="cpu", weights_only=True) - else: - state_dict = torch.load(train_state_filename, map_location="cpu", weights_only=True) - ts = TrainState() - ts.load_state_dict(state_dict) - state_obj[0] = ts - except Exception as e: - error_msg = f"ERROR: Unable to load train state file {train_state_filename}: {e}" - sys.stderr.write(error_msg + "\n") - state_obj[0] = {"error": True, "msg": error_msg} - - print_rank_0(f"Broadcasting TrainState from rank 0 to all {get_world_size_safe()} ranks") - torch.distributed.broadcast_object_list(state_obj, src=0) - - if isinstance(state_obj[0], dict) and state_obj[0].get("error", False): - raise RuntimeError(state_obj[0]["msg"]) - - return state_obj[0] - - try: - if MultiStorageClientFeature.is_enabled(): - msc = MultiStorageClientFeature.import_package() - state_dict = msc.torch.load(train_state_filename, map_location="cpu", weights_only=True) - else: - state_dict = torch.load(train_state_filename, map_location="cpu", weights_only=True) - ts = TrainState() - ts.load_state_dict(state_dict) - return ts - except Exception as e: - raise RuntimeError(f"Unable to load train state file {train_state_filename}: {e}") from e - - def _sanitize_run_config_object(obj: Any) -> Any: """Remove runtime-only objects from run config dictionaries. diff --git a/src/megatron/bridge/training/utils/train_utils.py b/src/megatron/bridge/training/utils/train_utils.py index 800a6da97f..723aa24843 100644 --- a/src/megatron/bridge/training/utils/train_utils.py +++ b/src/megatron/bridge/training/utils/train_utils.py @@ -706,7 +706,7 @@ def training_log( """ timers = global_state.timers train_state = global_state.train_state - iteration = train_state.step + iteration = train_state.iteration writer = global_state.tensorboard_logger wandb_writer = global_state.wandb_logger mlflow_logger = global_state.mlflow_logger @@ -1086,7 +1086,7 @@ def training_log( # # The main training loop folds each step's DP-exact FLOPS (Σ tokens for the # linear terms, Σᵢ sᵢ² for THD attention) into - # ``train_state.floating_point_operations_so_far`` every step, so the FLOPS + # ``train_state.num_floating_point_operations_so_far`` every step, so the FLOPS # performed over this logging interval is exactly the delta since the last # log. Dividing that interval total by the *interval* elapsed time (not the # per-iteration average) keeps numerator and denominator over the same window. @@ -1094,13 +1094,13 @@ def training_log( # log-boundary step alone) divided by an interval-averaged time would over- or # under-report whenever that step is heavier/lighter than the interval mean. # Using the cumulative delta is also the single source of truth — it cannot - # disagree with ``floating_point_operations_so_far`` and needs no extra reduce. + # disagree with ``num_floating_point_operations_so_far`` and needs no extra reduce. num_flops = None if hasattr(config.model, "kv_channels") and hasattr(config.model, "num_attention_heads"): - # Coerce to float — floating_point_operations_so_far is a float in + # Coerce to float — num_floating_point_operations_so_far is a float in # production, but getattr on a MagicMock test double (and an unset anchor) # returns a MagicMock; treat non-numerics as 0 so the math stays valid. - flops_so_far = train_state.floating_point_operations_so_far + flops_so_far = train_state.num_floating_point_operations_so_far if not isinstance(flops_so_far, (int, float)): flops_so_far = 0.0 prev_flops = getattr(global_state, "_flops_at_last_log", 0.0) @@ -1385,7 +1385,7 @@ def report_runtime( time_unit (str, optional): Time unit to use for `time` logging. Can be one of 'seconds', 'minutes', 'hours', or 'days'. Defaults to 'hours'. """ - elapsed_dur = train_state.step / train_iters + elapsed_dur = train_state.iteration / train_iters divider = 1 if time_unit == "seconds": @@ -1409,7 +1409,7 @@ def report_runtime( time_metrics["time/tokens"] = train_state.consumed_train_samples * seq_length time_metrics["time/samples"] = train_state.consumed_train_samples - time_metrics["time/batches"] = train_state.step + time_metrics["time/batches"] = train_state.iteration time_metrics["time/total"] = (time.time() - start_time) / divider return time_metrics @@ -1509,7 +1509,7 @@ def prepare_forward_step_func(forward_step_func: ForwardStepCallable, state: Glo Wrapping once is safe since: - functools.partial stores a reference to the state object, not a copy - - When state.train_state.step or other fields change, the partial sees those changes + - When state.train_state.iteration or other fields change, the partial sees those changes - No staleness issues because GlobalState is mutable and passed by reference Functor support: diff --git a/tests/functional_tests/test_groups/data/test_loaders.py b/tests/functional_tests/test_groups/data/test_loaders.py index e99a047431..533a9fae97 100644 --- a/tests/functional_tests/test_groups/data/test_loaders.py +++ b/tests/functional_tests/test_groups/data/test_loaders.py @@ -528,7 +528,7 @@ def test_build_data_loaders_sample_based(self, mock_broadcast, mock_get_rank, mo # Normal training state (no backward compatibility needed) train_state = TrainState() - train_state.step = 0 + train_state.iteration = 0 train_state.consumed_train_samples = 0 train_state.consumed_valid_samples = 0 diff --git a/tests/functional_tests/test_groups/megatron_mimo/test_pretrain_megatron_mimo.py b/tests/functional_tests/test_groups/megatron_mimo/test_pretrain_megatron_mimo.py index 8118e8c887..466537f9c8 100644 --- a/tests/functional_tests/test_groups/megatron_mimo/test_pretrain_megatron_mimo.py +++ b/tests/functional_tests/test_groups/megatron_mimo/test_pretrain_megatron_mimo.py @@ -245,7 +245,7 @@ def _build_data_iterators(cfg, megatron_mimo_infra, *, train_state=None, valid_s Accepts an optional ``train_state`` so consumed-sample offsets from a restored checkpoint are honored during resume. ``setup_megatron_mimo`` introspects the - signature and passes ``train_state`` when ``train_state.step > 0``. + signature and passes ``train_state`` when ``train_state.iteration > 0``. """ from megatron.bridge.data.megatron_mimo.loaders import build_megatron_mimo_data_loaders from megatron.bridge.training.state import TrainState @@ -712,7 +712,7 @@ def test_megatron_mimo_checkpoint_resume_dp1_both(self, tmp_path, save_rng): ``dist_checkpointing.save`` raises during phase 1 because each MegatronMIMO module emits an identically-keyed RNG ShardedObject. Asserts: - * ``train_state.step`` goes 0 → SAVE_STEPS after phase 1, SAVE_STEPS → TOTAL_STEPS after phase 2 + * ``train_state.iteration`` goes 0 → SAVE_STEPS after phase 1, SAVE_STEPS → TOTAL_STEPS after phase 2 * ``train_state.consumed_train_samples`` is correctly restored * Phase-1 and phase-2 consumed indices are disjoint (the resumed loader skips samples already seen before the crash) @@ -776,7 +776,7 @@ def test_megatron_mimo_checkpoint_resume_dp1_both(self, tmp_path, save_rng): ) phase1_indices = sorted(set(_RESUME_TEST_CONSUMED_INDICES)) phase1_consumed = state_save.train_state.consumed_train_samples - assert state_save.train_state.step == save_steps + assert state_save.train_state.iteration == save_steps assert phase1_consumed > 0 dist.barrier() @@ -805,8 +805,8 @@ def test_megatron_mimo_checkpoint_resume_dp1_both(self, tmp_path, save_rng): phase2_indices = sorted(set(_RESUME_TEST_CONSUMED_INDICES)) # Step counter continues from the checkpoint. - assert state_resume.train_state.step == total_steps, ( - f"Step continuity broken: phase 2 ended at step={state_resume.train_state.step}, expected {total_steps}" + assert state_resume.train_state.iteration == total_steps, ( + f"Step continuity broken: phase 2 ended at step={state_resume.train_state.iteration}, expected {total_steps}" ) # consumed_train_samples was restored and then incremented by the extra iters. diff --git a/tests/functional_tests/test_groups/quantization/test_qat_workflow.py b/tests/functional_tests/test_groups/quantization/test_qat_workflow.py index 3af80c1b22..124c80de35 100644 --- a/tests/functional_tests/test_groups/quantization/test_qat_workflow.py +++ b/tests/functional_tests/test_groups/quantization/test_qat_workflow.py @@ -19,12 +19,12 @@ from pathlib import Path import pytest +from megatron.training.utils.checkpoint_utils import get_checkpoint_train_state_filename from megatron.bridge.training.utils.checkpoint_utils import ( TRACKER_PREFIX, get_checkpoint_name, get_checkpoint_tracker_filename, - get_checkpoint_train_state_filename, ) from tests.functional_tests.utils import clear_directories diff --git a/tests/functional_tests/test_groups/training/test_local_checkpointing.py b/tests/functional_tests/test_groups/training/test_local_checkpointing.py index 47297a57d2..4e5e3ee61b 100644 --- a/tests/functional_tests/test_groups/training/test_local_checkpointing.py +++ b/tests/functional_tests/test_groups/training/test_local_checkpointing.py @@ -92,14 +92,14 @@ def __init__(self): self.steps_executed: int = 0 def on_train_start(self, context: CallbackContext) -> None: - self.start_step = context.state.train_state.step + self.start_step = context.state.train_state.iteration self.start_consumed_samples = context.state.train_state.consumed_train_samples def on_train_step_end(self, context: CallbackContext) -> None: self.steps_executed += 1 def on_train_end(self, context: CallbackContext) -> None: - self.end_step = context.state.train_state.step + self.end_step = context.state.train_state.iteration self.end_consumed_samples = context.state.train_state.consumed_train_samples diff --git a/tests/functional_tests/utils.py b/tests/functional_tests/utils.py index 2e9494ba03..498c3bfeb1 100644 --- a/tests/functional_tests/utils.py +++ b/tests/functional_tests/utils.py @@ -17,12 +17,12 @@ from pathlib import Path import torch +from megatron.training.utils.checkpoint_utils import get_checkpoint_train_state_filename from megatron.bridge.training.utils.checkpoint_utils import ( TRACKER_PREFIX, get_checkpoint_name, get_checkpoint_tracker_filename, - get_checkpoint_train_state_filename, ) diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py index 18b58a9c93..2e47732a8d 100644 --- a/tests/unit_tests/conftest.py +++ b/tests/unit_tests/conftest.py @@ -90,17 +90,15 @@ def check_gpu_requirements(request): def clear_lru_cache(): """Clear LRU cache before each test to ensure test isolation.""" # Import the functions that use @lru_cache - from megatron.bridge.training.utils.checkpoint_utils import read_run_config, read_train_state + from megatron.bridge.training.utils.checkpoint_utils import read_run_config # Clear the cache before each test read_run_config.cache_clear() - read_train_state.cache_clear() yield # Clear cache after each test as well read_run_config.cache_clear() - read_train_state.cache_clear() @pytest.fixture diff --git a/tests/unit_tests/models/qwen_omni/conftest.py b/tests/unit_tests/models/qwen_omni/conftest.py index 161a8004bf..b8186d4364 100644 --- a/tests/unit_tests/models/qwen_omni/conftest.py +++ b/tests/unit_tests/models/qwen_omni/conftest.py @@ -25,15 +25,13 @@ def clear_lru_cache(): behavior for isolation. """ try: - from megatron.bridge.training.utils.checkpoint_utils import read_run_config, read_train_state + from megatron.bridge.training.utils.checkpoint_utils import read_run_config except (ImportError, ModuleNotFoundError): yield return read_run_config.cache_clear() - read_train_state.cache_clear() yield read_run_config.cache_clear() - read_train_state.cache_clear() diff --git a/tests/unit_tests/training/megatron_mimo/test_megatron_mimo_checkpointing.py b/tests/unit_tests/training/megatron_mimo/test_megatron_mimo_checkpointing.py index 4d132fb72a..8aed9bf1d6 100644 --- a/tests/unit_tests/training/megatron_mimo/test_megatron_mimo_checkpointing.py +++ b/tests/unit_tests/training/megatron_mimo/test_megatron_mimo_checkpointing.py @@ -122,9 +122,9 @@ def _make_global_state( data_parallel_size=1, ), train_state=SimpleNamespace( - step=step, + iteration=step, consumed_train_samples=0, - floating_point_operations_so_far=0, + num_floating_point_operations_so_far=0, ), start_time=time.time(), signal_handler=Mock(), @@ -707,7 +707,7 @@ def test_exits_loop_when_checkpoint_and_decide_exit_returns_true( # Should have exited after 1 iteration, not 100 assert mock_train_step.call_count == 1 - assert state.train_state.step == 1 + assert state.train_state.iteration == 1 @patch("megatron.bridge.training.train_megatron_mimo.checkpoint_and_decide_exit", return_value=False) @patch("megatron.bridge.training.train_megatron_mimo.train_step_megatron_mimo") @@ -786,7 +786,7 @@ def test_persists_terminal_checkpoint_without_duplicate_interval_save( checkpoint_manager = MagicMock() saved_steps = [] checkpoint_manager.save.side_effect = lambda context, _callback_manager: saved_steps.append( - context.state.train_state.step + context.state.train_state.iteration ) with ( @@ -840,9 +840,9 @@ def _make_setup_output_for_load( pg_collections = {"language": Mock()} train_state = SimpleNamespace( - step=train_state_step, + iteration=train_state_step, consumed_train_samples=consumed_train_samples, - floating_point_operations_so_far=floating_point_operations_so_far, + num_floating_point_operations_so_far=floating_point_operations_so_far, ) timers_handle = Mock() timers = Mock(return_value=timers_handle) @@ -1123,7 +1123,7 @@ def test_floating_point_ops_preserved(self): setup_output=setup_output, ) _, kwargs = mocks["train_megatron_mimo"].call_args - assert kwargs["global_state"].train_state.floating_point_operations_so_far == 99999 + assert kwargs["global_state"].train_state.num_floating_point_operations_so_far == 99999 # --------------------------------------------------------------------------- @@ -1309,7 +1309,7 @@ def _run_setup(self, *, load_path=None, checkpoint_exists_return=False, checkpoi state = Mock() state.cfg = cfg state.start_time = time.time() - state.train_state = SimpleNamespace(step=0) + state.train_state = SimpleNamespace(iteration=0) mocks = {} @@ -1429,7 +1429,7 @@ class TestSetupMegatronMIMOResumeIterators: """Verify setup_megatron_mimo builds iterators with train_state when resuming.""" def test_train_state_passed_when_resuming(self): - """When train_state.step > 0, builder receives train_state kwarg.""" + """When train_state.iteration > 0, builder receives train_state kwarg.""" from megatron.bridge.training.setup_megatron_mimo import setup_megatron_mimo build_fn = MagicMock(return_value=(iter([]), None)) @@ -1470,7 +1470,7 @@ def _sig_fn(cfg, megatron_mimo_infra, *, train_state=None): state.cfg = cfg state.start_time = time.time() # Simulate resumed training — step > 0 - state.train_state = SimpleNamespace(step=10, consumed_train_samples=100) + state.train_state = SimpleNamespace(iteration=10, consumed_train_samples=100) mock_optimizer = MagicMock() mock_optimizer.module_infos = {} @@ -1517,7 +1517,7 @@ def _sig_fn(cfg, megatron_mimo_infra, *, train_state=None): assert kwargs["train_state"].step == 10 def test_no_train_state_when_not_resuming(self): - """When train_state.step == 0, builder is called without train_state.""" + """When train_state.iteration == 0, builder is called without train_state.""" from megatron.bridge.training.setup_megatron_mimo import setup_megatron_mimo build_fn = Mock(return_value=(iter([]), None)) @@ -1552,7 +1552,7 @@ def test_no_train_state_when_not_resuming(self): state = Mock() state.cfg = cfg state.start_time = time.time() - state.train_state = SimpleNamespace(step=0) + state.train_state = SimpleNamespace(iteration=0) mock_optimizer = MagicMock() mock_optimizer.module_infos = {} @@ -1637,7 +1637,7 @@ def legacy_builder(cfg, megatron_mimo_infra): state = Mock() state.cfg = cfg state.start_time = time.time() - state.train_state = SimpleNamespace(step=10) + state.train_state = SimpleNamespace(iteration=10) mock_optimizer = MagicMock() mock_optimizer.module_infos = {} diff --git a/tests/unit_tests/training/megatron_mimo/test_pretrain_megatron_mimo.py b/tests/unit_tests/training/megatron_mimo/test_pretrain_megatron_mimo.py index b2222e5fad..c1c16f661b 100644 --- a/tests/unit_tests/training/megatron_mimo/test_pretrain_megatron_mimo.py +++ b/tests/unit_tests/training/megatron_mimo/test_pretrain_megatron_mimo.py @@ -24,7 +24,7 @@ def _make_cfg(): def _make_setup_output(module_to_grid_map): global_state = MagicMock() - global_state.train_state.step = 0 + global_state.train_state.iteration = 0 mock_checkpoint_manager = MagicMock() mock_checkpoint_manager.checkpointing_context = None return SimpleNamespace( diff --git a/tests/unit_tests/training/post_training/test_checkpointing.py b/tests/unit_tests/training/post_training/test_checkpointing.py index b4c9acfa49..969789e056 100644 --- a/tests/unit_tests/training/post_training/test_checkpointing.py +++ b/tests/unit_tests/training/post_training/test_checkpointing.py @@ -397,7 +397,7 @@ def test_current_torch_dist_save_detect_restore_contract( modelopt_state_path = iteration_path / "modelopt_state" iteration_path.mkdir(parents=True) newer_iteration_path.mkdir() - torch.save(TrainState(step=100).state_dict(), checkpoint_path / "latest_train_state.pt") + torch.save(TrainState(iteration=100).state_dict(), checkpoint_path / "latest_train_state.pt") # The torch_dist strategy synchronizes CUDA even though this fixture only stores objects. with ( diff --git a/tests/unit_tests/training/test_checkpoint_resume_correctness.py b/tests/unit_tests/training/test_checkpoint_resume_correctness.py index 6bf3155518..49a26dee00 100644 --- a/tests/unit_tests/training/test_checkpoint_resume_correctness.py +++ b/tests/unit_tests/training/test_checkpoint_resume_correctness.py @@ -55,7 +55,7 @@ def test_quantized_resume_requires_successful_main_parameter_restore( cfg.ddp.fp4_param_gather = precision == "fp4" if mode == "conversion": cfg.ddp = None - state.train_state = TrainState(step=3, consumed_train_samples=48) + state.train_state = TrainState(iteration=3, consumed_train_samples=48) optimizer, scheduler = fixtures["mock_optimizer"], fixtures["mock_scheduler"] optimizer.is_stub_optimizer = mode == "stub" optimizer.param_groups = [{"max_lr": 0.2, "min_lr": 0.1}] diff --git a/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py index 250627383d..2a0d4161d1 100644 --- a/tests/unit_tests/training/test_checkpointing.py +++ b/tests/unit_tests/training/test_checkpointing.py @@ -31,6 +31,7 @@ TorchDistSaveShardedStrategy, ) from megatron.core.msc_utils import MultiStorageClientFeature +from megatron.training.utils.checkpoint_utils import get_checkpoint_train_state_filename from nvidia_resiliency_ext.checkpointing.async_ckpt.core import AsyncRequest as NVRxAsyncRequest from megatron.bridge.training.checkpointing import ( @@ -69,7 +70,6 @@ get_checkpoint_name, get_checkpoint_run_config_filename, get_checkpoint_tracker_filename, - get_checkpoint_train_state_filename, get_rng_state, init_checkpointing_context, load_checkpoint, @@ -165,18 +165,6 @@ def test_find_checkpoint_rank_0(self, mock_dist_ckpt): expected = "/checkpoints/release" assert result == expected - @pytest.mark.parametrize( - "checkpoints_path,prefix,expected", - [ - ("/checkpoints", None, "/checkpoints/train_state.pt"), - ("/checkpoints", "latest", "/checkpoints/latest_train_state.pt"), - ], - ) - def test_get_checkpoint_train_state_filename(self, checkpoints_path, prefix, expected): - """Test train state filename generation.""" - result = get_checkpoint_train_state_filename(checkpoints_path, prefix) - assert result == expected - def test_get_checkpoint_run_config_filename(self): """Test run config filename generation.""" result = get_checkpoint_run_config_filename("/checkpoints") @@ -660,7 +648,7 @@ def save_checkpoint_fixtures(): """Fixture for save checkpoint tests.""" mock_state = Mock(spec=GlobalState) mock_state.train_state = Mock(spec=TrainState) - mock_state.train_state.step = 1000 + mock_state.train_state.iteration = 1000 # Make state_dict() return a real dictionary that supports item assignment mock_state.train_state.state_dict.return_value = { "step": torch.tensor(1000), @@ -721,7 +709,7 @@ def test_metadata_failure_does_not_publish_incomplete_checkpoint(self, tmp_path, state.cfg.checkpoint.most_recent_k = -1 state.cfg.to_yaml.side_effect = OSError("metadata write failed") - previous_state = TrainState(step=500) + previous_state = TrainState(iteration=500) latest_train_state = tmp_path / "latest_train_state.pt" torch.save(previous_state.state_dict(), latest_train_state) legacy_tracker = tmp_path / "latest_checkpointed_iteration.txt" @@ -1088,7 +1076,7 @@ def test_sync_persistent_save_honors_retain_interval( current_checkpoint = tmp_path / "iter_0000030" state = save_checkpoint_fixtures["mock_state"] - state.train_state.step = 30 + state.train_state.iteration = 30 state.train_state.state_dict.return_value = {"step": torch.tensor(30)} state.cfg.checkpoint.save = str(tmp_path) state.cfg.checkpoint.async_save = False @@ -1154,7 +1142,7 @@ def test_async_persistent_save_defers_retain_interval_cleanup(self, tmp_path, sa current_checkpoint = tmp_path / "iter_0000030" state = save_checkpoint_fixtures["mock_state"] - state.train_state.step = 30 + state.train_state.iteration = 30 state.train_state.state_dict.return_value = {"step": torch.tensor(30)} state.cfg.checkpoint.save = str(tmp_path) state.cfg.checkpoint.async_save = True @@ -1301,7 +1289,7 @@ def test_async_checkpoint_loggers_use_scheduled_step(self, save_checkpoint_fixtu pg_collection=pg_collection, ) - state.train_state.step = 1001 + state.train_state.iteration = 1001 for finalize_fn in finalize_fns: finalize_fn() @@ -1334,7 +1322,7 @@ def test_async_global_non_persistent_overlap_preserves_tracker_checkpoint( torch.save({"step": torch.tensor(10)}, latest_train_state) state = save_checkpoint_fixtures["mock_state"] - state.train_state.step = 30 + state.train_state.iteration = 30 state.train_state.state_dict.return_value = {"step": torch.tensor(30)} state.cfg.checkpoint.save = str(save_dir) state.cfg.checkpoint.non_persistent_ckpt_type = "global" @@ -1429,7 +1417,7 @@ def test_sync_global_non_persistent_honors_configured_retention( future_incomplete_checkpoint.mkdir() state = save_checkpoint_fixtures["mock_state"] - state.train_state.step = 60 + state.train_state.iteration = 60 state.train_state.state_dict.return_value = {"step": torch.tensor(60)} state.cfg.checkpoint.save = str(save_dir) state.cfg.checkpoint.non_persistent_ckpt_type = "global" @@ -1964,7 +1952,7 @@ def test_load_checkpoint_from_parent_save_dir_ignores_iter_hf_export( mock_get_pg_collection.return_value = Mock() mock_check_dist_ckpt.return_value = False mock_file_exists.return_value = True - mock_read_train_state.return_value = Mock(step=1000) + mock_read_train_state.return_value = Mock(iteration=1000) mock_is_hf_checkpoint_dir.side_effect = lambda path: str(path).endswith("/hf") mock_load_base.return_value = (None, "", False, None) @@ -2168,8 +2156,8 @@ def test_torch_dist_omitted_optimizer_keeps_load_state_fresh(self, load_checkpoi cfg.peft = None state = load_checkpoint_fixtures["mock_state"] - state.train_state.step = 0 - state.train_state.floating_point_operations_so_far = 0 + state.train_state.iteration = 0 + state.train_state.num_floating_point_operations_so_far = 0 optimizer = load_checkpoint_fixtures["mock_optimizer"] optimizer.is_stub_optimizer = False @@ -2360,8 +2348,8 @@ def test_load_checkpoint_found( mock_exists_os.return_value = True # train_state.pt exists (normal case) mock_train_state = Mock() - mock_train_state.step = 1000 - mock_train_state.floating_point_operations_so_far = 500000 + mock_train_state.iteration = 1000 + mock_train_state.num_floating_point_operations_so_far = 500000 mock_read_state.return_value = mock_train_state # Mock utility functions @@ -2462,7 +2450,7 @@ def test_get_non_persistent_iteration_global(self, mock_isfile, mock_read_state) mock_isfile.return_value = True mock_train_state = Mock() - mock_train_state.step = 1500 + mock_train_state.iteration = 1500 mock_read_state.return_value = mock_train_state result = _get_non_persistent_iteration("/np_dir", non_persistent_ckpt_type) @@ -2786,7 +2774,7 @@ def test_newer_global_non_persistent_checkpoint_is_loaded_from_distinct_save_dir non_persistent_dir = save_dir / "non_persistent" non_persistent_dir.mkdir(parents=True) torch.save( - TrainState(step=200).state_dict(), + TrainState(iteration=200).state_dict(), get_checkpoint_train_state_filename(str(non_persistent_dir), prefix="latest"), ) @@ -2830,7 +2818,7 @@ def test_ckpt_step_overrides_newer_non_persistent_checkpoint( non_persistent_dir = tmp_path / "recovery" non_persistent_dir.mkdir(parents=True) torch.save( - TrainState(step=200).state_dict(), + TrainState(iteration=200).state_dict(), get_checkpoint_train_state_filename(str(non_persistent_dir), prefix="latest"), ) @@ -2879,11 +2867,11 @@ def test_newer_global_non_persistent_checkpoint_under_load_remains_available( load_non_persistent_dir.mkdir(parents=True) save_non_persistent_dir.mkdir(parents=True) torch.save( - TrainState(step=200).state_dict(), + TrainState(iteration=200).state_dict(), get_checkpoint_train_state_filename(str(load_non_persistent_dir), prefix="latest"), ) torch.save( - TrainState(step=100).state_dict(), + TrainState(iteration=100).state_dict(), get_checkpoint_train_state_filename(str(save_non_persistent_dir), prefix="latest"), ) @@ -2944,7 +2932,7 @@ def test_load_base_checkpoint_non_distributed_error( mock_file_exists.return_value = True # train_state file exists mock_train_state = Mock() - mock_train_state.step = 1000 + mock_train_state.iteration = 1000 mock_read_state.return_value = mock_train_state mock_dist_ckpt.check_is_distributed_checkpoint.return_value = False @@ -3997,11 +3985,11 @@ def mock_exists_side_effect(path): # Verify that the legacy train state was created correctly train_state = mock_state.train_state - assert train_state.step == 2000 + assert train_state.iteration == 2000 assert train_state.consumed_train_samples == 100000 assert train_state.skipped_train_samples == 50 assert train_state.consumed_valid_samples == 10000 - assert train_state.floating_point_operations_so_far == 5000000 + assert train_state.num_floating_point_operations_so_far == 5000000 assert train_state.do_train is False assert train_state.do_valid is False assert train_state.do_test is False @@ -4037,7 +4025,7 @@ def test_get_train_state_complete_state_dict(self): assert result.consumed_train_samples == 150000 assert result.skipped_train_samples == 250 assert result.consumed_valid_samples == 12000 - assert result.floating_point_operations_so_far == 7500000 + assert result.num_floating_point_operations_so_far == 7500000 assert result.do_train is False assert result.do_valid is False assert result.do_test is False @@ -4064,7 +4052,7 @@ def test_get_train_state_missing_iteration(self): assert result.consumed_train_samples == 100000 assert result.skipped_train_samples == 50 assert result.consumed_valid_samples == 8000 - assert result.floating_point_operations_so_far == 5000000 + assert result.num_floating_point_operations_so_far == 5000000 def test_get_train_state_missing_args(self): """Test creating TrainState when args is missing.""" @@ -4083,7 +4071,7 @@ def test_get_train_state_missing_args(self): assert result.consumed_train_samples == 0 # fallback assert result.skipped_train_samples == 0 # fallback assert result.consumed_valid_samples == 0 # fallback - assert result.floating_point_operations_so_far == 4000000 + assert result.num_floating_point_operations_so_far == 4000000 def test_get_train_state_missing_flops(self): """Test creating TrainState when floating point operations count is missing.""" @@ -4107,7 +4095,7 @@ def test_get_train_state_missing_flops(self): assert result.consumed_train_samples == 75000 assert result.skipped_train_samples == 30 assert result.consumed_valid_samples == 6000 - assert result.floating_point_operations_so_far == 0 # default + assert result.num_floating_point_operations_so_far == 0 # default def test_get_train_state_partial_args(self): """Test creating TrainState when args has only some attributes.""" @@ -4136,7 +4124,7 @@ def __init__(self): assert result.consumed_train_samples == 200000 # from args assert result.skipped_train_samples == 0 # default from getattr assert result.consumed_valid_samples == 0 # default from getattr - assert result.floating_point_operations_so_far == 9000000 + assert result.num_floating_point_operations_so_far == 9000000 def test_get_train_state_empty_state_dict(self): """Test creating TrainState from an empty state_dict.""" @@ -4151,7 +4139,7 @@ def test_get_train_state_empty_state_dict(self): assert result.consumed_train_samples == 0 assert result.skipped_train_samples == 0 assert result.consumed_valid_samples == 0 - assert result.floating_point_operations_so_far == 0 + assert result.num_floating_point_operations_so_far == 0 assert result.do_train is False assert result.do_valid is False assert result.do_test is False @@ -4173,7 +4161,7 @@ def test_get_train_state_args_none(self): assert result.consumed_train_samples == 0 # fallback assert result.skipped_train_samples == 0 # fallback assert result.consumed_valid_samples == 0 # fallback - assert result.floating_point_operations_so_far == 1000000 + assert result.num_floating_point_operations_so_far == 1000000 def test_get_train_state_large_values(self): """Test creating TrainState with large numerical values.""" @@ -4197,7 +4185,7 @@ def test_get_train_state_large_values(self): assert result.consumed_train_samples == 999999999 assert result.skipped_train_samples == 1000000 assert result.consumed_valid_samples == 50000000 - assert result.floating_point_operations_so_far == 999999999999 + assert result.num_floating_point_operations_so_far == 999999999999 def test_get_train_state_zero_values(self): """Test creating TrainState with zero values.""" @@ -4221,7 +4209,7 @@ def test_get_train_state_zero_values(self): assert result.consumed_train_samples == 0 assert result.skipped_train_samples == 0 assert result.consumed_valid_samples == 0 - assert result.floating_point_operations_so_far == 0 + assert result.num_floating_point_operations_so_far == 0 # Boolean flags should still be False assert result.do_train is False assert result.do_valid is False @@ -4255,11 +4243,11 @@ def test_inprocess_restart_refreshes_mutable_bridge_tracker(self, tmp_path): from megatron.bridge.training.checkpointing import _resolve_checkpoint_iteration tracker = get_checkpoint_train_state_filename(str(tmp_path), prefix="latest") - first_state = TrainState(step=5) + first_state = TrainState(iteration=5) torch.save(first_state.state_dict(), tracker) assert _resolve_checkpoint_iteration(str(tmp_path), None) == (5, False) - latest_state = TrainState(step=10) + latest_state = TrainState(iteration=10) torch.save(latest_state.state_dict(), tracker) GlobalState().reset_for_restart() @@ -4273,7 +4261,7 @@ def test_loads_from_bridge_tracker(self, mock_read_train_state, mock_file_exists mock_file_exists.return_value = True train_state = TrainState() - train_state.step = 1000 + train_state.iteration = 1000 mock_read_train_state.return_value = train_state iteration, release = _resolve_checkpoint_iteration( @@ -4390,7 +4378,7 @@ def test_ignore_ckpt_step_loads_latest(self, mock_read_train_state, mock_file_ex mock_file_exists.return_value = True train_state = TrainState() - train_state.step = 9999 # Latest checkpoint + train_state.iteration = 9999 # Latest checkpoint mock_read_train_state.return_value = train_state # Pass None to _resolve_checkpoint_iteration (simulating ignore_ckpt_step=True) @@ -5025,7 +5013,7 @@ def test_fsdp_dtensor_ckpt_step_prepares_requested_iteration( MultiStorageClientFeature.disable() mock_is_iter_dir.return_value = False mock_file_exists.side_effect = lambda path: path.endswith("latest_train_state.pt") - mock_read_train_state.return_value = TrainState(step=200) + mock_read_train_state.return_value = TrainState(iteration=200) mock_metadata = Mock() mock_metadata.state_dict_metadata = {} @@ -5716,8 +5704,8 @@ def test_load_local_checkpoint_calls_load_state_dict_from_file( mock_data_distribution_group.rank.return_value = 5 mock_train_state = Mock() - mock_train_state.step = 500 - mock_train_state.floating_point_operations_so_far = 0 + mock_train_state.iteration = 500 + mock_train_state.num_floating_point_operations_so_far = 0 mock_read_state.return_value = mock_train_state mock_read_config.return_value = { @@ -5837,8 +5825,8 @@ def test_load_global_checkpoint_uses_standard_load_state_dict( mock_get_pg_collection.return_value = mock_pg_collection mock_train_state = Mock() - mock_train_state.step = 1000 - mock_train_state.floating_point_operations_so_far = 0 + mock_train_state.iteration = 1000 + mock_train_state.num_floating_point_operations_so_far = 0 mock_read_state.return_value = mock_train_state mock_read_config.return_value = { diff --git a/tests/unit_tests/training/test_eval.py b/tests/unit_tests/training/test_eval.py index 5a7d30d001..5c03f2668b 100644 --- a/tests/unit_tests/training/test_eval.py +++ b/tests/unit_tests/training/test_eval.py @@ -30,7 +30,7 @@ def _make_state(): wandb_logger=None, mlflow_logger=None, comet_logger=None, - train_state=SimpleNamespace(step=0, consumed_train_samples=0), + train_state=SimpleNamespace(iteration=0, consumed_train_samples=0), cfg=SimpleNamespace(logger=SimpleNamespace(log_validation_ppl_to_tensorboard=False)), ) diff --git a/tests/unit_tests/training/test_functor_support.py b/tests/unit_tests/training/test_functor_support.py index 3fd7798c91..a646f7c2f8 100644 --- a/tests/unit_tests/training/test_functor_support.py +++ b/tests/unit_tests/training/test_functor_support.py @@ -222,7 +222,7 @@ def test_pretrain_with_four_arg_functor( setup_output.state = MagicMock() setup_output.state.cfg = container setup_output.state.train_state.do_train = True - setup_output.state.train_state.step = 0 + setup_output.state.train_state.iteration = 0 setup_output.state.train_state.do_valid = False setup_output.state.train_state.do_test = False @@ -278,7 +278,7 @@ def test_pretrain_with_stateful_functor( setup_output.state = MagicMock() setup_output.state.cfg = container setup_output.state.train_state.do_train = True - setup_output.state.train_state.step = 0 + setup_output.state.train_state.iteration = 0 setup_output.state.train_state.do_valid = False setup_output.state.train_state.do_test = False diff --git a/tests/unit_tests/training/test_peft_checkpointing.py b/tests/unit_tests/training/test_peft_checkpointing.py index e56e2e99b1..9ab6a784ca 100644 --- a/tests/unit_tests/training/test_peft_checkpointing.py +++ b/tests/unit_tests/training/test_peft_checkpointing.py @@ -540,8 +540,8 @@ def test_load_checkpoint_peft_resume_detection( mock_state.train_state.consumed_train_samples = 0 mock_state.train_state.skipped_train_samples = 0 mock_state.train_state.consumed_valid_samples = 0 - mock_state.train_state.step = 1000 # Set to integer for comparisons - mock_state.train_state.floating_point_operations_so_far = 50000 + mock_state.train_state.iteration = 1000 # Set to integer for comparisons + mock_state.train_state.num_floating_point_operations_so_far = 50000 mock_cfg.ddp = Mock() mock_cfg.ddp.use_megatron_fsdp = False @@ -666,8 +666,8 @@ def test_load_checkpoint_non_peft_regular_loading(self, mock_dist_ckpt, mock_che mock_state.train_state.consumed_train_samples = 0 mock_state.train_state.skipped_train_samples = 0 mock_state.train_state.consumed_valid_samples = 0 - mock_state.train_state.step = 1000 # Set to integer for comparisons - mock_state.train_state.floating_point_operations_so_far = 50000 + mock_state.train_state.iteration = 1000 # Set to integer for comparisons + mock_state.train_state.num_floating_point_operations_so_far = 50000 mock_cfg.ddp = Mock() mock_cfg.ddp.use_megatron_fsdp = False @@ -807,8 +807,8 @@ def test_load_checkpoint_peft_resume_multi_model( mock_state.train_state.consumed_train_samples = 0 mock_state.train_state.skipped_train_samples = 0 mock_state.train_state.consumed_valid_samples = 0 - mock_state.train_state.step = 1000 # Set to integer for comparisons - mock_state.train_state.floating_point_operations_so_far = 50000 + mock_state.train_state.iteration = 1000 # Set to integer for comparisons + mock_state.train_state.num_floating_point_operations_so_far = 50000 mock_cfg.ddp = Mock() mock_cfg.ddp.use_megatron_fsdp = False diff --git a/tests/unit_tests/training/test_pg_collection_wiring.py b/tests/unit_tests/training/test_pg_collection_wiring.py index f323c453ff..302b4aead3 100644 --- a/tests/unit_tests/training/test_pg_collection_wiring.py +++ b/tests/unit_tests/training/test_pg_collection_wiring.py @@ -70,7 +70,7 @@ def _get_data_distribution_group(pg_collection, received_model_config): monkeypatch.setattr(train_module, "_dummy_train_step", lambda *args, **kwargs: None) # Pre-check counters - assert state.train_state.step == 0 + assert state.train_state.iteration == 0 assert state.train_state.consumed_train_samples == 0 assert state.train_state.skipped_train_samples == 0 @@ -81,7 +81,7 @@ def _get_data_distribution_group(pg_collection, received_model_config): assert did_skip is True assert group_calls == [(fake_pg, model_config)] # One iteration skipped - assert state.train_state.step == 1 + assert state.train_state.iteration == 1 # Batch size = dp.size * micro_batch_size * num_microbatches = 3 * 4 * 2 = 24 expected_batch = 3 * 4 * 2 assert state.train_state.consumed_train_samples == expected_batch @@ -130,10 +130,10 @@ def test_train_stops_nsys_profiler_when_skipped_iteration_reaches_profile_end(mo state = SimpleNamespace( cfg=config, train_state=SimpleNamespace( - step=0, + iteration=0, consumed_train_samples=0, skipped_train_samples=0, - floating_point_operations_so_far=0, + num_floating_point_operations_so_far=0, ), timers=Mock(), straggler_timer=Mock(), @@ -236,10 +236,10 @@ def test_train_stops_nsys_profiler_when_rerun_requests_exit(monkeypatch): state = SimpleNamespace( cfg=config, train_state=SimpleNamespace( - step=0, + iteration=0, consumed_train_samples=0, skipped_train_samples=0, - floating_point_operations_so_far=0, + num_floating_point_operations_so_far=0, ), timers=Mock(), straggler_timer=Mock(), @@ -347,10 +347,10 @@ def test_first_skipped_iteration_preserves_cuda_graph_hook_bootstrap(monkeypatch state = SimpleNamespace( cfg=config, train_state=SimpleNamespace( - step=0, + iteration=0, consumed_train_samples=0, skipped_train_samples=0, - floating_point_operations_so_far=0, + num_floating_point_operations_so_far=0, do_valid=False, ), timers=Mock(), @@ -451,6 +451,6 @@ def enable_forward_pre_hook(*args, **kwargs): pg_collection=pg_collection, ) - assert state.train_state.step == 3 + assert state.train_state.iteration == 3 assert graph_state["created"] is True helper.cuda_graph_set_manual_hooks.assert_called_once() diff --git a/tests/unit_tests/training/test_runtime_callbacks.py b/tests/unit_tests/training/test_runtime_callbacks.py index 5e5a4192ee..0e98461da3 100644 --- a/tests/unit_tests/training/test_runtime_callbacks.py +++ b/tests/unit_tests/training/test_runtime_callbacks.py @@ -106,7 +106,7 @@ def _run_setup( start_time=0.0, tensorboard_logger=None, timers=Mock(return_value=timer), - train_state=SimpleNamespace(step=1), + train_state=SimpleNamespace(iteration=1), wandb_logger=None, ) pg_collection = SimpleNamespace( diff --git a/tests/unit_tests/training/test_setup.py b/tests/unit_tests/training/test_setup.py index ca8d52a4a8..9922f0e807 100644 --- a/tests/unit_tests/training/test_setup.py +++ b/tests/unit_tests/training/test_setup.py @@ -237,7 +237,7 @@ def test_data_init_warmup_preserves_checkpoint_restored_rng_state(): start_time=0.0, tensorboard_logger=None, timers=Mock(return_value=timer), - train_state=SimpleNamespace(step=1), + train_state=SimpleNamespace(iteration=1), wandb_logger=None, ) pg_collection = SimpleNamespace(dp=object()) diff --git a/tests/unit_tests/training/test_state.py b/tests/unit_tests/training/test_state.py index 7e947b290c..458cfc400c 100644 --- a/tests/unit_tests/training/test_state.py +++ b/tests/unit_tests/training/test_state.py @@ -19,7 +19,6 @@ from unittest.mock import MagicMock, patch import pytest -import torch from megatron.bridge.training.state import FaultToleranceState, GlobalState, TrainState @@ -39,167 +38,6 @@ def _send_sigterm_with_handler_disabled() -> None: os.kill(os.getpid(), signal.SIGTERM) -class TestTrainState: - """Test suite for TrainState class.""" - - def test_initialization_defaults(self): - """Test that TrainState initializes with correct default values.""" - state = TrainState() - - assert state.step == 0 - assert state.consumed_train_samples == 0 - assert state.skipped_train_samples == 0 - assert state.consumed_valid_samples == 0 - assert state.floating_point_operations_so_far == 0 - assert state.do_train is False - assert state.do_valid is False - assert state.do_test is False - - def test_initialization_with_custom_values(self): - """Test initialization with custom values.""" - state = TrainState( - step=100, - consumed_train_samples=1000, - skipped_train_samples=50, - consumed_valid_samples=200, - floating_point_operations_so_far=1500, - do_train=True, - do_valid=True, - do_test=False, - ) - - assert state.step == 100 - assert state.consumed_train_samples == 1000 - assert state.skipped_train_samples == 50 - assert state.consumed_valid_samples == 200 - assert state.floating_point_operations_so_far == 1500 - assert state.do_train is True - assert state.do_valid is True - assert state.do_test is False - - def test_state_dict_structure_and_types(self): - """Test that state_dict returns correct structure and tensor types.""" - state = TrainState( - step=42, - consumed_train_samples=500, - skipped_train_samples=25, - consumed_valid_samples=100, - floating_point_operations_so_far=2000, - do_train=True, - do_valid=False, - do_test=True, - ) - - state_dict = state.state_dict() - - # Check all expected keys are present - expected_keys = { - "step", - "consumed_train_samples", - "skipped_train_samples", - "consumed_valid_samples", - "floating_point_operations_so_far", - "do_train", - "do_valid", - "do_test", - } - assert set(state_dict.keys()) == expected_keys - - # Check tensor types (int64 to avoid overflow on long training runs; - # torch.uint64 is not supported by torch's legacy pickle path used inside - # broadcast_object_list during dist-checkpoint save, so use int64 instead). - assert state_dict["step"].dtype == torch.int64 - assert state_dict["consumed_train_samples"].dtype == torch.int64 - assert state_dict["skipped_train_samples"].dtype == torch.int64 - assert state_dict["consumed_valid_samples"].dtype == torch.int64 - assert state_dict["floating_point_operations_so_far"].dtype == torch.float64 - assert state_dict["do_train"].dtype == torch.bool - assert state_dict["do_valid"].dtype == torch.bool - assert state_dict["do_test"].dtype == torch.bool - - # Check values - assert state_dict["step"].item() == 42 - assert state_dict["consumed_train_samples"].item() == 500 - assert state_dict["skipped_train_samples"].item() == 25 - assert state_dict["consumed_valid_samples"].item() == 100 - assert state_dict["floating_point_operations_so_far"].item() == 2000 - assert state_dict["do_train"].item() is True - assert state_dict["do_valid"].item() is False - assert state_dict["do_test"].item() is True - - def test_load_state_dict(self): - """Test loading state from state dictionary.""" - # Create a state dict manually - state_dict = { - "step": torch.tensor(75, dtype=torch.int32), - "consumed_train_samples": torch.tensor(800, dtype=torch.int32), - "skipped_train_samples": torch.tensor(40, dtype=torch.int32), - "consumed_valid_samples": torch.tensor(150, dtype=torch.int32), - "floating_point_operations_so_far": torch.tensor(3000.5, dtype=torch.float64), - "do_train": torch.tensor(False, dtype=torch.bool), - "do_valid": torch.tensor(True, dtype=torch.bool), - "do_test": torch.tensor(False, dtype=torch.bool), - } - - state = TrainState() - state.load_state_dict(state_dict) - - assert state.step == 75 - assert state.consumed_train_samples == 800 - assert state.skipped_train_samples == 40 - assert state.consumed_valid_samples == 150 - assert state.floating_point_operations_so_far == 3000.5 - assert state.do_train is False - assert state.do_valid is True - assert state.do_test is False - - def test_round_trip_serialization(self): - """Test that state_dict -> load_state_dict preserves all values.""" - # Create original state with various values - original_state = TrainState( - step=123, - consumed_train_samples=2500, - skipped_train_samples=100, - consumed_valid_samples=500, - floating_point_operations_so_far=12345.67, - do_train=True, - do_valid=False, - do_test=True, - ) - - # Serialize to state dict - state_dict = original_state.state_dict() - - # Create new state and load from dict - new_state = TrainState() - new_state.load_state_dict(state_dict) - - # Verify all values are preserved - assert new_state.step == original_state.step - assert new_state.consumed_train_samples == original_state.consumed_train_samples - assert new_state.skipped_train_samples == original_state.skipped_train_samples - assert new_state.consumed_valid_samples == original_state.consumed_valid_samples - assert new_state.floating_point_operations_so_far == original_state.floating_point_operations_so_far - assert new_state.do_train == original_state.do_train - assert new_state.do_valid == original_state.do_valid - assert new_state.do_test == original_state.do_test - - def test_state_dict_with_defaults(self): - """Test state_dict with default values.""" - state = TrainState() - state_dict = state.state_dict() - - # All default values should be properly serialized - assert state_dict["step"].item() == 0 - assert state_dict["consumed_train_samples"].item() == 0 - assert state_dict["skipped_train_samples"].item() == 0 - assert state_dict["consumed_valid_samples"].item() == 0 - assert state_dict["floating_point_operations_so_far"].item() == 0 - assert state_dict["do_train"].item() is False - assert state_dict["do_valid"].item() is False - assert state_dict["do_test"].item() is False - - class TestFaultToleranceState: """Test suite for FaultToleranceState class.""" @@ -511,7 +349,7 @@ def test_train_state_property_lazy_initialization(self): def test_train_state_property_setter(self): """Test train_state property setter.""" state = GlobalState() - custom_train_state = TrainState(step=100, consumed_train_samples=1000) + custom_train_state = TrainState(iteration=100, consumed_train_samples=1000) state.train_state = custom_train_state diff --git a/tests/unit_tests/training/test_train.py b/tests/unit_tests/training/test_train.py index 1d597cd526..75626e7a85 100644 --- a/tests/unit_tests/training/test_train.py +++ b/tests/unit_tests/training/test_train.py @@ -93,7 +93,7 @@ def finalize_async_saves(self, state, blocking=False, terminate=False): ) state = SimpleNamespace( cfg=config, - train_state=SimpleNamespace(step=1, floating_point_operations_so_far=0.0), + train_state=SimpleNamespace(iteration=1, num_floating_point_operations_so_far=0.0), timers=Mock(return_value=Mock()), straggler_timer=Mock(), energy_monitor=None, @@ -599,9 +599,9 @@ def _make_state(self) -> tuple[SimpleNamespace, Mock]: checkpoint=SimpleNamespace(async_save=False), ), train_state=SimpleNamespace( - floating_point_operations_so_far=0, + num_floating_point_operations_so_far=0, consumed_train_samples=0, - step=1, + iteration=1, ), start_time=0, ) @@ -683,7 +683,7 @@ def _create_mock_state( mock_state.cfg.checkpoint.non_persistent_save_interval = None # Mock train state - mock_state.train_state.step = step + mock_state.train_state.iteration = step # Mock start_time mock_state.start_time = start_time if start_time is not None else time.time() @@ -1273,7 +1273,7 @@ class TestIterationSkipping: def _create_mock_global_state(self, step=0, iterations_to_skip=None, micro_batch_size=4): """Helper method to create a mock global state.""" mock_state = Mock() - mock_state.train_state.step = step + mock_state.train_state.iteration = step mock_state.train_state.consumed_train_samples = 0 mock_state.train_state.skipped_train_samples = 0 @@ -1314,7 +1314,7 @@ def test_should_skip_iteration_when_step_in_skip_list(self, mock_get_microbatche mock_dummy_step.assert_called_once_with(global_state, train_data_iterator, fake_pg) # Verify state updates - assert global_state.train_state.step == 5 # incremented + assert global_state.train_state.iteration == 5 # incremented expected_batch_size = 2 * 4 * 4 # dp_world_size * micro_batch_size * num_microbatches assert global_state.train_state.consumed_train_samples == expected_batch_size assert global_state.train_state.skipped_train_samples == expected_batch_size @@ -1335,7 +1335,7 @@ def test_should_not_skip_iteration_when_step_not_in_skip_list(self, mock_dummy_s mock_dummy_step.assert_not_called() # Verify state not modified - assert global_state.train_state.step == 7 # unchanged + assert global_state.train_state.iteration == 7 # unchanged assert global_state.train_state.consumed_train_samples == 0 # unchanged assert global_state.train_state.skipped_train_samples == 0 # unchanged diff --git a/tests/unit_tests/training/utils/test_checkpoint_utils.py b/tests/unit_tests/training/utils/test_checkpoint_utils.py index 92f897dec4..546be10bee 100644 --- a/tests/unit_tests/training/utils/test_checkpoint_utils.py +++ b/tests/unit_tests/training/utils/test_checkpoint_utils.py @@ -15,72 +15,24 @@ import os import threading import time -from dataclasses import dataclass -from typing import Any, Dict from unittest.mock import mock_open, patch import pytest import yaml +from megatron.training.utils.checkpoint_utils import TRAIN_STATE_FILE from megatron.bridge.training.utils.checkpoint_utils import ( CONFIG_FILE, TRACKER_PREFIX, - TRAIN_STATE_FILE, checkpoint_exists, get_checkpoint_run_config_filename, - get_checkpoint_train_state_filename, get_hf_model_id_from_checkpoint, is_checkpoint_iteration_directory, is_hf_checkpoint_dir, read_run_config, - read_train_state, ) -@dataclass -class MockTrainState: - """Mock train state class for testing.""" - - iteration: int = 0 - epoch: int = 0 - step: int = 0 - - def load_state_dict(self, state_dict: Dict[str, Any]) -> None: - """Load state from dictionary.""" - self.iteration = state_dict.get("iteration", 0) - self.epoch = state_dict.get("epoch", 0) - self.step = state_dict.get("step", 0) - - -@dataclass -class ComplexTrainState: - """More complex train state class for advanced testing.""" - - iteration: int = 0 - epoch: int = 0 - step: int = 0 - learning_rate: float = 0.0 - loss: float = 0.0 - metrics: Dict[str, float] = None - optimizer_state: Dict[str, Any] = None - - def __post_init__(self): - if self.metrics is None: - self.metrics = {} - if self.optimizer_state is None: - self.optimizer_state = {} - - def load_state_dict(self, state_dict: Dict[str, Any]) -> None: - """Load state from dictionary.""" - self.iteration = state_dict.get("iteration", 0) - self.epoch = state_dict.get("epoch", 0) - self.step = state_dict.get("step", 0) - self.learning_rate = state_dict.get("learning_rate", 0.0) - self.loss = state_dict.get("loss", 0.0) - self.metrics = state_dict.get("metrics", {}) - self.optimizer_state = state_dict.get("optimizer_state", {}) - - class TestCheckpointUtils: """Test suite for checkpoint utility functions.""" @@ -164,23 +116,6 @@ def test_get_hf_model_id_from_checkpoint_invalid_path(self, tmp_path): with pytest.raises(NotADirectoryError): get_hf_model_id_from_checkpoint(file_path) - def test_get_checkpoint_train_state_filename_without_prefix(self, tmp_path): - """Test get_checkpoint_train_state_filename without prefix.""" - checkpoint_dir = str(tmp_path / "checkpoints") - expected_path = os.path.join(checkpoint_dir, TRAIN_STATE_FILE) - - result = get_checkpoint_train_state_filename(checkpoint_dir) - assert result == expected_path - - def test_get_checkpoint_train_state_filename_with_prefix(self, tmp_path): - """Test get_checkpoint_train_state_filename with prefix.""" - checkpoint_dir = str(tmp_path / "checkpoints") - prefix = "custom_prefix" - expected_path = os.path.join(checkpoint_dir, f"{prefix}_{TRAIN_STATE_FILE}") - - result = get_checkpoint_train_state_filename(checkpoint_dir, prefix) - assert result == expected_path - @patch("megatron.bridge.training.utils.checkpoint_utils.get_rank_safe") @patch("megatron.bridge.training.utils.checkpoint_utils.torch.distributed.is_initialized") def test_read_run_config_rank_0_success(self, mock_is_initialized, mock_get_rank): @@ -276,73 +211,6 @@ def test_read_run_config_sanitizes_runtime_only_targets(self, mock_is_initialize assert result["model"]["keep"]["_target_"] == "some.other.Component" assert result["model"]["nested"][1]["other"]["_target_"] == "another.Component" - @patch("megatron.bridge.training.utils.checkpoint_utils.get_rank_safe") - @patch("megatron.bridge.training.utils.checkpoint_utils.torch.distributed.is_initialized") - @patch("megatron.bridge.training.utils.checkpoint_utils.torch.load") - def test_read_train_state_rank_0_success(self, mock_torch_load, mock_is_initialized, mock_get_rank): - """Test read_train_state successful read on rank 0.""" - # Setup mocks - mock_get_rank.return_value = 0 - mock_is_initialized.return_value = False - - # Mock train state data - state_dict = {"iteration": 100, "epoch": 5, "step": 1000} - mock_torch_load.return_value = state_dict - - with patch("megatron.bridge.training.utils.checkpoint_utils.TrainState", return_value=MockTrainState()): - result = read_train_state("train_state.pt") - - assert isinstance(result, MockTrainState) - assert result.iteration == 100 - assert result.epoch == 5 - assert result.step == 1000 - mock_torch_load.assert_called_once_with("train_state.pt", map_location="cpu", weights_only=True) - - @patch("megatron.bridge.training.utils.checkpoint_utils.get_rank_safe") - @patch("megatron.bridge.training.utils.checkpoint_utils.get_world_size_safe") - @patch("megatron.bridge.training.utils.checkpoint_utils.torch.distributed.is_initialized") - @patch("megatron.bridge.training.utils.checkpoint_utils.torch.distributed.broadcast_object_list") - @patch("torch.distributed.get_rank") # Mock the direct torch.distributed.get_rank call - def test_read_train_state_distributed_success( - self, mock_torch_get_rank, mock_broadcast, mock_is_initialized, mock_get_world_size, mock_get_rank - ): - """Test read_train_state with distributed broadcasting.""" - # Setup mocks for distributed scenario - rank = 2 # Non-rank 0 - mock_get_rank.return_value = rank - mock_torch_get_rank.return_value = rank # Mock torch.distributed.get_rank as well - mock_get_world_size.return_value = 4 - mock_is_initialized.return_value = True - - # Mock the broadcast to simulate receiving train state from rank 0 - train_state = MockTrainState() - train_state.iteration = 200 - train_state.epoch = 10 - - def broadcast_side_effect(obj_list, src): - obj_list[0] = train_state - - mock_broadcast.side_effect = broadcast_side_effect - - result = read_train_state("train_state.pt") - - assert result == train_state - assert result.iteration == 200 - assert result.epoch == 10 - mock_broadcast.assert_called_once() - - @patch("megatron.bridge.training.utils.checkpoint_utils.get_rank_safe") - @patch("megatron.bridge.training.utils.checkpoint_utils.torch.distributed.is_initialized") - @patch("megatron.bridge.training.utils.checkpoint_utils.torch.load") - def test_read_train_state_load_error(self, mock_torch_load, mock_is_initialized, mock_get_rank): - """Test read_train_state handles torch.load error.""" - mock_get_rank.return_value = 0 - mock_is_initialized.return_value = False - mock_torch_load.side_effect = RuntimeError("Corrupted file") - - with pytest.raises(RuntimeError, match="Unable to load train state file"): - read_train_state("corrupted.pt") - def test_caching_behavior_read_run_config(self): """Test that read_run_config uses caching properly.""" with ( @@ -364,33 +232,8 @@ def test_caching_behavior_read_run_config(self): # File should only be opened once due to caching assert mock_file.call_count == 1 - def test_caching_behavior_read_train_state(self): - """Test that read_train_state uses caching properly.""" - with ( - patch("megatron.bridge.training.utils.checkpoint_utils.get_rank_safe", return_value=0), - patch( - "megatron.bridge.training.utils.checkpoint_utils.torch.distributed.is_initialized", return_value=False - ), - patch("megatron.bridge.training.utils.checkpoint_utils.torch.load") as mock_load, - ): - state_dict = {"iteration": 100, "epoch": 5} - mock_load.return_value = state_dict - - # First call - with patch("megatron.bridge.training.utils.checkpoint_utils.TrainState", return_value=MockTrainState()): - result1 = read_train_state("train_state.pt") - # Second call should use cache - with patch("megatron.bridge.training.utils.checkpoint_utils.TrainState", return_value=MockTrainState()): - result2 = read_train_state("train_state.pt") - - assert result1 == result2 - assert result1.iteration == 100 - # torch.load should only be called once due to caching - assert mock_load.call_count == 1 - def test_constants_are_correct(self): """Test that module constants have expected values.""" - assert TRAIN_STATE_FILE == "train_state.pt" assert TRACKER_PREFIX == "latest" assert CONFIG_FILE == "run_config.yaml" @@ -509,39 +352,6 @@ def read_config_worker(): assert len(results) == 10 assert all(result == config_data for result in results) - def test_memory_usage_with_complex_train_state(self): - """Test memory efficiency with complex train state objects.""" - # Create a complex state with large nested data - complex_state_dict = { - "iteration": 10000, - "epoch": 50, - "step": 100000, - "learning_rate": 0.0001, - "loss": 1.23, - "metrics": {f"metric_{i}": i * 0.1 for i in range(1000)}, - "optimizer_state": { - "param_groups": [{"params": list(range(10000))} for _ in range(10)], - "state": {i: {"momentum": [0.1] * 100} for i in range(1000)}, - }, - } - - with ( - patch("megatron.bridge.training.utils.checkpoint_utils.get_rank_safe", return_value=0), - patch( - "megatron.bridge.training.utils.checkpoint_utils.torch.distributed.is_initialized", return_value=False - ), - patch("megatron.bridge.training.utils.checkpoint_utils.torch.load", return_value=complex_state_dict), - patch("megatron.bridge.training.utils.checkpoint_utils.TrainState", return_value=ComplexTrainState()), - ): - result = read_train_state("complex_state.pt") - - # Verify the complex state is loaded correctly - assert result.iteration == 10000 - assert result.epoch == 50 - assert len(result.metrics) == 1000 - assert len(result.optimizer_state["param_groups"]) == 10 - assert len(result.optimizer_state["state"]) == 1000 - def test_stress_test_many_different_files(self, tmp_path): """Stress test with many different config files to test cache behavior.""" num_files = 50 @@ -590,19 +400,6 @@ def test_permission_error_handling(self, mock_is_initialized, mock_get_rank): with pytest.raises(RuntimeError, match="Unable to load config file"): read_run_config("restricted_config.yaml") - @patch("megatron.bridge.training.utils.checkpoint_utils.get_rank_safe") - @patch("megatron.bridge.training.utils.checkpoint_utils.torch.distributed.is_initialized") - @patch("megatron.bridge.training.utils.checkpoint_utils.torch.load") - def test_out_of_memory_error_handling(self, mock_torch_load, mock_is_initialized, mock_get_rank): - """Test handling of out-of-memory errors during torch.load.""" - mock_get_rank.return_value = 0 - mock_is_initialized.return_value = False - mock_torch_load.side_effect = RuntimeError("CUDA out of memory") - - with patch("megatron.bridge.training.utils.checkpoint_utils.TrainState", return_value=MockTrainState()): - with pytest.raises(RuntimeError, match="Unable to load train state file"): - read_train_state("large_state.pt") - def test_unicode_and_special_characters_in_config(self, tmp_path): """Test handling of Unicode and special characters in config files.""" config_data = { diff --git a/tests/unit_tests/training/utils/test_train_utils.py b/tests/unit_tests/training/utils/test_train_utils.py index 5b8ec41d63..c757d8b907 100644 --- a/tests/unit_tests/training/utils/test_train_utils.py +++ b/tests/unit_tests/training/utils/test_train_utils.py @@ -239,7 +239,7 @@ def mock_global_state(self): global_state = mock.MagicMock() # Mock train state - global_state.train_state.step = 100 + global_state.train_state.iteration = 100 global_state.train_state.consumed_train_samples = 12800 global_state.train_state.skipped_train_samples = 0 @@ -291,7 +291,7 @@ def test_basic_logging_without_skip( mock_get_world_size.return_value = 32 # Override iteration to avoid log interval reset (101 % 5 != 0) - mock_global_state.train_state.step = 101 + mock_global_state.train_state.iteration = 101 # Call the function result = training_log( @@ -351,7 +351,7 @@ def test_skipped_iterations( mock_get_world_size.return_value = 32 # Override iteration to avoid log interval reset (101 % 5 != 0) - mock_global_state.train_state.step = 101 + mock_global_state.train_state.iteration = 101 # Call the function with skipped iteration result = training_log( @@ -407,7 +407,7 @@ def test_nan_detection( mock_get_world_size.return_value = 32 # Override iteration to avoid log interval reset (101 % 5 != 0) - mock_global_state.train_state.step = 101 + mock_global_state.train_state.iteration = 101 # Create loss dict with NaN values nan_loss_dict = { @@ -468,7 +468,7 @@ def test_tensorboard_logging_interval( mock_get_world_size.return_value = 32 # Set iteration to match tensorboard logging interval - mock_global_state.train_state.step = 100 # Should trigger tensorboard logging (100 % 10 == 0) + mock_global_state.train_state.iteration = 100 # Should trigger tensorboard logging (100 % 10 == 0) mock_config.logger.tensorboard_log_interval = 10 training_log( @@ -515,7 +515,7 @@ def test_comet_only_logs_general_training_scalars(self, monkeypatch, mock_config mock_config.logger.log_loss_scale_to_tensorboard = False mock_config.logger.log_world_size_to_tensorboard = False - mock_global_state.train_state.step = 1 + mock_global_state.train_state.iteration = 1 mock_global_state.tensorboard_logger = None mock_global_state.wandb_logger = None mock_global_state.mlflow_logger = None @@ -589,7 +589,7 @@ def write(self, names, writer, iteration, normalizer=1.0, reset=True, barrier=Fa mock_config.logger.log_loss_scale_to_tensorboard = False mock_config.logger.log_world_size_to_tensorboard = False - mock_global_state.train_state.step = 1 + mock_global_state.train_state.iteration = 1 mock_global_state.timers = timers mock_global_state.tensorboard_logger = None mock_global_state.wandb_logger = None @@ -649,7 +649,7 @@ def test_max_attention_logit_logged( mock_reduce_lr.return_value = 1e-4 mock_get_world_size.return_value = 32 - mock_global_state.train_state.step = 100 + mock_global_state.train_state.iteration = 100 mock_config.logger.tensorboard_log_interval = 10 training_log( @@ -709,7 +709,7 @@ def test_throughput_uses_interval_flops_delta( """Logged TFLOP/s is the interval FLOPs delta ÷ interval time ÷ world size. Regression guard for the THD logging fix: ``training_log`` must derive - throughput from the cumulative ``floating_point_operations_so_far`` delta + throughput from the cumulative ``num_floating_point_operations_so_far`` delta since the last log (over the full interval elapsed time), not from a single step's FLOPs over the per-iteration average. It must also advance the ``_flops_at_last_log`` anchor. @@ -723,12 +723,12 @@ def test_throughput_uses_interval_flops_delta( mock_get_world_size.return_value = 8 # Log boundary (10 % 5 == 0). interval-time mock returns 0.5s (see fixture). - mock_global_state.train_state.step = 10 + mock_global_state.train_state.iteration = 10 mock_config.logger.log_throughput_to_tensorboard = True # Interval FLOPs = so_far - anchor = 8e12; per_gpu_tf = 8e12 / 0.5 / 8 / 1e12 = 2.0 prev_flops = 100.0e12 mock_global_state._flops_at_last_log = prev_flops - mock_global_state.train_state.floating_point_operations_so_far = prev_flops + 8.0e12 + mock_global_state.train_state.num_floating_point_operations_so_far = prev_flops + 8.0e12 expected_per_gpu_tf = 8.0e12 / 0.5 / 8 / 1e12 # == 2.0 training_log( @@ -791,7 +791,7 @@ def test_timing_log_level_1( # Set timing_log_level to 1 mock_config.logger.timing_log_level = 1 - mock_global_state.train_state.step = 100 + mock_global_state.train_state.iteration = 100 mock_config.logger.tensorboard_log_interval = 10 training_log( @@ -859,7 +859,7 @@ def test_timing_log_level_2( # Set timing_log_level to 2 mock_config.logger.timing_log_level = 2 - mock_global_state.train_state.step = 100 + mock_global_state.train_state.iteration = 100 mock_config.logger.tensorboard_log_interval = 10 training_log( @@ -927,7 +927,7 @@ def test_memory_reporting( mock_get_rank.return_value = 0 # Set iteration to match log interval for memory reporting - mock_global_state.train_state.step = 5 + mock_global_state.train_state.iteration = 5 mock_config.logger.log_interval = 5 # Call the function with memory reporting enabled @@ -1012,7 +1012,7 @@ def test_memory_reporting_cutoff( mock_config.model.cuda_graph_impl = cuda_graph_impl mock_config.model.vision_cuda_graph_impl = vision_cuda_graph_impl mock_config.optimizer.optimizer_cuda_graph = optimizer_cuda_graph - mock_global_state.train_state.step = iteration + mock_global_state.train_state.iteration = iteration mock_config.logger.log_interval = 1 result = training_log( @@ -1065,7 +1065,7 @@ def test_memory_reporting_checkpoint_resume( mock_get_rank.return_value = 0 # First iteration after resume from checkpoint at iteration 100 - mock_global_state.train_state.step = 101 + mock_global_state.train_state.iteration = 101 mock_config.logger.log_interval = 1 result = training_log( @@ -1620,7 +1620,7 @@ def test_energy_monitoring( mock_global_state.energy_monitor = mock_energy_monitor # Set iteration to match log interval - mock_global_state.train_state.step = 5 + mock_global_state.train_state.iteration = 5 mock_config.logger.log_interval = 5 training_log( @@ -1651,10 +1651,10 @@ def test_energy_monitoring( # Verify tensorboard logging for energy metrics mock_global_state.tensorboard_logger.add_scalar.assert_any_call( - "iter-energy/gpu", mock.ANY, mock_global_state.train_state.step + "iter-energy/gpu", mock.ANY, mock_global_state.train_state.iteration ) mock_global_state.tensorboard_logger.add_scalar.assert_any_call( - "power/gpu", mock.ANY, mock_global_state.train_state.step + "power/gpu", mock.ANY, mock_global_state.train_state.iteration ) @mock.patch("megatron.bridge.training.utils.train_utils.get_num_microbatches") @@ -1715,7 +1715,7 @@ def test_profiling_memory_snapshot( mock_config.logger.tensorboard_dir = "/tmp/tb" # Set iteration to the configured snapshot dump step. - mock_global_state.train_state.step = 10 + mock_global_state.train_state.iteration = 10 training_log( loss_dict=loss_dict, @@ -1788,7 +1788,7 @@ def test_profiling_memory_snapshot_skips_non_stop_step( mock_profiling_config.profile_step_end = 10 mock_config.profiling = mock_profiling_config mock_config.logger.tensorboard_dir = "/tmp/tb" - mock_global_state.train_state.step = 9 + mock_global_state.train_state.iteration = 9 training_log( loss_dict=loss_dict, @@ -1846,7 +1846,7 @@ def test_wandb_specific_logging( mock_get_world_size.return_value = 32 # Set iteration to match tensorboard logging interval - mock_global_state.train_state.step = 10 + mock_global_state.train_state.iteration = 10 mock_config.logger.tensorboard_log_interval = 10 training_log( @@ -1908,7 +1908,7 @@ def test_no_loggers_present( mock_global_state.comet_logger = None # Set iteration to match logging intervals - mock_global_state.train_state.step = 10 + mock_global_state.train_state.iteration = 10 mock_config.logger.tensorboard_log_interval = 10 mock_config.logger.log_interval = 5 @@ -1981,7 +1981,7 @@ def test_memory_tensorboard_logging( mock_config.logger.log_memory_to_tensorboard = True # Set iteration to match tensorboard logging interval - mock_global_state.train_state.step = 10 + mock_global_state.train_state.iteration = 10 mock_config.logger.tensorboard_log_interval = 10 training_log( @@ -2043,7 +2043,7 @@ def test_report_runtime(self): seq_length = 2048 train_iters = 1000 - train_state = MockTrainState(step=step, consumed_train_samples=consumed_train_samples) + train_state = MockTrainState(iteration=step, consumed_train_samples=consumed_train_samples) runtime_report = report_runtime( train_state=train_state, start_time=start_time, @@ -2598,10 +2598,10 @@ def test_prepare_with_state_parameter_injects(self): """Test prepare_forward_step_func with function that needs state injection.""" def forward_with_state(state: GlobalState, data_iterator, model): - return state.train_state.step + return state.train_state.iteration mock_state = mock.MagicMock() - mock_state.train_state.step = 42 + mock_state.train_state.iteration = 42 result = prepare_forward_step_func(forward_with_state, mock_state) @@ -2633,11 +2633,11 @@ def __init__(self): def __call__(self, state: GlobalState, data_iterator, model): self.call_count += 1 - return state.train_state.step + self.call_count + return state.train_state.iteration + self.call_count functor = ForwardFunctor() mock_state = mock.MagicMock() - mock_state.train_state.step = 10 + mock_state.train_state.iteration = 10 result = prepare_forward_step_func(functor, mock_state) @@ -2677,10 +2677,10 @@ def test_prepare_sees_state_mutations(self): """Test that prepared function sees mutations to GlobalState.""" def forward_with_state(state: GlobalState, data_iterator, model): - return state.train_state.step + return state.train_state.iteration mock_state = mock.MagicMock() - mock_state.train_state.step = 10 + mock_state.train_state.iteration = 10 # Prepare once wrapped = prepare_forward_step_func(forward_with_state, mock_state) @@ -2689,13 +2689,13 @@ def forward_with_state(state: GlobalState, data_iterator, model): assert wrapped(None, None) == 10 # Mutate state (simulates training loop incrementing step) - mock_state.train_state.step = 20 + mock_state.train_state.iteration = 20 # Call again - should see mutated value assert wrapped(None, None) == 20 # Further mutation - mock_state.train_state.step = 100 + mock_state.train_state.iteration = 100 # Still sees current value assert wrapped(None, None) == 100