runtime/litgpt/pretrain.py

# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.

import math
import os
import pprint
import time
import warnings
from dataclasses import asdict
from datetime import timedelta
from functools import partial
from pathlib import Path
from typing import Literal

import lightning as L
import torch
import torch.nn as nn
from lightning.fabric.strategies import FSDPStrategy
from lightning.fabric.utilities.throughput import ThroughputMonitor, measure_flops
from torch.utils.data import DataLoader
from torchmetrics.aggregation import RunningMean

from litgpt import Tokenizer
from litgpt.args import EvalArgs, LogArgs, TrainArgs
from litgpt.config import name_to_config
from litgpt.constants import _TORCH_EQUAL_2_7, _TORCH_EQUAL_2_8
from litgpt.data import DataModule, TinyLlama
from litgpt.model import (
    GPT,
    AdaptiveDSwiGLUMLP,
    Block,
    BlockSparseAdaptiveDSwiGLUMLP,
    CausalSelfAttention,
    Config,
    DSwiGLUMLP,
    LLaMAMLP,
    TileRoutedChannelGroupStateMemoryDSwiGLUMLP,
    TileRoutedChannelMemoryDSwiGLUMLP,
    TileRoutedChannelStateMemoryDSwiGLUMLP,
    TileRoutedCenteredChannelGroupStateMemoryDSwiGLUMLP,
    TileRoutedCenteredChannelStateMemoryDSwiGLUMLP,
    TileRoutedDSwiGLUMLP,
    TileRoutedGroupStateMemoryDSwiGLUMLP,
    TileRoutedKeyedChannelMemoryDSwiGLUMLP,
    TileRoutedDSwiGLUMLPStaticA2,
    TileRoutedDSwiGLUMLPStaticGPTS,
)
from litgpt.optim import AdaptiveRowTAdamW, HybridMuonAdamW, OfficialMuonWithAuxAdam, RMNPGrouped
from litgpt.parser_config import save_hyperparameters
from litgpt.types import LoggerChoice
from litgpt.utils import (
    CycleIterator,
    capture_hparams,
    check_nvlink_connectivity,
    choose_logger,
    chunked_cross_entropy,
    copy_config_files,
    extend_checkpoint_dir,
    find_resume_path,
    get_default_supported_precision,
    init_out_dir,
    instantiate_torch_optimizer,
    num_parameters,
    parse_devices,
    reset_parameters,
    save_config,
)

torch.set_float32_matmul_precision("high")
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(True)


def setup(
    model_name: str,
    model_config: Config | None = None,
    out_dir: Path = Path("out/pretrain"),
    precision: Literal["bf16-true", "bf16-mixed", "32-true", None] = None,
    initial_checkpoint_dir: Path | None = None,
    resume: bool | Literal["auto"] | Path = False,
    data: DataModule | None = None,
    train: TrainArgs = TrainArgs(
        save_interval=1000,
        log_interval=1,
        global_batch_size=512,
        micro_batch_size=4,
        max_tokens=int(3e12),  # 3 trillion
        max_norm=1.0,
        min_lr=4e-5,
        lr_warmup_steps=2000,
        tie_embeddings=False,
    ),
    eval: EvalArgs = EvalArgs(interval=1000, max_iters=100),
    log: LogArgs = LogArgs(),
    optimizer: str | dict = "AdamW",
    devices: int | str = "auto",
    num_nodes: int = 1,
    tokenizer_dir: Path | None = None,
    logger_name: LoggerChoice = "tensorboard",
    seed: int = 42,
):
    """Pretrain a model.

    Arguments:
        model_name: The name of the model to pretrain. Choose from names in ``litgpt.config``. Use "list" to list the supported models.
        model_config: A ``litgpt.Config`` object to define the model architecture. Mutually exclusive with
            ``model_config``. Overrides the `model_name` if specified.
        out_dir: Directory in which to save checkpoints and logs. If running in a Lightning Studio Job, look for it in
            /teamspace/jobs/<job-name>/share.
        precision: The precision to use for finetuning. Determines a compatible precision setting by default.
        initial_checkpoint_dir: Optional path to a checkpoint directory to initialize the model from.
            Useful for continued pretraining. Mutually exclusive with ``resume``.
        resume: Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
            from the latest checkpoint in ``out_dir``. An error will be raised if no checkpoint is found. Passing
            ``'auto'`` will resume from the latest checkpoint but not error if no checkpoint exists.
        data: Data-related arguments. If not provided, the default is ``litgpt.data.TinyLlama``.
        train: Training-related arguments. See ``litgpt.args.TrainArgs`` for details.
        eval: Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details.
        optimizer: An optimizer name (such as "AdamW") or config.

        devices: How many devices/GPUs to use. Uses all GPUs by default.
        num_nodes: How many nodes the code is being run on.
        tokenizer_dir: Optional path to the tokenizer dir that was used for preprocessing the dataset. Only some data
            module require this.
        logger_name: The name of the logger to send metrics to.
        seed: The random seed to use for reproducibility.
    """
    if model_name == "list":
        available_models = "\n".join(sorted(name_to_config))
        print(f"Available values:\n{available_models}")
        quit()

    if initial_checkpoint_dir is not None:
        initial_checkpoint_dir = extend_checkpoint_dir(initial_checkpoint_dir)

    if tokenizer_dir is not None:
        tokenizer_dir = extend_checkpoint_dir(tokenizer_dir)

    if model_config is None:
        # Support both model_name options: meta-llama/Meta-Llama-3-8B & Meta-Llama-3-8B
        try:
            model_config = Config.from_name(model_name)
        except ValueError:
            print(f"Model name {model_name} is not supported.\n")
            available_models = "\n".join(sorted(name_to_config))
            print(f"Available values:\n{available_models}")
            quit()

    hparams = capture_hparams()
    data = TinyLlama() if data is None else data

    config = Config.from_name(model_name) if model_config is None else model_config
    precision = precision or get_default_supported_precision(training=True)
    devices = parse_devices(devices)
    out_dir = init_out_dir(out_dir)
    # in case the dataset requires the Tokenizer
    tokenizer = Tokenizer(tokenizer_dir) if tokenizer_dir is not None else None

    logger = choose_logger(
        logger_name,
        out_dir,
        name=f"pretrain-{config.name}",
        resume=bool(resume),
        log_interval=train.log_interval,
        log_args=asdict(log),
    )

    if devices * num_nodes > 1:
        strategy = FSDPStrategy(auto_wrap_policy={Block}, state_dict_type="full", sharding_strategy="HYBRID_SHARD")
    else:
        strategy = "auto"

    fabric = L.Fabric(devices=devices, num_nodes=num_nodes, strategy=strategy, precision=precision, loggers=[logger])

    if torch.cuda.is_available() and devices > 1:
        check_nvlink_connectivity(fabric)

    fabric.launch()

    fabric.print(pprint.pformat(hparams))
    if logger_name in ("tensorboard", "wandb", "mlflow"):
        fabric.logger.log_hyperparams(hparams)

    main(
        fabric=fabric,
        devices=devices,
        num_nodes=num_nodes,
        seed=seed,
        initial_checkpoint_dir=initial_checkpoint_dir,
        resume=resume,
        config=config,
        data=data,
        out_dir=out_dir,
        tokenizer_dir=tokenizer_dir,
        tokenizer=tokenizer,
        train=train,
        eval=eval,
        optimizer=optimizer,
    )


