remove warmup callback

Author: denys-fridman

moe_pretraining/nemo/callback.py

@@ -16,7 +16,6 @@
 import os
 import time
 from functools import wraps
-from pprint import pprint
 from typing import Any, Callable, List, Optional, Protocol, Union
 
 import torch
@@ -40,10 +39,6 @@
 logger = logging.getLogger(__name__)
 
 
-# =============================================================================
-# Utility Functions
-# =============================================================================
-
 def get_last_pp_rank():
     """Check if current rank is the last pipeline parallel rank."""
     is_last_pp = mpu.is_pipeline_last_stage(ignore_virtual=True)
@@ -69,22 +64,9 @@ def broadcast_loss(loss_reduced):
     return loss_synced.item()
 
 
-def get_rank():
-    """Get current process rank for warmup callback."""
-    return int(os.getenv("SLURM_PROCID", 0))
-
-
-# =============================================================================
-# MLPerf Logger
-# =============================================================================
-
 mllogger = MLLoggerWrapper(PyTCommunicationHandler())
 
 
-# =============================================================================
-# Timer Utility
-# =============================================================================
-
 class DeltaTimer:
     """Timer for measuring time deltas."""
 
@@ -378,10 +360,6 @@ def install_callbacks() -> None:
         pretrain_module.train = train_module.train
 
 
-# =============================================================================
-# MLPerf Logging Callback
-# =============================================================================
-
 class MLPerfLoggingCallback:
     """MLPerf logging callback for compliance logging."""
 
@@ -575,10 +553,6 @@ def _get_samples_count(self, global_state):
         return self._get_step(global_state) * self.global_batch_size
 
 
-# =============================================================================
-# Delta Timing Callback
-# =============================================================================
-
 class DeltaTimingCallback:
     """Callback for tracking training step timing."""
 
@@ -633,157 +607,3 @@ def on_validation_end(
     ):
         """Reset timer after validation to avoid including validation time in first train step."""
         self.t0 = time.time()
