train.py

import contextlib
import gc
import json
import os
import sys
import time
from collections import defaultdict
from dataclasses import dataclass
from datetime import timedelta
from functools import partial
from pathlib import Path
from timeit import default_timer as timer
from typing import Any

import numpy as np
import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch.distributed.checkpoint.stateful import Stateful
from torch.distributed.elastic.multiprocessing.errors import record
from torch.distributed.tensor.parallel import loss_parallel

from transformers import AutoTokenizer, PreTrainedTokenizerFast

from maester.checkpoint import CheckpointManager
from maester.config import Config
from maester.data_monitor import DataMonitor
from maester.datasets.experimental_otf import build_experimental_data_loader
# from maester.datasets.experimental import build_experimental_data_loader # TODO: clean up datasets integration
from maester.log_utils import init_logger, logger
from maester.lr_scheduling import get_lr_scheduler
from maester.memory import cleanup_before_training
from maester.metrics import build_gpu_memory_monitor, build_metric_logger, register_logits_monitoring, WeightScaleMonitor
from maester.data_monitor import DataMonitor
from maester.models import (
    model_name_to_cls,
    models_config,
    model_name_to_parallelize,
    model_name_to_optimizers_builder,
)
from maester.parallelisms import ParallelDims
from maester.profiling import (maybe_enable_memory_snapshot,
                               maybe_enable_profiling)
from maester.sft import build_sft_data_loader
from maester.utils import (clean_param_name, clip_grad_norm, dist_max, dist_mean, get_num_flop_per_token,
                           get_num_params, get_peak_flops, init_distributed,
                           set_pg_timeouts)
from nccl_preflight import run_nccl_preflight

# Training state that is saved in checkpoints
@dataclass
class TrainState(Stateful):
    step: int = 0

    def state_dict(self) -> dict[str, Any]:
        return {
            "step": torch.tensor(self.step, dtype=torch.int32),
        }

    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
        self.step = state_dict["step"].item()


# TODO: do these do much/anything?
# torch._inductor.config.coordinate_descent_tuning = True # type: ignore
torch._inductor.config.triton.unique_kernel_names = True # type: ignore
torch._inductor.config.fx_graph_cache = True # Experimental feature to reduce compilation times, will be on by default in future # type: ignore