def main(
    fabric: L.Fabric,
    devices: int,
    seed: int,
    initial_checkpoint_dir: Path | None,
    resume: bool | Literal["auto"] | Path,
    config: Config,
    data: DataModule,
    out_dir: Path,
    tokenizer_dir: Path | None,
    tokenizer: Tokenizer | None,
    train: TrainArgs,
    eval: EvalArgs,
    optimizer: str | dict,
    num_nodes: int = 1,
) -> None:
    validate_args(train, eval, initial_checkpoint_dir, resume)

    if fabric.global_rank == 0:
        out_dir.mkdir(parents=True, exist_ok=True)

    fabric.seed_everything(seed)  # same seed for every process to init model (FSDP)

    t0 = time.perf_counter()
    with fabric.init_module(empty_init=True):
        model = GPT(config)

    initialize_weights(fabric, model, n_layer=config.n_layer, n_embd=config.n_embd)

    if train.tie_embeddings and not config.mamba3_hierarchical_vocab:
        model.transformer.wte.weight = model.lm_head.weight
    if train.max_seq_length:
        model.max_seq_length = train.max_seq_length
    if config.sliding_window_size is not None:
        model.prepare_sliding_window_masks(model.max_seq_length, fabric.device)

    fabric.print(f"Time to instantiate model: {time.perf_counter() - t0:.02f} seconds.")
    fabric.print(f"Total parameters: {num_parameters(model):,}")
    if bool(train.mtp_loss_weight) != bool(config.mtp_num_layers):
        raise ValueError(
            "Paper-faithful MTP requires both model.mtp_num_layers=1 and "
            "train.mtp_loss_weight>0, or both disabled."
        )

    if train.compile_model:
        if config.kda_enabled:
            # FLA stores a distinct integer ``layer_idx`` on each KDA module.
            # Dynamo otherwise specializes the same forward on every layer,
            # hits its recompile limit after eight variants, and leaves later
            # layers on a slower fallback. The attribute is only a cache key;
            # allowing it to remain symbolic preserves the exact computation.
            torch._dynamo.config.allow_unspec_int_on_nn_module = True
        compile_kwargs = {}
        if train.compile_mode != "default":
            compile_kwargs["mode"] = train.compile_mode
        fabric.print(f"Compiling model with torch.compile({compile_kwargs or 'default'})")
        model = torch.compile(model, **compile_kwargs)
    model = fabric.setup(model)
    if train.mtp_loss_weight:
        model.mark_forward_method('mtp_forward')

    extra_kwargs = {"fused": fabric.device.type == "cuda"}
    if isinstance(optimizer, dict) and "fused" in optimizer.get("init_args", {}):
        # Respect an explicit optimizer setting. This is required for Gefen's
        # documented pure-PyTorch fallback when a matching nvcc is unavailable.
        extra_kwargs = {}
    optimizer_class_path = optimizer.get("class_path", "") if isinstance(optimizer, dict) else ""
    if optimizer_class_path == "litgpt.optim.OfficialMuonWithAuxAdam":
        if train.channel_memory_lr_mult != 1.0:
            raise ValueError("`train.channel_memory_lr_mult` is not supported with OfficialMuonWithAuxAdam.")
        optimizer = instantiate_official_muon_optimizer(optimizer, model)
    elif optimizer_class_path == "litgpt.optim.AdaptiveRowTAdamW":
        if train.channel_memory_lr_mult != 1.0:
            raise ValueError("`train.channel_memory_lr_mult` is not supported with AdaptiveRowTAdamW.")
        optimizer = instantiate_adaptive_row_tadamw(optimizer, model)
    elif optimizer_class_path == "litgpt.optim.RMNPGrouped":
        if train.channel_memory_lr_mult != 1.0:
            raise ValueError("`train.channel_memory_lr_mult` is not supported with RMNPGrouped.")
        optimizer = instantiate_rmnp_optimizer(optimizer, model)
    elif config.mlp_matrix_optimizer == "himuon_tile32":
        if train.channel_memory_lr_mult != 1.0:
            raise ValueError("`train.channel_memory_lr_mult` is only supported for torch optimizers.")
        optimizer = instantiate_hybrid_mlp_optimizer(optimizer, model, **extra_kwargs)
    else:
        if optimizer_class_path == "gefen.Gefen":
            import gefen.gefen as gefen_impl

            backend = os.environ.get("GEFEN_FIND_PERIOD_BACKEND")
            if backend:
                if backend not in {"cpu", "gpu", "cuda_kernel"}:
                    raise ValueError(f"Unsupported GEFEN_FIND_PERIOD_BACKEND={backend!r}")
                gefen_impl.FIND_PERIOD_BACKEND = backend
                fabric.print(f"Gefen period backend: {backend}")
            if os.environ.get("GEFEN_FUSE_HISTOGRAM_FOR_EXACT") == "0":
                gefen_impl.FUSE_HISTOGRAM_FOR_EXACT = False
                fabric.print("Gefen exact histogram fusion: disabled")
        # Gefen accepts named parameters and uses their names in its learned
        # partition diagnostics. Preserve that information for grouped tensors.
        model_parameters = model.named_parameters() if optimizer_class_path == "gefen.Gefen" else model.parameters()
        if config.rwkv7_enabled:
            if optimizer_class_path == "gefen.Gefen":
                raise ValueError("RWKV-7 parameter rules are not implemented for Gefen.")
            if not isinstance(optimizer, dict):
                raise ValueError("RWKV-7 optimizer groups require optimizer.init_args.")
            init_args = optimizer.get("init_args", {})
            base_lr = init_args.get("lr")
            weight_decay = init_args.get("weight_decay", 0.0)
            if base_lr is None:
                raise ValueError("RWKV-7 optimizer groups require optimizer.init_args.lr.")
            decay_parameters = []
            no_decay_parameters = []
            decay_lr2_parameters = []
            decay_suffixes = (
                ".receptance.weight",
                ".key.weight",
                ".value.weight",
                ".output.weight",
                ".local_screen.query_weight",
                ".local_screen.key_weight",
                ".local_screen.value_weight",
                ".local_screen.gate_weight",
                ".local_screen.output_weight",
            )
            for name, parameter in model.named_parameters():
                if not parameter.requires_grad:
                    continue
                if name.endswith(".w0"):
                    decay_lr2_parameters.append(parameter)
                elif (
                    name in {"transformer.wte.weight", "lm_head.weight"}
                    or name.endswith(decay_suffixes)
                ):
                    decay_parameters.append(parameter)
                else:
                    no_decay_parameters.append(parameter)
            model_parameters = [
                {"params": no_decay_parameters, "weight_decay": 0.0},
                {"params": decay_parameters, "weight_decay": float(weight_decay)},
                {
                    "params": decay_lr2_parameters,
                    "weight_decay": 0.0,
                    "lr": 2.0 * float(base_lr),
                },
            ]
            fabric.print(
                "RWKV-7 optimizer groups: "
                f"{len(decay_parameters)} large tensors at weight_decay={float(weight_decay):g}; "
                f"{len(no_decay_parameters)} tensors without decay; "
                f"{len(decay_lr2_parameters)} time-decay tensors at 2x LR"
            )
        elif config.mamba3_enabled:
            if optimizer_class_path == "gefen.Gefen":
                raise ValueError("Mamba-3 parameter rules are not implemented for Gefen.")
            if not isinstance(optimizer, dict):
                raise ValueError("Mamba-3 optimizer groups require optimizer.init_args.")
            init_args = optimizer.get("init_args", {})
            weight_decay = float(init_args.get("weight_decay", 0.0))
            decay_parameters = []
            no_decay_parameters = []
            for name, parameter in model.named_parameters():
                if not parameter.requires_grad:
                    continue
                no_weight_decay = bool(getattr(parameter, "_no_weight_decay", False))
                is_norm_or_bias = (
                    parameter.ndim < 2
                    or ".norm." in name
                    or name.endswith(".bias")
                    or "basis_coefficients" in name
                )
                (no_decay_parameters if no_weight_decay or is_norm_or_bias else decay_parameters).append(parameter)
            model_parameters = [
                {"params": no_decay_parameters, "weight_decay": 0.0},
                {"params": decay_parameters, "weight_decay": weight_decay},
            ]
            fabric.print(
                "Mamba-3 optimizer groups: "
                f"{len(decay_parameters)} matrices at weight_decay={weight_decay:g}; "
                f"{len(no_decay_parameters)} state, norm, and bias tensors without decay"
            )
        use_memory_group = train.channel_memory_lr_mult != 1.0
        use_grouped_mlp_group = (
            train.grouped_mlp_lr_mult != 1.0
            or train.grouped_mlp_weight_decay_mult != 1.0
        )
        controller_parameter_names = (
            "cvr_energy_coeff",
            "cvr_agreement_coeff",
            "caet_temperature_raw",
            "adaptive_commit_router",
        )
        use_controller_group = config.sandwich_adaptive_commit > 0 or config.mlp_class_name in (
            "TileRoutedContextualValueRotationEulerGLUMLP",
            "TileRoutedContextAdaptiveEulerTemperatureGLUMLP",
        )
        if (
            use_memory_group or use_grouped_mlp_group or use_controller_group
        ) and not (config.rwkv7_enabled or config.mamba3_enabled):
            if not isinstance(optimizer, dict):
                raise ValueError("Optimizer parameter groups require optimizer.init_args.")
            init_args = optimizer.get("init_args", {})
            base_lr = init_args.get("lr")
            base_weight_decay = init_args.get("weight_decay")
            if base_lr is None:
                raise ValueError("Optimizer parameter groups require optimizer.init_args.lr.")
            if use_grouped_mlp_group and base_weight_decay is None:
                raise ValueError("Grouped MLP decay requires optimizer.init_args.weight_decay.")

            base_parameters = []
            memory_parameters = []
            grouped_mlp_parameters = []
            controller_parameters = []
            memory_parameter_names = (
                "channel_memory_gain_raw",
                "group_state_memory_gain_raw",
                "keyed_memory_key",
                "keyed_memory_value_raw",
                "grouped_mlp_mix_alpha",
                "grouped_mlp_mixer",
                "shared_alpha",
                "shared_path",
            )
            grouped_mlp_parameter_names = (
                ".mlp.up_weight",
                ".mlp.value_weight",
                ".mlp.down_weight",
            )
            for name, parameter in model.named_parameters():
                if not parameter.requires_grad:
                    continue
                if use_controller_group and name.endswith(controller_parameter_names):
                    controller_parameters.append(parameter)
                elif use_grouped_mlp_group and name.endswith(grouped_mlp_parameter_names):
                    grouped_mlp_parameters.append(parameter)
                elif use_memory_group and any(
                    memory_name in name for memory_name in memory_parameter_names
                ):
                    memory_parameters.append(parameter)
                else:
                    base_parameters.append(parameter)
            if use_memory_group and not memory_parameters:
                raise ValueError(
                    "`train.channel_memory_lr_mult` was set, but no trainable channel memory gains were found."
                )
            if use_grouped_mlp_group and not grouped_mlp_parameters:
                raise ValueError(
                    "Grouped MLP optimizer settings were set, but no grouped MLP projection tensors were found."
                )
            if use_controller_group and not controller_parameters:
                raise ValueError("Adaptive model was selected, but no controller parameters were found.")

            model_parameters = [{"params": base_parameters}]
            if use_controller_group:
                fabric.print(
                    "Adaptive controller group: "
                    f"{len(controller_parameters)} tensors at lr={float(base_lr):g}, weight_decay=0"
                )
                model_parameters.append(
                    {"params": controller_parameters, "lr": float(base_lr), "weight_decay": 0.0}
                )
            if use_memory_group:
                memory_lr = float(base_lr) * float(train.channel_memory_lr_mult)
                fabric.print(
                    "Channel memory LR multiplier: "
                    f"{train.channel_memory_lr_mult:g}x ({len(memory_parameters)} tensors at lr={memory_lr:g})"
                )
                model_parameters.append({"params": memory_parameters, "lr": memory_lr})
            if use_grouped_mlp_group:
                grouped_mlp_lr = float(base_lr) * float(train.grouped_mlp_lr_mult)
                grouped_mlp_weight_decay = float(base_weight_decay) * float(train.grouped_mlp_weight_decay_mult)
                fabric.print(
                    "Grouped MLP AdamW group: "
                    f"{len(grouped_mlp_parameters)} tensors at lr={grouped_mlp_lr:g}, "
                    f"weight_decay={grouped_mlp_weight_decay:g}"
                )
                model_parameters.append(
                    {
                        "params": grouped_mlp_parameters,
                        "lr": grouped_mlp_lr,
                        "weight_decay": grouped_mlp_weight_decay,
                    }
                )
        optimizer = instantiate_torch_optimizer(optimizer, model_parameters, **extra_kwargs)
    optimizer = fabric.setup_optimizers(optimizer)

    train_dataloader, val_dataloader = get_dataloaders(fabric, data, tokenizer, train, model.max_seq_length)
    train_dataloader, val_dataloader = fabric.setup_dataloaders(train_dataloader, val_dataloader)

    if initial_checkpoint_dir:
        fabric.load_raw(initial_checkpoint_dir / "lit_model.pth", model)

    state = {
        "model": model,
        "optimizer": optimizer,
        "train_dataloader": train_dataloader,
        "iter_num": 0,
        "step_count": 0,
    }

    resume = find_resume_path(resume, out_dir)
    if resume:
        fabric.print(f"Resuming training from {resume}")
        fabric.load(resume, state)

    train_time = time.perf_counter()

    # work around PyTorch issue https://github.com/pytorch/pytorch/issues/152162
    # which does not like the lazy initialization to be called in dynamo.
    # TODO: Happens with PyTorch 2.7+
    if (
        (_TORCH_EQUAL_2_7 or _TORCH_EQUAL_2_8)
        and (model._forward_module.__class__.__name__ == "OptimizedModule")
        and (model._forward_module._orig_mod.__class__.__name__ == "FullyShardedDataParallel")
    ):
        from torch.distributed.fsdp._runtime_utils import _root_pre_forward

        _root_pre_forward(model._forward_module._orig_mod, model._forward_module._orig_mod, [], {})

    fit(
        fabric=fabric,
        devices=devices,
        num_nodes=num_nodes,
        state=state,
        train_dataloader=train_dataloader,
        val_dataloader=val_dataloader,
        out_dir=out_dir,
        tokenizer_dir=tokenizer_dir,
        train=train,
        eval=eval,
    )

    # Save final checkpoint
    save_checkpoint(fabric, state, tokenizer_dir, out_dir / "final" / "lit_model.pth")

    total_tokens = state["iter_num"] * train.micro_batch_size * model.max_seq_length * fabric.world_size

    # Print formatted output
    separator = "-" * 40
    fabric.print(separator)
    fabric.print("| Performance")
    fabric.print(f"| - Total tokens  : {total_tokens:,}")
    training_elapsed = time.perf_counter() - train_time
    fabric.print(f"| - Training Time : {training_elapsed:.2f} s")
    fabric.print(f"| - Tok/sec       : {total_tokens / training_elapsed:.2f} tok/s")
    fabric.print("| " + "-" * 40)

    if fabric.device.type == "cuda":
        memory_used = torch.cuda.max_memory_allocated() / 1e9
        fabric.print("| Memory Usage")
        fabric.print(f"| - Memory Used   : {memory_used:.2f} GB")
    fabric.print(separator)