-
-
-# =============================================================================
-# Warmup Callback
-# =============================================================================
-
-def get_mock_data(config):
-    """Get mock data configuration for warmup."""
-    from megatron.bridge.training.config import MockGPTDatasetConfig
-
-    return MockGPTDatasetConfig(
-        sequence_length=config.model.encoder_seq_length,
-        random_seed=config.model.seed,
-        dataloader_type="single",
-        num_workers=config.model.data.num_workers,
-        reset_position_ids=False,
-        reset_attention_mask=False,
-        eod_mask_loss=False,
-        path_to_cache=None,
-        split="900,50,50",
-    )
-
-
-class WarmupCallback:
-    """Callback for performing training and validation warmup."""
-
-    def __init__(self, cfg):
-        self.cfg = cfg
-        self.train_steps = cfg.model.custom.warmup_train_steps
-        self.val_steps = cfg.model.custom.warmup_validation_steps
-
-    def on_train_start(
-        self,
-        global_state,
-        forward_step_func,
-        model,
-        optimizer,
-        scheduler,
-    ):
-        if get_rank() == 0:
-            print("\nMCore config:", flush=True)
-            pprint(model[0].config)
-
-        torch.distributed.barrier()
-        enable_cuda_graph = int(os.getenv("MCORE_CUDA_GRAPH", "0")) == 1
-        cuda_graph_scope = self.cfg.model.overwritten_attributes.cuda_graph_scope
-
-        forward_backward_func = get_forward_backward_func()
-        if enable_cuda_graph and cuda_graph_scope == "full_iteration":
-            forward_backward_func = FullCudaGraphWrapper(
-                forward_backward_func,
-                cuda_graph_warmup_steps=1,
-            )
-
-        skip_val_warmup = enable_cuda_graph and cuda_graph_scope != "full_iteration"
-
-        warmup_mock_cfg = get_mock_data(self.cfg)
-
-        train_dataloader, val_dataloader, _ = warmup_mock_cfg.build_dataloaders()
-        data_iterator = iter(train_dataloader)
-        eval_data_iterator = iter(val_dataloader)
-
-        pp_group = mpu.get_pipeline_model_parallel_group()
-        torch.distributed.barrier(pp_group)
-
-        for group in optimizer.param_groups:
-            group["betas_"] = group["betas"]
-            group["bias_correction_"] = group["bias_correction"]
-            group["betas"] = [1.0, 1.0]
-            group["bias_correction"] = False
-
-        if torch.distributed.get_rank() == 0:
-            logger.info("Starting training warmup")
-        start = time.time()
-        for step_idx in range(self.train_steps):
-            if torch.distributed.get_rank() == 0:
-                logger.info(f"    Starting warmup step {step_idx}")
-                step_timer = time.time()
-            torch.cuda.synchronize()
-            torch.distributed.barrier()
-            forward_backward_func(
-                forward_step_func=forward_step_func,
-                data_iterator=data_iterator,
-                model=model,
-                num_microbatches=get_num_microbatches(),
-                seq_length=self.cfg.model.encoder_seq_length,
-                micro_batch_size=self.cfg.model.micro_batch_size,
-                decoder_seq_length=self.cfg.model.encoder_seq_length,
-                forward_only=False,
-            )
-            optimizer.zero_grad()
-            optimizer.step()
-            optimizer.zero_grad()
-            torch.cuda.synchronize()
-
-            for chunk in model:
-                chunk.module.zero_grad_buffer()
-                chunk.module.zero_grad()
-
-            if torch.distributed.get_rank() == 0:
-                logger.info(f"    Finished warmup step {step_idx}, takes {time.time() - step_timer} s")
-
-        torch.cuda.synchronize()
-        torch.distributed.barrier()
-        if torch.distributed.get_rank() == 0:
-            logger.info(f"Finished training warmup: {time.time() - start} s. ")
-
-        for group in optimizer.param_groups:
-            group["betas"] = group["betas_"]
-            group["bias_correction"] = group["bias_correction_"]
-            del group["betas_"]
-            del group["bias_correction_"]
-            if "step" in group:
-                if isinstance(group["step"], torch.Tensor):
-                    group["step"].fill_(1)
-                else:
-                    group["step"] = 1
-
-        if not skip_val_warmup:
-            start = time.time()
-            if torch.distributed.get_rank() == 0:
-                logger.info("Starting validation warmups")
-            for model_module in model:
-                model_module.eval()
-            with torch.no_grad():
-                for _ in range(self.val_steps):
-                    torch.cuda.synchronize()
-                    torch.distributed.barrier()
-
-                    forward_backward_func(
-                        forward_step_func=forward_step_func,
-                        data_iterator=eval_data_iterator,
-                        model=model,
-                        num_microbatches=get_num_microbatches(),
-                        seq_length=self.cfg.model.encoder_seq_length,
-                        micro_batch_size=self.cfg.model.micro_batch_size,
-                        forward_only=True,
-                    )
-                    torch.cuda.synchronize()
-
-            torch.distributed.barrier()
-            if torch.distributed.get_rank() == 0:
-                logger.info(f"Finished validation warmup: {time.time() - start} s. ")
-
-        for chunk in model:
-            chunk.module.zero_grad_buffer()
-            chunk.module.zero_grad()
-        if torch.distributed.get_rank() == 0:
-            logger.info(f"Finished training warmup: {time.time() - start} s. ")
-
-        torch.cuda.synchronize()
-        torch.distributed.barrier()
-        if torch.distributed.get_rank() == 0:
-            logger.info(f"Time spent in run_training_warmup: {time.time() - start}s")

moe_pretraining/nemo/pretrain_deepseek_v3_671b.py

@@ -27,7 +27,6 @@
 from callback import (
     MLPerfLoggingCallback,
     DeltaTimingCallback,
-    WarmupCallback,
     mllogger,
     install_callbacks,
     register_callback,