# Enable debug tracing on failure: https://pytorch.org/docs/stable/elastic/errors.html
@record
def main():
    init_logger()
    logger.info(f"Starting training.")

    if len(sys.argv) > 1: 
        config_path = Path(sys.argv[1]) / "config.json"
        if not config_path.exists():
            raise ValueError(f"Config not found: {config_path}")
        logger.info(f"Loading config from {config_path}")
        with open(config_path, 'r') as f:
            cfg = Config(**json.load(f))
    else:
        logger.info("Using configuration from config.py")
        cfg = Config()

    # SFT imports if enabled
    if cfg.sft is not None:
        logger.info("SFT mode enabled")
    
    # take control of garbage collection to avoid stragglers
    gc.disable()
    gc.collect(1)

    # init world mesh
    world_size = int(os.environ["WORLD_SIZE"])
    parallel_dims = ParallelDims(
        dp_shard=cfg.data_parallel_shard_degree,
        dp_replicate=cfg.data_parallel_replicate_degree,
        tp=cfg.tensor_parallel_degree,
        ep=cfg.expert_parallel_degree,
        world_size=world_size,
        enable_loss_parallel=cfg.enable_loss_parallel,
    )
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
    init_distributed(cfg)

    if os.environ.get("RUN_NCCL_PREFLIGHT", "0").lower() in {"1", "true", "yes"}:
        logger.info("Running NCCL preflight checks (unset RUN_NCCL_PREFLIGHT or set to 0 to skip)")
        run_nccl_preflight()

    train_state = TrainState()
    with maybe_enable_memory_snapshot(
        cfg, global_step=train_state.step
    ) as memory_profiler:

        # build meshes
        world_mesh = parallel_dims.world_mesh
        if parallel_dims.dp_enabled:
            dp_mesh = world_mesh["dp"]
            dp_degree = dp_mesh.size()
            dp_rank = dp_mesh.get_local_rank()
        else:
            dp_degree, dp_rank = 1, 0
        logger.info(f"world mesh: {world_mesh}")
        if parallel_dims.dp_enabled:
            logger.info(f"{dp_mesh=}")

        # Get tokenizer to determine vocab size
        if os.path.isfile(cfg.tokenizer_name):
            tokenizer = PreTrainedTokenizerFast(tokenizer_file=cfg.tokenizer_name)
        else:
            tokenizer = AutoTokenizer.from_pretrained(cfg.tokenizer_name)

        # build model w/ meta init
        model_cls = model_name_to_cls[cfg.model_name]
        model_config = models_config[cfg.model_name][cfg.flavor]
        # set the model configs from training inputs:
        # 1. norm type to decide which norm layer to use
        # 2. vocab size from tokenizer
        # 3. max_seq_len base on inputs
        model_config.norm_type = cfg.norm_type
        if not hasattr(model_config, 'vocab_size') or model_config.vocab_size <= 0: 
            model_config.vocab_size = len(tokenizer)
        model_config.max_seq_len = cfg.seq_len
        if cfg.enable_mup:
            model_config.enable_mup = True
            model_config.mup_input_alpha = cfg.mup_input_alpha
            model_config.mup_output_alpha = cfg.mup_output_alpha
            model_config.mup_width_mul = cfg.model_width / cfg.base_model_width
            model_config.dim = cfg.model_width
            head_dim = 128
            model_config.n_heads = cfg.model_width // head_dim
            if model_config.n_kv_heads:
                model_config.n_kv_heads = min(model_config.n_kv_heads, model_config.n_heads)

        with torch.device("meta"):
            logger.info(
                f"Building {cfg.model_name} {cfg.flavor} with {model_config}"
            )
            model = model_cls.from_model_args(model_config)

        # log model size
        # model_param_count = get_num_params(model)
        # model_param_count_without_embedding = get_num_params(model, exclude_embedding=True)
        # num_flop_per_token = get_num_flop_per_token(
        #     model_param_count if model.model_args.tied_embeddings else model_param_count_without_embedding, # count lm head matmul only
        #     model_config,
        #     cfg.seq_len,
        # )
        model_param_count, num_flop_per_token = model_config.get_nparams_and_flops(model, cfg.seq_len)
        model_param_count_without_embedding = 0
        logger.info(
            f"Model {cfg.model_name} {cfg.flavor} "
            f"size: {model_param_count:,} total parameters ({model_param_count_without_embedding:,} without embeddings)"
        )

        # initialize GPU memory monitor before applying parallelisms to the model
        gpu_memory_monitor = build_gpu_memory_monitor()
        # obtain the peak flops of bf16 type for MFU calculation
        gpu_peak_flops = get_peak_flops(torch.cuda.get_device_properties(0).name)

        # Choose parallelization function based on model type
        parallelize = model_name_to_parallelize[cfg.model_name]
        parallelize(model, world_mesh, parallel_dims, cfg)
        logger.info(f"Model after parallelization {model=}\n")

        # allocate sharded model on GPU and initialize weights via DTensor
        model.to_empty(device="cuda")
        model.init_weights()
        
        # register hooks after compile?
        # get_logits_metrics, cleanup_monitoring, reinit_storage = register_logits_monitoring(
        #     model, 
        #     train_state,
        #     log_freq=cfg.log_freq,
        #     monitor_attention=False  # TODO: configurable?
        # )

        # reinit_storage() # reinitialize logits storage tensors on the gpu 

        gpu_mem_stats = gpu_memory_monitor.get_peak_stats()
        logger.info(
            f"GPU memory usage for model: "
            f"{gpu_mem_stats.max_reserved_gib:.2f}GiB"
            f"({gpu_mem_stats.max_reserved_pct:.2f}%)"
        )

        if cfg.enable_mup and cfg.mup_log_coord_check:
            activation_hooks = []
            activation_stats = defaultdict(list)
            
            def fw_hook(mod: torch.nn.Module, inp, out, key: str):
                if train_state.step % cfg.log_freq == 0:
                    activation_stats[key].append(out.abs().mean().item())
            for module_name, module in model.named_modules():
                if module_name == 'tok_embeddings':
                    activation_hooks.append(
                        module.register_forward_hook(partial(fw_hook, key='tok_embed'))
                    )
                elif 'attention.' in module_name and module_name.endswith(('.wq', '.wk', '.wv', '.wo')):
                    activation_hooks.append(
                        module.register_forward_hook(partial(fw_hook, key='attn'))
                    )
                elif 'feed_forward.' in module_name and module_name.endswith(('.w1', '.w2', '.w3')):
                    activation_hooks.append(
                        module.register_forward_hook(partial(fw_hook, key='ffn'))
                    )
                elif module_name == 'output':
                    activation_hooks.append(
                        module.register_forward_hook(partial(fw_hook, key='output'))
                    )
                else:
                    logger.info(f"No activation hook registered for {module_name}")
            logger.info(f"Activation hooks registered for {len(activation_hooks)} modules")

        # Create appropriate dataloader based on mode
        if cfg.sft is not None:
            data_loader = build_sft_data_loader(cfg, rank=dp_rank, world_size=dp_degree)
        else:
            data_loader = build_experimental_data_loader(cfg, rank=dp_rank, world_size=dp_degree)

        # data_monitor = DataMonitor(train_state, log_freq=cfg.log_freq)

        # Build optimizers using model-specific builder
        optimizers_builder = model_name_to_optimizers_builder[cfg.model_name]
        optimizers = optimizers_builder(model, cfg, parallel_dims)
        
        scheduler = get_lr_scheduler(optimizers, cfg)

        metric_logger = build_metric_logger(cfg)

        # loss_parallel enables dispatching to efficient loss operators
        loss_parallel_ctx = (
            loss_parallel if parallel_dims.loss_parallel_enabled else contextlib.nullcontext
        )

        def loss_fn(pred, labels):
            return F.cross_entropy(pred.flatten(0, 1).float(), labels.flatten(0, 1))

        def opt_step():
            optimizers.step()
            scheduler.step()
        
        # if cfg.compile:
        #     loss_fn = torch.compile(loss_fn)
        #     opt_step = torch.compile(opt_step)

        # training loop
        cleanup_before_training()
        model.train()
        if hasattr(optimizers, 'train'): # some optimizers need to be put in train mode (e.g. schedule free)
            optimizers.train() # type: ignore (.train obviously exists)

        weight_scale_monitor = WeightScaleMonitor(model, log_freq=cfg.log_freq)

        # checkpointing
        checkpoint = CheckpointManager(
            model=model,
            optimizer=optimizers,
            lr_scheduler=scheduler,
            dataloader=data_loader,
            states={"train_state": train_state},
            cfg=cfg,
        )
        checkpoint.load()

        # TODO: do we want to checkpoint metrics?

        data_iterator = iter(data_loader)

        logger.info(f"Training starts at step {train_state.step}")
        with maybe_enable_profiling(
            cfg, global_step=train_state.step
        ) as torch_profiler:
            checkpoint.reset()

            # variables for metric logging
            losses_since_last_log: list[torch.Tensor] = []
            padding_lengths_since_last_log: list[torch.Tensor] = []
            ntokens_since_last_log = 0
            total_tokens = 0
            data_loading_times: list[float] = []
            time_last_log = timer()
            gpu_memory_monitor.reset_peak_stats()

            grad_accum_steps = max(1, cfg.gradient_accumulation_steps)
            fsdp_can_toggle_sync = hasattr(model, "set_requires_gradient_sync")
            skip_sync_during_accum = (
                grad_accum_steps > 1 and not cfg.gradient_accumulation_sync_each_step
            )
            global_micro_step = train_state.step * grad_accum_steps

            while train_state.step < cfg.train_num_steps:
                optimizers.zero_grad(set_to_none=True)

                for micro_idx in range(grad_accum_steps):
                    global_micro_step += 1
                    torch.manual_seed(global_micro_step + dp_rank)

                    data_load_start = timer()
                    batch = next(data_iterator)

                    input_ids = batch["input_ids"]
                    labels = batch["labels"]

                    # Get position_ids if available (currently only from packed SFT data)
                    # TODO: Consider generating position_ids for all data loaders for consistency
                    position_ids = batch.get("position_ids", None)

                    # Get document_ids if available (for flex attention document masking in packed data)
                    document_ids = batch.get("document_ids", None)

                    # Collect padding stats if available (SFT mode)
                    if "stats" in batch and "actual_lengths" in batch["stats"]:
                        padding_lengths_since_last_log.append(batch["stats"]["actual_lengths"])

                    ntokens_since_last_log += labels.numel()
                    data_loading_times.append(timer() - data_load_start)

                    input_ids = input_ids.cuda()
                    labels = labels.cuda()
                    if position_ids is not None:
                        position_ids = position_ids.cuda()
                    if document_ids is not None:
                        document_ids = document_ids.cuda()

                    sync_grads_now = True
                    if skip_sync_during_accum:
                        sync_grads_now = micro_idx == grad_accum_steps - 1

                    if fsdp_can_toggle_sync and grad_accum_steps > 1:
                        model.set_requires_gradient_sync(sync_grads_now)

                    with loss_parallel_ctx():
                        if cfg.enable_cut_cross_entropy:
                            loss = model(
                                input_ids,
                                labels,
                                position_ids=position_ids,
                                document_ids=document_ids,
                            )
                        else:
                            pred = model(
                                input_ids,
                                position_ids=position_ids,
                                document_ids=document_ids,
                            )
                            loss = loss_fn(pred, labels)
                            del pred

                        losses_since_last_log.append(loss.detach())
                        scaled_loss = loss / grad_accum_steps
                        scaled_loss.backward()

                    if (
                        fsdp_can_toggle_sync
                        and grad_accum_steps > 1
                        and skip_sync_during_accum
                        and not sync_grads_now
                    ):
                        model.set_requires_gradient_sync(True)

                # TODO: re-enable grad clipping (broken w/ MoE) and/or monitoring?
                # grad_norms = clip_grad_norm( # note: maester.utils.clip_grad_norm, not torch.nn.utils.clip_grad_norm_
                #     model.parameters(), cfg.max_grad_norm, foreach=True
                # )
                #optimizers.step()
                #scheduler.step()
                opt_step()
                train_state.step += 1

                if train_state.step % cfg.gc_freq == 0:
                    gc.collect(1)

                weight_scale_stats = weight_scale_monitor.step_monitor()

                # log metrics
                if train_state.step == 1 or train_state.step % cfg.log_freq == 0:
                    losses = [l.detach().item() for l in losses_since_last_log]
                    avg_loss, max_loss = (
                        np.mean(losses),
                        np.max(losses),
                    )
                    if parallel_dims.dp_enabled:
                        global_avg_loss, global_max_loss = (
                            dist_mean(avg_loss, dp_mesh).item(), # type: ignore (dp_mesh exists)
                            dist_max(max_loss, dp_mesh).item() # type: ignore (dp_mesh exists)
                        )
                    else:
                        global_avg_loss, global_max_loss = avg_loss, max_loss

                    # TODO: re-enable grad norm logging?
                    # param_to_name = {param: name for name, param in model.named_parameters()}
                    # exp_avgs, exp_avg_sqs, param_names = [], [], []
                    # for group in optimizers.param_groups:
                    #     for p in group['params']:
                    #         if p.grad is None:
                    #             continue
                    #         state = optimizers.state[p]
                    #         if 'exp_avg' in state:  # Check if states initialized
                    #             exp_avgs.append(state['exp_avg'])
                    #             exp_avg_sqs.append(state['exp_avg_sq'])
                    #             param_names.append(param_to_name[p])
                    # exp_avg_norms = torch._foreach_norm(exp_avgs, 2)
                    # exp_avg_sq_norms = torch._foreach_norm(exp_avg_sqs, 2)

                    time_delta = timer() - time_last_log

                    total_tokens += ntokens_since_last_log
                    tps = ntokens_since_last_log / (time_delta * parallel_dims.model_parallel_size)
                    mfu = 100 * num_flop_per_token * tps / gpu_peak_flops

                    time_end_to_end = time_delta / cfg.log_freq
                    time_data_loading = np.mean(data_loading_times)
                    time_data_loading_pct = 100 * np.sum(data_loading_times) / time_delta
                    
                    # Aggregate data loading times across ALL ranks (TP ranks load redundantly)
                    # Flatten world mesh to get all ranks
                    global_mesh = world_mesh._flatten() if hasattr(world_mesh, '_flatten') else world_mesh
                    global_avg_data_loading = dist_mean(time_data_loading, global_mesh).item()
                    global_max_data_loading = dist_max(time_data_loading, global_mesh).item()
                    global_avg_data_loading_pct = dist_mean(time_data_loading_pct, global_mesh).item()
                    global_max_data_loading_pct = dist_max(time_data_loading_pct, global_mesh).item()

                    gpu_mem_stats = gpu_memory_monitor.get_peak_stats()

                    # TODO: add data metrics?
                    metrics = {
                        # "epoch": epoch,
                        "loss/global_avg": global_avg_loss,
                        "loss/global_max": global_max_loss,
                        "tps": tps,
                        "mfu(%)": mfu,
                        "data/total_tokens": total_tokens * parallel_dims.dp_shard * parallel_dims.dp_replicate,
                        "time/end_to_end(s)": time_end_to_end,
                        "time/data_loading_avg(s)": global_avg_data_loading,
                        "time/data_loading_max(s)": global_max_data_loading,
                        "time/data_loading_avg(%)": global_avg_data_loading_pct,
                        "time/data_loading_max(%)": global_max_data_loading_pct,
                        "memory/max_active(GiB)": gpu_mem_stats.max_active_gib,
                        "memory/max_active(%)": gpu_mem_stats.max_active_pct,
                        "memory/max_reserved(GiB)": gpu_mem_stats.max_reserved_gib,
                        "memory/max_reserved(%)": gpu_mem_stats.max_reserved_pct,
                        "memory/num_alloc_retries": gpu_mem_stats.num_alloc_retries,
                        "memory/num_ooms": gpu_mem_stats.num_ooms,
                    }
                    
                    # Add padding stats if available (SFT mode)
                    if padding_lengths_since_last_log:
                        all_lengths = torch.cat(padding_lengths_since_last_log)
                        seq_len = input_ids.shape[1]  # Max sequence length
                        
                        # Calculate efficiency: what % of tokens are actual content (not padding)
                        total_actual_tokens = all_lengths.sum().item()
                        total_batch_tokens = all_lengths.numel() * seq_len
                        efficiency = total_actual_tokens / total_batch_tokens
                        
                        metrics.update({
                            "padding/efficiency": efficiency,  # % of tokens that are actual content
                            "padding/avg_length": all_lengths.float().mean().item(),
                            "padding/std_length": all_lengths.float().std().item(),
                        })
                    # for i in range(len(optimizer.param_groups)):
                    #     metrics[f"lr/group{i}"] = scheduler.get_last_lr()[i]
                    # for gn, (name, _) in zip(grad_norms, model.named_parameters()):
                    #     cn = clean_param_name(name)
                    #     metrics[f"{cn}/grad_norm"] = gn
                    # for exp_avg_norm, exp_avg_sq_norm, name in zip(exp_avg_norms, exp_avg_sq_norms, param_names):
                    #     cn = clean_param_name(name)
                    #     metrics[f"{cn}/exp_avg_norm"] = exp_avg_norm
                    #     metrics[f"{cn}/exp_avg_sq_norm"] = exp_avg_sq_norm
                    if cfg.enable_mup and cfg.mup_log_coord_check:
                        for key in activation_stats: # type: ignore
                            if activation_stats[key]: # type: ignore
                                metrics[f'act/{key}_abs_mean'] = np.mean(activation_stats[key]) # type: ignore
                        activation_stats = defaultdict(list) # reset
                    # metrics.update(get_logits_metrics())
                    if weight_scale_stats is not None:
                        metrics.update(weight_scale_stats)
                    
                    # Collect optimizer hook statistics if available
                    # TODO: This is a temporary solution - consider refactoring for cleaner API
                    if hasattr(optimizers, '_hook_stats') and optimizers._hook_stats:
                        metrics.update(optimizers._hook_stats)
                    
                    if metric_logger is not None:
                        metric_logger.log(metrics, step=train_state.step)

                    logger.info(
                        f"Step {train_state.step:2}: "
                        f"lr={scheduler.get_last_lr()[0]:.2E}, "
                        f"loss={global_avg_loss:7.4f} (max={global_max_loss:7.4f}), "
                        f"tps={round(tps):}, "
                        f"mfu={mfu:.2f}%, "
                        f"memory: {gpu_mem_stats.max_reserved_gib:5.2f}GiB"
                        f"({gpu_mem_stats.max_reserved_pct:.2f}%) "
                        f"time/data_loading={global_avg_data_loading:.2f}s (max={global_max_data_loading:.2f}s, {global_max_data_loading_pct:.2f}%)"
                    )

                    losses_since_last_log.clear()
                    padding_lengths_since_last_log.clear()
                    ntokens_since_last_log = 0
                    data_loading_times.clear()
                    time_last_log = timer()
                    gpu_memory_monitor.reset_peak_stats()

                checkpoint.save(
                    train_state.step, force=(train_state.step == cfg.train_num_steps)
                )

                # signals the profiler that the next profiling step has started
                if torch_profiler:
                    torch_profiler.step()
                if memory_profiler:
                    memory_profiler.step()

                # TODO: Reduce timeout after first train step for faster signal (assumes lazy init, compile are finished)
                if train_state.step == 1:
                    set_pg_timeouts(
                        timeout=timedelta(seconds=cfg.train_timeout_seconds),
                        world_mesh=world_mesh,
                    )

        if dist.get_rank() == 0:
            logger.info("Sleeping 2 seconds for other ranks to complete")
            time.sleep(2)
        if metric_logger is not None:
            metric_logger.close()
        logger.info("Training successfully completed!")
        dist.destroy_process_group()

if __name__ == '__main__':
    main()