def instantiate_hybrid_mlp_optimizer(optimizer_config, model: nn.Module, **kwargs) -> HybridMuonAdamW:
    if isinstance(optimizer_config, str):
        init_args = {"lr": 1e-3}
    else:
        init_args = dict(optimizer_config.get("init_args", {}))
    init_args.update({key: value for key, value in kwargs.items() if key in {"lr", "weight_decay", "betas", "eps"}})
    init_args.pop("fused", None)
    muon_params = []
    adamw_params = []
    for name, param in model.named_parameters():
        is_mlp_matrix = ".mlp." in name and param.ndim >= 2
        if is_mlp_matrix:
            muon_params.append(param)
        else:
            adamw_params.append(param)
    param_groups = []
    if adamw_params:
        param_groups.append({"params": adamw_params, "use_muon": False})
    if muon_params:
        param_groups.append({"params": muon_params, "use_muon": True, "tile_size": 32})
    return HybridMuonAdamW(param_groups, **init_args)


def instantiate_rmnp_optimizer(optimizer_config, model: nn.Module) -> RMNPGrouped:
    """Map LitGPT parameter names onto the authors' RMNP/Adam branches."""
    if not isinstance(optimizer_config, dict):
        raise ValueError("RMNPGrouped requires optimizer.init_args.")
    init_args = dict(optimizer_config.get("init_args", {}))
    init_args.pop("fused", None)
    rmnp_params = []
    adam_params = []
    for name, param in model.named_parameters():
        # The authors assign every hidden matrix to RMNP and keep embeddings,
        # lm_head, vectors, and scalars on Adam. LitGPT calls its embedding
        # `transformer.wte`, so exclude that exact semantic equivalent too.
        is_embedding_or_head = "embed" in name or ".wte." in name or name.endswith("lm_head.weight")
        if param.ndim >= 2 and not is_embedding_or_head:
            rmnp_params.append(param)
        else:
            adam_params.append(param)
    param_groups = []
    if rmnp_params:
        param_groups.append({"params": rmnp_params, "is_rmnp": True})
    if adam_params:
        param_groups.append({"params": adam_params, "is_rmnp": False})
    if not rmnp_params:
        raise ValueError("RMNPGrouped found no hidden matrix parameters.")
    return RMNPGrouped(param_groups, **init_args)


def instantiate_adaptive_row_tadamw(optimizer_config, model: nn.Module) -> AdaptiveRowTAdamW:
    """Apply adaptive row correction only to hidden matrix parameters."""
    if not isinstance(optimizer_config, dict):
        raise ValueError("AdaptiveRowTAdamW requires optimizer.init_args.")
    init_args = dict(optimizer_config.get("init_args", {}))
    init_args.pop("fused", None)
    matrix_params = []
    base_params = []
    for name, param in model.named_parameters():
        is_embedding_or_head = "embed" in name or ".wte." in name or name.endswith("lm_head.weight")
        if param.ndim >= 2 and not is_embedding_or_head:
            matrix_params.append(param)
        else:
            base_params.append(param)
    param_groups = []
    if matrix_params:
        param_groups.append({"params": matrix_params, "use_row_blend": True})
    if base_params:
        param_groups.append({"params": base_params, "use_row_blend": False})
    return AdaptiveRowTAdamW(param_groups, **init_args)


def instantiate_official_muon_optimizer(optimizer_config, model: nn.Module) -> OfficialMuonWithAuxAdam:
    """Map hidden matrices to Muon and embeddings/head/vectors to auxiliary Adam."""
    if not isinstance(optimizer_config, dict):
        raise ValueError("OfficialMuonWithAuxAdam requires optimizer.init_args.")
    init_args = dict(optimizer_config.get("init_args", {}))
    init_args.pop("fused", None)
    muon_params = []
    per_head_muon_params: list[tuple[nn.Parameter, int]] = []
    adam_params = []
    for name, param in model.named_parameters():
        is_embedding_or_head = "embed" in name or ".wte." in name or name.endswith("lm_head.weight")
        if param.ndim >= 2 and not is_embedding_or_head:
            # Kimi K3 section 2.5 orthogonalizes each Q/K/V head block
            # independently instead of coupling all heads in one matrix.
            # FLA's KDA projections lay heads contiguously on output axis 0.
            is_kda_qkv = (
                ".attn.attn." in name
                and name.endswith(("q_proj.weight", "k_proj.weight", "v_proj.weight"))
                and param.ndim == 2
            )
            if is_kda_qkv:
                num_heads = int(model.config.kda_num_heads)
                per_head_muon_params.append((param, num_heads))
            else:
                muon_params.append(param)
        else:
            adam_params.append(param)
    param_groups = []
    if muon_params:
        param_groups.append({"params": muon_params, "use_muon": True})
    for param, num_heads in per_head_muon_params:
        param_groups.append(
            {
                "params": [param],
                "use_muon": True,
                "muon_head_splits": num_heads,
            }
        )
    if adam_params:
        param_groups.append({"params": adam_params, "use_muon": False})
    return OfficialMuonWithAuxAdam(param_groups, **init_args)

def fit(
    fabric: L.Fabric,
    devices: int,
    state: dict,
    train_dataloader: DataLoader,
    val_dataloader: DataLoader,
    out_dir: Path,
    tokenizer_dir: Path | None,
    train: TrainArgs,
    eval: EvalArgs,
    num_nodes: int = 1,
) -> None:
    model = state["model"]
    optimizer = state["optimizer"]

    if eval.initial_validation:
        val_loss = validate(fabric, model, val_dataloader, max_iters=eval.max_iters)
        val_loss = f"{val_loss:.3f}"
    else:
        fabric.print("Verifying settings ...")
        validate(fabric, model, val_dataloader, max_iters=2, verbose=False)  # sanity check
        val_loss = "n/a"

    throughput = ThroughputMonitor(fabric, window_size=5)

    if (
        model.config.mlp_class_name == "BlockSparseAdaptiveDSwiGLUMLP"
        or model.config.multiscreen_enabled
        or model.config.mamba3_enabled
        or model.config.rwkv7_enabled
        or model.config.kda_enabled
    ):
        measured_flops = 0
        fabric.print("Measured TFLOPs: skipped for custom recurrent/screening kernel")
    else:
        with torch.device("meta"):
            meta_model = GPT(model.config)
            x = torch.randint(0, 1, (train.micro_batch_size, meta_model.max_seq_length))
            model_fwd = lambda: meta_model(x)  # noqa: F821
            model_loss = lambda y: chunked_cross_entropy(y, x, chunk_size=0)  # noqa: F821
            measured_flops = measure_flops(meta_model, model_fwd, model_loss)
            fabric.print(f"Measured TFLOPs: {measured_flops * fabric.world_size / 1e12:.2f}")
            del meta_model, x

    max_tokens_per_device = train.max_tokens // fabric.world_size
    tokens_per_iter = train.micro_batch_size * model.max_seq_length
    max_iters = max_tokens_per_device // tokens_per_iter
    log_iter_interval = train.log_interval * train.gradient_accumulation_iters(devices, num_nodes)
    initial_iter = state["iter_num"]
    train_iterator = CycleIterator(train_dataloader)

    running_loss = RunningMean(window=train.gradient_accumulation_iters(devices, num_nodes), sync_on_compute=False).to(
        fabric.device
    )
    fabric.barrier()
    total_t0 = time.perf_counter()

    warmup_iters = train.warmup_iters(devices, num_nodes, max_iters, train_dataloader)
    sliding_window_enabled: bool | None = None

    for train_data in train_iterator:
        if state["iter_num"] >= max_iters:
            break

        if model.config.sliding_window_curriculum == "qg_global_qg":
            progress = (state["iter_num"] + 1) / max_iters
            next_sliding_window_enabled = (
                progress <= model.config.sliding_window_curriculum_early_fraction
                or progress >= model.config.sliding_window_curriculum_late_fraction
            )
            if next_sliding_window_enabled != sliding_window_enabled:
                changed = model.set_sliding_window_enabled(next_sliding_window_enabled)
                phase = "QG-windowed" if next_sliding_window_enabled else "all-global"
                fabric.print(
                    f"Sliding-window curriculum -> {phase} at progress={progress:.6f} "
                    f"(attention modules changed={changed})"
                )
                sliding_window_enabled = next_sliding_window_enabled

        fade_start = model.config.depth_memory_fade_start_fraction
        if model.config.depth_memory_mode == "adaptive" and fade_start < 1.0:
            progress = (state["iter_num"] + 1) / max_iters
            if progress <= fade_start:
                depth_memory_scale = 1.0
            else:
                fade_progress = (progress - fade_start) / (1.0 - fade_start)
                depth_memory_scale = 0.5 * (1.0 + math.cos(math.pi * fade_progress))
            model.depth_memory_scale.fill_(depth_memory_scale)

        # determine and set the learning rate for this iteration
        lr = get_lr(
            optimizer.defaults["lr"],
            state["iter_num"],
            warmup_iters,
            max_iters,
            train.min_lr,
            train.lr_schedule,
            train.lr_decay_start_fraction,
        )
        for param_group in optimizer.param_groups:
            group_base_lr = param_group.get("schedule_base_lr")
            if group_base_lr is None:
                param_group["lr"] = lr
            else:
                reference_lr = optimizer.defaults["lr"]
                group_min_lr = train.min_lr * group_base_lr / reference_lr
                param_group["lr"] = get_lr(
                    group_base_lr,
                    state["iter_num"],
                    warmup_iters,
                    max_iters,
                    group_min_lr,
                    train.lr_schedule,
                    train.lr_decay_start_fraction,
                )

        state["iter_num"] += 1
        iter_t0 = time.perf_counter()

        input_ids = train_data[:, 0 : model.max_seq_length].contiguous().long()
        targets = train_data[:, 1 : (model.max_seq_length + 1)].contiguous().long()

        is_accumulating = state["iter_num"] % train.gradient_accumulation_iters(devices, num_nodes) != 0
        with fabric.no_backward_sync(model, enabled=is_accumulating):
            use_multiscreen_chunked_loss = (
                model.config.multiscreen_chunked_lm_loss
                and train.mtp_loss_weight == 0.0
                and train.nta_margin_loss_weight == 0.0
            )
            if use_multiscreen_chunked_loss:
                loss = model(input_ids, targets=targets)
                logits = None
            elif (
                train.mtp_loss_weight > 0.0
                and model.config.multiscreen_chunked_lm_loss
            ):
                loss, hidden = model(
                    input_ids,
                    return_hidden=True,
                    targets=targets,
                )
                logits = None
            elif train.mtp_loss_weight > 0.0:
                logits, hidden = model(input_ids, return_hidden=True)
            else:
                logits = model(input_ids)
            if (
                not use_multiscreen_chunked_loss
                and not (
                    train.mtp_loss_weight > 0.0
                    and model.config.multiscreen_chunked_lm_loss
                )
            ):
                loss = chunked_cross_entropy(logits, targets)
            if train.mtp_loss_weight > 0.0:
                mtp_targets = train_data[:, 2 : (model.max_seq_length + 1)].contiguous().long()
                mtp_inputs = train_data[:, 1 : model.max_seq_length].contiguous().long()
                if model.config.multiscreen_chunked_lm_loss:
                    mtp_loss = model.mtp_forward(
                        hidden[:, : mtp_targets.shape[1], :],
                        mtp_inputs,
                        targets=mtp_targets,
                    )
                else:
                    mtp_logits = model.mtp_forward(
                        hidden[:, : mtp_targets.shape[1], :],
                        mtp_inputs,
                    )
                    mtp_loss = chunked_cross_entropy(mtp_logits, mtp_targets)
                loss = loss + train.mtp_loss_weight * mtp_loss
            if train.nta_margin_loss_weight > 0.0:
                loss = loss + train.nta_margin_loss_weight * next_token_margin_loss(
                    logits,
                    targets,
                    train.nta_margin,
                    error_only=train.nta_margin_error_only,
                )
            fabric.backward(loss / train.gradient_accumulation_iters(devices, num_nodes))

        running_loss.update(loss.detach())

        if not is_accumulating:
            fabric.clip_gradients(model, optimizer, max_norm=train.max_norm)
            optimizer.step()
            optimizer.zero_grad()
            state["step_count"] += 1

        if state["iter_num"] % log_iter_interval == 0:
            loss = running_loss.compute().item()  # expensive device-to-host synchronization
            t1 = time.perf_counter()
            throughput.update(
                time=(t1 - total_t0),
                flops=(measured_flops * log_iter_interval),
                batches=state["iter_num"],
                samples=(state["iter_num"] * train.micro_batch_size),
                lengths=(state["iter_num"] * train.micro_batch_size * model.max_seq_length),
            )
            metrics = {
                "loss": loss,
                "iter": state["iter_num"],
                "step": state["step_count"],
                "epoch": train_iterator.epoch,
                "iter_time": t1 - iter_t0,
                "remaining_time": (
                    (t1 - total_t0) / (state["iter_num"] - initial_iter) * (max_iters - state["iter_num"])
                ),
                "tokens": state["iter_num"] * train.micro_batch_size * model.max_seq_length,
                "total_tokens": (state["iter_num"] * train.micro_batch_size * model.max_seq_length * fabric.world_size),
                "learning_rate": lr,
            }
            if isinstance(val_loss, float):
                val_loss = f"{val_loss:.3f}"
            fabric.print(
                f"Epoch {metrics['epoch'] + 1} | iter {metrics['iter']} step {metrics['step']} |"
                f" loss train: {metrics['loss']:.3f},"
                f" val: {val_loss} |"
                f" iter time: {metrics['iter_time'] * 1000:.2f} ms"
                f"{' (step)' if not is_accumulating else ''}"
                f" remaining time: {timedelta(seconds=int(metrics['remaining_time']))!s}"
            )

            throughput_metrics = throughput.compute()
            metrics.update(throughput_metrics)
            fabric.log_dict(metrics, step=state["iter_num"] - 1)

        if val_dataloader is not None and not is_accumulating and state["step_count"] % eval.interval == 0:
            t0 = time.perf_counter()
            val_loss = validate(fabric, model, val_dataloader, max_iters=eval.max_iters)
            val_loss = val_loss.item()
            td = time.perf_counter() - t0

            fabric.print(f"iter {state['iter_num']}: val loss {val_loss:.4f}, val time: {td * 1000:.2f} ms")
            metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
            fabric.log_dict(metrics, step=state["iter_num"] - 1)
            fabric.barrier()

        if train.save_interval is not None and not is_accumulating and state["step_count"] % train.save_interval == 0:
            save_checkpoint(fabric, state, tokenizer_dir, out_dir / f"step-{state['step_count']:08d}" / "lit_model.pth")

    # Final validation
    if eval.final_validation:
        val_loss = validate(fabric, model, val_dataloader, max_iters=eval.max_iters)
        metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
        fabric.log_dict(metrics, step=state["iter_num"])
        fabric.print(f"Final evaluation | val loss: {val_loss.item():.3f} | val ppl: {math.exp(val_loss):.3f}")


@torch.no_grad()
def validate(
    fabric: L.Fabric, model: nn.Module, val_dataloader: DataLoader, max_iters: int, verbose: bool = True
) -> torch.Tensor:
    fabric.barrier()
    if verbose:
        fabric.print("Validating ...")
    model.eval()

    losses = []
    for k, batch in enumerate(val_dataloader):
        if k >= max_iters:
            break
        input_ids = batch[:, 0 : model.max_seq_length].contiguous().long()
        targets = batch[:, 1 : (model.max_seq_length + 1)].contiguous().long()
        if model.config.mamba3_hierarchical_vocab:
            # The two 64-way cross-entropies are exactly the NLL of the
            # Cartesian-product 4096-way distribution. Avoid materializing a
            # B x T x 4096 validation tensor at throughput microbatches.
            loss = model(input_ids, targets=targets)
        else:
            logits = model(input_ids)
            loss = chunked_cross_entropy(logits, targets)
        losses.append(loss)

    val_loss = torch.stack(losses).mean()
    model.train()
    fabric.barrier()
    return val_loss


def get_dataloaders(
    fabric: L.Fabric, data: DataModule, tokenizer: Tokenizer, train: TrainArgs, block_size: int
) -> tuple[DataLoader, DataLoader]:
    data.connect(tokenizer=tokenizer, batch_size=train.micro_batch_size, max_seq_length=block_size)
    with fabric.rank_zero_first():
        data.prepare_data()
    data.setup()
    train_dataloader = data.train_dataloader()
    val_dataloader = data.val_dataloader()
    return train_dataloader, val_dataloader


# learning rate decay scheduler
def next_token_margin_loss(
    logits: torch.Tensor,
    targets: torch.Tensor,
    margin: float = 0.0,
    *,
    error_only: bool = False,
) -> torch.Tensor:
    """Hinge loss that rewards the target logit outranking the strongest wrong logit.

    ``error_only`` concentrates the auxiliary gradient on positions currently
    counted as NTA errors, avoiding extra margin pressure on already-correct
    predictions while cross-entropy remains active everywhere.
    """
    top_values, top_indices = logits.float().topk(2, dim=-1)
    target_logits = logits.float().gather(-1, targets.unsqueeze(-1)).squeeze(-1)
    top_is_target = top_indices[..., 0].eq(targets)
    strongest_wrong = torch.where(top_is_target, top_values[..., 1], top_values[..., 0])
    penalties = torch.relu(float(margin) + strongest_wrong - target_logits)
    if not error_only:
        return penalties.mean()
    error_mask = (~top_is_target).to(dtype=penalties.dtype)
    return (penalties * error_mask).sum() / error_mask.sum().clamp_min(1.0)



def get_lr(
    learning_rate: float,
    it: int,
    warmup_iters: int,
    max_iters: int,
    min_lr: float,
    schedule: str = "cosine",
    decay_start_fraction: float = 0.7,
) -> float:
    if schedule == "wsd":
        if warmup_iters > 0 and it < warmup_iters:
            return learning_rate * it / warmup_iters
        decay_start = max(warmup_iters, int(max_iters * decay_start_fraction))
        if it < decay_start:
            return learning_rate
        if it > max_iters:
            return min_lr
        decay_span = max(1, max_iters - decay_start)
        decay_ratio = (it - decay_start) / decay_span
        decay_ratio = min(max(decay_ratio, 0.0), 1.0)
        coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
        return min_lr + coeff * (learning_rate - min_lr)

    if schedule == "onecycle":
        if warmup_iters > 0 and it < warmup_iters:
            return min_lr + (learning_rate - min_lr) * it / warmup_iters
        if it > max_iters:
            return min_lr
        decay_span = max(1, max_iters - warmup_iters)
        decay_ratio = (it - warmup_iters) / decay_span
        decay_ratio = min(max(decay_ratio, 0.0), 1.0)
        coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
        return min_lr + coeff * (learning_rate - min_lr)

    # 1) linear warmup for warmup_iters steps
    if it < warmup_iters:
        return learning_rate * it / warmup_iters
    # 2) if it > max_iters, return min learning rate
    if it > max_iters:
        return min_lr
    # 3) in between, use cosine decay down to min learning rate
    decay_ratio = (it - warmup_iters) / (max_iters - warmup_iters)
    assert 0 <= decay_ratio <= 1
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))  # coeff ranges 0..1
    return min_lr + coeff * (learning_rate - min_lr)


def initialize_weights(fabric: L.Fabric, model: GPT, n_layer: int, n_embd: int) -> None:
    """GPT-NeoX weight initialization (https://arxiv.org/abs/2204.06745)."""
    # Adapted from https://github.com/jzhang38/TinyLlama

    def init_weights(module, std):
        nn.init.normal_(module.weight, mean=0.0, std=std)
        if getattr(module, "bias", None) is not None:
            nn.init.zeros_(module.bias)

    def init_block_sparse_weights(module, std, down_std):
        nn.init.normal_(module.up_weight, mean=0.0, std=std)
        nn.init.normal_(module.value_weight, mean=0.0, std=std)
        nn.init.normal_(module.down_weight, mean=0.0, std=down_std)
        if hasattr(module, "router"):
            init_weights(module.router, std)
        if hasattr(module, "threshold"):
            nn.init.zeros_(module.threshold.weight)
            init = min(max(module.config.drelu_threshold_init, 1e-4), 1 - 1e-4)
            nn.init.constant_(module.threshold.bias, math.log(init / (1 - init)))
        if hasattr(module, "direction_input"):
            # LoRA-style zero-output start: the base Euler+Sandwich function is
            # exact at step zero, while random input/router factors break route
            # symmetry and allow the conditional residual to grow by gradient.
            nn.init.normal_(module.direction_input, mean=0.0, std=0.02)
            nn.init.zeros_(module.direction_up)
            nn.init.zeros_(module.direction_value)
            init_weights(module.direction_router, 0.02)
            module.direction_strength_raw.data.fill_(-2.0)

    if model.config.initializer_range is not None:
        std = model.config.initializer_range
        for mod in model.modules():
            if isinstance(mod, (nn.Embedding, nn.Linear)):
                mod.reset_parameters = partial(init_weights, mod, std=std)
            elif isinstance(
                mod,
                (
                    BlockSparseAdaptiveDSwiGLUMLP,
                    TileRoutedChannelGroupStateMemoryDSwiGLUMLP,
                    TileRoutedChannelMemoryDSwiGLUMLP,
                    TileRoutedChannelStateMemoryDSwiGLUMLP,
                    TileRoutedCenteredChannelGroupStateMemoryDSwiGLUMLP,
                    TileRoutedCenteredChannelStateMemoryDSwiGLUMLP,
                    TileRoutedDSwiGLUMLP,
                    TileRoutedGroupStateMemoryDSwiGLUMLP,
                    TileRoutedKeyedChannelMemoryDSwiGLUMLP,
                    TileRoutedDSwiGLUMLPStaticA2,
                    TileRoutedDSwiGLUMLPStaticGPTS,
                ),
            ):
                mod.reset_parameters = partial(init_block_sparse_weights, mod, std=std, down_std=std)

        if not isinstance(fabric.strategy, FSDPStrategy):
            reset_parameters(model)
        return

    for mod in model.modules():
        if isinstance(mod, (nn.Embedding, nn.Linear)):
            mod.reset_parameters = partial(init_weights, mod, std=math.sqrt(2.0 / 5 / n_embd))

    if model.config.mamba3_enabled:
        def init_mamba_input(module):
            nn.Linear.reset_parameters(module)

        def init_mamba_output(module):
            nn.Linear.reset_parameters(module)
            with torch.no_grad():
                module.weight.div_(math.sqrt(n_layer))

        # Preserve the official Mamba initialization instead of replacing its
        # selective-state projections with GPT-NeoX normal initialization.
        for block in model.transformer.h:
            block.mixer.in_proj.reset_parameters = partial(init_mamba_input, block.mixer.in_proj)
            block.mixer.out_proj.reset_parameters = partial(init_mamba_output, block.mixer.out_proj)

    # need a separate loop because the residual projection linears below are `nn.Linear` too
    for mod in model.modules():
        if isinstance(mod, BlockSparseAdaptiveDSwiGLUMLP):
            mod.reset_parameters = partial(
                init_block_sparse_weights,
                mod,
                std=math.sqrt(2.0 / 5 / n_embd),
                down_std=(1 / math.sqrt(n_embd) / n_layer),
            )
        if isinstance(
            mod,
            (
                TileRoutedChannelGroupStateMemoryDSwiGLUMLP,
                TileRoutedChannelMemoryDSwiGLUMLP,
                TileRoutedChannelStateMemoryDSwiGLUMLP,
                TileRoutedCenteredChannelGroupStateMemoryDSwiGLUMLP,
                TileRoutedCenteredChannelStateMemoryDSwiGLUMLP,
                TileRoutedDSwiGLUMLP,
                TileRoutedGroupStateMemoryDSwiGLUMLP,
                TileRoutedKeyedChannelMemoryDSwiGLUMLP,
                TileRoutedDSwiGLUMLPStaticA2,
                TileRoutedDSwiGLUMLPStaticGPTS,
            ),
        ):
            mod.reset_parameters = partial(
                init_block_sparse_weights,
                mod,
                std=math.sqrt(2.0 / 5 / n_embd),
                down_std=(1 / math.sqrt(n_embd) / n_layer),
            )
        if isinstance(
            mod,
            (LLaMAMLP, DSwiGLUMLP, AdaptiveDSwiGLUMLP, CausalSelfAttention),
        ):
            mod.proj.reset_parameters = partial(init_weights, mod.proj, std=(1 / math.sqrt(n_embd) / n_layer))

    if not isinstance(fabric.strategy, FSDPStrategy):
        reset_parameters(model)


def save_checkpoint(fabric, state, tokenizer_dir, checkpoint_file):
    model = state["model"]
    checkpoint_file.parent.mkdir(parents=True, exist_ok=True)
    fabric.print(f"Saving checkpoint to {str(checkpoint_file)!r}")
    fabric.save(checkpoint_file, state)
    if fabric.global_rank == 0:
        save_hyperparameters(setup, checkpoint_file.parent)
        if tokenizer_dir is not None:
            copy_config_files(tokenizer_dir, checkpoint_file.parent)
        save_config(model.config, checkpoint_file.parent)


def validate_args(train: TrainArgs, eval: EvalArgs, initial_checkpoint_dir, resume) -> None:
    issues = []
    unsupported = [(train, ["epochs"]), (eval, ["max_new_tokens"])]
    for args, names in unsupported:
        for name in names:
            if getattr(args, name) is not None:
                issues.append(f"{__file__} doesn't support the {name!r} argument. This is set in {args}")
    if train.max_steps is not None:
        warnings.warn(
            "`train.max_steps` is intended for profiling or debug runs only. "
            "For full pretraining runs, prefer `train.max_tokens` or `train.max_time`.",
            UserWarning,
        )
    required = [(train, ["max_tokens", "max_norm"])]
    for args, names in required:
        for name in names:
            if getattr(args, name) is None:
                issues.append(f"{__file__} requires the {name!r} argument. This is set in {args}")
    if initial_checkpoint_dir and resume:
        issues.append("Can't provide both `--resume` and `--initial_checkpoint_dir`. Choose one.")
    if issues:
        raise ValueError("\n".join(issues))