diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 00000000..e3a3d828 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,17 @@ +# Build context for alto/Dockerfile.multinode is this directory (ALTO/). +# Keep large / irrelevant paths out of the context sent to the Docker daemon. +# gptoss_chkpt alone is ~9.6TB and must never be tarred into the build context. +.git +gptoss_chkpt/ +logs/ +comm_traces/ +plotting/ +examples/ +tests/ +slurm-*.out +*.out +error*.txt +sendit.toml +wait_and_run.sh +**/__pycache__/ +**/*.pyc diff --git a/.gitignore b/.gitignore index 74a24a16..0d7f7cfa 100644 --- a/.gitignore +++ b/.gitignore @@ -212,3 +212,19 @@ __marimo__/ datasets/ /models/ comm_traces/ + + +# Slurm +*.out + +# Checkpoints +*.distcp +gptoss_chkpt/ +plotting/plots + +*.safetensors + +docker/ +gpt_*/ +*/plotting/ +docs/ diff --git a/.gitmodules b/.gitmodules index 826c33ce..f892dbd9 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,7 @@ [submodule "3rdparty/torchtitan"] path = 3rdparty/torchtitan url = https://github.com/hann-wang/torchtitan +[submodule "3rdparty/adahop"] + path = 3rdparty/adahop + url = git@github.com:AMDResearch/low-precision-training.git + branch = AdaHOP_C42 diff --git a/3rdparty/adahop b/3rdparty/adahop new file mode 160000 index 00000000..38b47508 --- /dev/null +++ b/3rdparty/adahop @@ -0,0 +1 @@ +Subproject commit 38b475083df70366912c73d72069910acc2cb5c7 diff --git a/Dockerfile.multinode b/Dockerfile.multinode new file mode 100644 index 00000000..66a60c7a --- /dev/null +++ b/Dockerfile.multinode @@ -0,0 +1,98 @@ +# torch 2.11 (+ matching torchvision/triton) is required by the bundled torchtitan, +# which uses torch>=2.11 APIs (torch.nn.attention.varlen, the +# wrap_inductor_compiled_regions inductor option, ...). There is no prebuilt +# torch 2.11 + ROCm 7.2 image, so this pins the validated ROCm 7.14 + torch 2.11.0 +# build. NOTE: this bumps ROCm 7.2 -> 7.14; the Slurm hosts' amdgpu/KFD driver must +# support a ROCm 7.14 container userspace. +FROM rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.11.0 + +ARG DEBIAN_FRONTEND=noninteractive + +# Port 80 to archive.ubuntu.com/security.ubuntu.com is blocked in this build +# environment, but 443 works. Rewrite every apt source (including the deb822 +# /etc/apt/sources.list.d/ubuntu.sources used on noble) from http:// to https://. + +# Packages +# ibverbs-utils for debug via to `rdma link show` +# iproute2 debug ip addresses such as RDMA ips +# perftest RDMA debug + +RUN find /etc/apt -type f \ + \( -name '*.list' -o -name '*.sources' \) \ + -exec sed -i 's|http://|https://|g' {} + \ + && apt-get \ + -o Acquire::ForceIPv4=true \ + -o Acquire::Retries=5 \ + -o Acquire::https::Timeout=30 \ + update \ + && apt-get install -y \ + git-lfs \ + pkg-config \ + clang \ + libclang-dev \ + libunwind-dev \ + libnl-3-dev \ + libnl-route-3-dev \ + libibverbs-dev \ + ibverbs-utils \ + iproute2 \ + perftest \ + ibverbs-providers \ + cmake \ + pciutils \ + && (update-pciids || true) +# /var/lib/apt/lists is intentionally NOT removed, to keep apt usable for debugging. + +RUN pip install --no-cache-dir huggingface_hub "datasets>=3.6.0" \ + transformers tabulate wandb fsspec tyro "tokenizers>=0.15.0" safetensors \ + tensorboard pre-commit yapf pybind11 meson-python torchdata pytablewriter \ + "antlr4-python3-runtime==4.11.0" sympy math_verify more_itertools peft \ + accelerate pillow "numpy<2" opencv-python-headless scipy \ + numba huggingface-hub[cli,hf_transfer] "packaging>=24.2" \ + "setuptools>=77.0.3,<80.0.0" "setuptools-scm>=8" \ + protobuf-protoc-bin fmt + +RUN mkdir -p /var/lib/jenkins && \ + cd /var/lib/jenkins && \ + git clone --depth 1 https://github.com/EleutherAI/lm-evaluation-harness && \ + cd lm-evaluation-harness && \ + pip install -e . + +ENV PYTORCH_ROCM_ARCH="gfx90a;gfx942;gfx950" + +# RDMA providers (the libibverbs plugins, e.g. libmlx5 / libbnxt_re) are ABI-tied +# to each host's rdma-core and kernel, which vary across our machines, so they are +# deliberately NOT baked into the image. rdma_tests/start_container.sh mounts the +# host's libibverbs userspace (library + provider modules + /etc/libibverbs.d) +# read-only at runtime, so the in-container RDMA stack always matches the host it +# runs on. The ibverbs-providers/libibverbs-dev apt packages above remain as a +# self-contained fallback for same-ABI hosts (USE_HOST_RDMA=0). + +RUN FSDP_PARAM=$(python3 -c "import torch, os; print(os.path.join(os.path.dirname(torch.__file__), 'distributed/fsdp/_fully_shard/_fsdp_param.py'))") && \ + sed -i 's/self.sharded_param = nn.Parameter(self.to_sharded_dtensor(sharded_param))/self.sharded_param = nn.Parameter(self.to_sharded_dtensor(sharded_param), requires_grad=param.requires_grad)/' "$FSDP_PARAM" && \ + sed -i 's/ self.sharded_param.requires_grad_(param.requires_grad)//' "$FSDP_PARAM" + +# Install torchtitan and ALTO training deps at build time (as root) to avoid +# /opt/venv permission errors when the container runs as a non-root user. +COPY 3rdparty/torchtitan /tmp/torchtitan_src +# Pin the ROCm torch/torchvision/triton already in the base image via a pip +# constraints file so no transitive dependency (compressed_tensors -> transformers, +# torchao, torchtitan, ...) can pull a CUDA build of torch and the nvidia-* wheels. +# Replacing the ROCm torch breaks the pre-compiled ROCm torchvision ops at runtime +# ("RuntimeError: operator torchvision::nms does not exist" -> ABI mismatch). +# Constraints still let other packages upgrade; they just cannot move torch et al. +# torchao is also installed with --no-deps as belt-and-suspenders. meson-python/ +# pybind11/ninja are torchtitan's declared build backend (build-backend = "mesonpy") +# and must be present here because the torchtitan install uses --no-build-isolation. +# Read versions from package metadata (not `pip freeze`): the ROCm torch is a +# locally-built wheel that pip freeze renders as "torch @ file://...", which a +# "^torch==" grep would miss, producing an empty constraints file. +RUN python3 -c "import importlib.metadata as md; have={d.metadata['Name'].lower() for d in md.distributions()}; want=['torch','torchvision','triton','pytorch-triton-rocm']; open('/tmp/torch-constraints.txt','w').write(''.join(f'{p}=={md.version(p)}\n' for p in want if p.lower() in have))" && \ + cat /tmp/torch-constraints.txt && \ + pip install --no-cache-dir -c /tmp/torch-constraints.txt aim \ + compressed_tensors easydict loguru \ + poetry-core "poetry-dynamic-versioning>=1.0.0,<2.0.0" \ + meson-python pybind11 ninja tyro && \ + pip install --no-cache-dir -c /tmp/torch-constraints.txt --no-deps torchao && \ + pip install --no-cache-dir --no-build-isolation --no-deps /tmp/torchtitan_src && \ + rm -rf /tmp/torchtitan_src /tmp/torch-constraints.txt diff --git a/RDMA.md b/RDMA.md new file mode 100644 index 00000000..674bfa4b --- /dev/null +++ b/RDMA.md @@ -0,0 +1,30 @@ +# Testing RDMA between two nodes + +```bash +`ibv_devices` # list RDMA devices (e.g. mlx5_0) +`ibv_devinfo` # PortState: PORT_ACTIVE and note the link layer (IB vs. Ethernet/ROCE) +`rdma link show` # link state per device +``` + +# Raw RDMA loopback between two nodes + +```bash +# On node A (server): +ib_write_bw -d mlx5_0 -F --report_gbits + +# On node B (client), point at node A's IP: +ib_write_bw -d mlx5_0 -F --report_gbits + +``` + +You can get `` for a given RDMA interface like this: + +```bash +# Map RDMA device -> netdev: +ibdev2netdev # look for a line like: "mlx5_0 port 1 ==> rdma0 (Up)" +rdma link show # alternatively, look here. e.g. mlx5_0/1 ... netdev rdma0 + +# Then get that interface's IP: +ip -4 addr show rdma0 # look for the "inet x.x.x.x" line + +``` \ No newline at end of file diff --git a/alto/_adahop_bridge.py b/alto/_adahop_bridge.py new file mode 100644 index 00000000..3b7809a9 --- /dev/null +++ b/alto/_adahop_bridge.py @@ -0,0 +1,172 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Bridge to the AdaHOP submodule at ``3rdparty/adahop``. + +Loads selected AdaHOP modules by file path under non-colliding names so that +AdaHOP's vendored ``torchtitan`` package never enters ``sys.modules`` and +cannot shadow ALTO's own torchtitan submodule. +""" + +import importlib.util +import sys +from pathlib import Path +from typing import Any + +_ALTO_ROOT = Path(__file__).resolve().parent.parent +_ADAHOP_ROOT = _ALTO_ROOT / "3rdparty" / "adahop" +_HT_DIR = _ADAHOP_ROOT / "torchtitan" / "experiments" / "kernels" / "hadamard_transform" +_MXFP4_DIR = _ADAHOP_ROOT / "torchtitan" / "experiments" / "kernels" / "mxfp4" +_TC_PATH = _MXFP4_DIR / "transform_config.py" + + +def _load_module(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ImportError(f"Cannot load module {name} from {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _load_package(name: str, pkg_dir: Path): + init_path = pkg_dir / "__init__.py" + spec = importlib.util.spec_from_file_location( + name, + init_path, + submodule_search_locations=[str(pkg_dir)], + ) + if spec is None or spec.loader is None: + raise ImportError(f"Cannot load package {name} from {pkg_dir}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +if not _ADAHOP_ROOT.exists(): + raise ImportError(f"AdaHOP submodule not found at {_ADAHOP_ROOT}. " + "Run `git submodule update --init --recursive` from the ALTO root.") + +_ht = _load_package("_alto_adahop_ht", _HT_DIR) +_tc = _load_module("_alto_adahop_transform_config", _TC_PATH) + + +def _alias_ht_at_absolute_path() -> None: + """AdaHOP's ``iht_quantization`` does an *absolute* import: + ``from torchtitan.experiments.kernels.hadamard_transform.hadamard import _build_H_b32`` + ALTO's vendored ``torchtitan`` package does not contain AdaHOP's + ``hadamard_transform`` subpackage, so the import fails. Alias the + bridge-loaded hadamard package at the absolute path the file expects. + Also stub the intermediate ``torchtitan.experiments`` and + ``torchtitan.experiments.kernels`` packages if they aren't already + populated, so the dotted lookup resolves. + """ + import types + # Build out the chain torchtitan -> .experiments -> .kernels -> .hadamard_transform + # without clobbering whatever ALTO already has installed. + if "torchtitan" not in sys.modules: + sys.modules["torchtitan"] = types.ModuleType("torchtitan") + if "torchtitan.experiments" not in sys.modules: + sys.modules["torchtitan.experiments"] = types.ModuleType("torchtitan.experiments") + if "torchtitan.experiments.kernels" not in sys.modules: + sys.modules["torchtitan.experiments.kernels"] = types.ModuleType("torchtitan.experiments.kernels") + sys.modules["torchtitan.experiments.kernels.hadamard_transform"] = _ht + # The submodule that gets cherry-picked too: + if hasattr(_ht, "hadamard"): + sys.modules["torchtitan.experiments.kernels.hadamard_transform.hadamard"] = _ht.hadamard + + +_alias_ht_at_absolute_path() + + +def _load_mxfp4_package() -> Any: + """Load AdaHOP's ``mxfp4/`` package without triggering the colliding + ``torch.ops.torchtitan.*`` registrations from ``mxfp_linear.py`` and + ``mxfp_quantization.py``. + + The package's ``__init__.py`` does ``from .mxfp_linear import MXFP4Linear`` + and ``from .mxfp_quantization import convert_to_mxfp4, convert_from_mxfp4, + BLOCK_SIZE_DEFAULT`` at import time. Both files register ``@triton_op``s + in the ``torchtitan::`` namespace that ALTO already owns (collision would + raise ``RuntimeError: Tried to register operator ... twice``). We install + stubs at the expected ``sys.modules`` names BEFORE exec'ing the package + so those import lines succeed without touching the real source. + + What we actually need from the package: ``iht_quantization``, + ``foid.OUTLIER_K``, ``foid.prepare_outlier_clean_row/column``, + ``outlier_extract.inner_outlier_extract_left_cdna4`` and + ``..._right_cdna4``. None of these go through the stubbed modules at + call time — ``outlier_extract`` calls ``torch.ops.torchtitan.`` + ``blockwise_mxfp4_gemm`` (ALTO's registration, cdna-gated) directly, + and ``iht_quantization`` self-gates cdna3/cdna4 via ``is_cdna4()`` from + ``fp4_common``. + """ + pkg_name = "_alto_adahop_mxfp4" + import types + + # Stub mxfp_linear (only exports MXFP4Linear, referenced by __init__). + mxfp_linear_stub = types.ModuleType(f"{pkg_name}.mxfp_linear") + mxfp_linear_stub.MXFP4Linear = object # placeholder, never instantiated + sys.modules[f"{pkg_name}.mxfp_linear"] = mxfp_linear_stub + + # Stub mxfp_quantization with the three symbols __init__ imports. + mxfp_quant_stub = types.ModuleType(f"{pkg_name}.mxfp_quantization") + mxfp_quant_stub.BLOCK_SIZE_DEFAULT = 32 # matches AdaHOP's real value + mxfp_quant_stub.convert_to_mxfp4 = None # never called via this stub + mxfp_quant_stub.convert_from_mxfp4 = None + sys.modules[f"{pkg_name}.mxfp_quantization"] = mxfp_quant_stub + + # Stub mxfp_grouped_gemm (only referenced by the lazy mxfp4_grouped_gemm + # wrapper, never imported eagerly). + mxfp_grouped_stub = types.ModuleType(f"{pkg_name}.mxfp_grouped_gemm") + mxfp_grouped_stub.mxfp4_grouped_gemm = None + sys.modules[f"{pkg_name}.mxfp_grouped_gemm"] = mxfp_grouped_stub + + pkg = _load_package(pkg_name, _MXFP4_DIR) + # __init__.py doesn't import foid / outlier_extract; force-load them + # via the package's submodule search so they hang off `pkg` as attributes. + import importlib + pkg.foid = importlib.import_module(f"{pkg_name}.foid") + pkg.outlier_extract = importlib.import_module(f"{pkg_name}.outlier_extract") + return pkg + + +_mxfp4 = _load_mxfp4_package() + +HadamardFactory = _ht.HadamardFactory +HadamardTransform = _ht.HadamardTransform +detect_outlier_pattern = _ht.detect_outlier_pattern + +transform_config = _tc +configure_global_transforms = _tc.configure_global_transforms +configure_layer_transforms = _tc.configure_layer_transforms +get_layer_transform_config = _tc.get_layer_transform_config +should_apply_transform = _tc.should_apply_transform +clear_all_configs = _tc.clear_all_configs + +iht_quantization = _mxfp4.iht_quantization +OUTLIER_K = _mxfp4.foid.OUTLIER_K +prepare_outlier_clean_row = _mxfp4.foid.prepare_outlier_clean_row +prepare_outlier_clean_column = _mxfp4.foid.prepare_outlier_clean_column +inner_outlier_extract_left_cdna4 = _mxfp4.outlier_extract.inner_outlier_extract_left_cdna4 +inner_outlier_extract_right_cdna4 = _mxfp4.outlier_extract.inner_outlier_extract_right_cdna4 + +__all__ = [ + "HadamardFactory", + "HadamardTransform", + "detect_outlier_pattern", + "transform_config", + "configure_global_transforms", + "configure_layer_transforms", + "get_layer_transform_config", + "should_apply_transform", + "clear_all_configs", + "iht_quantization", + "OUTLIER_K", + "prepare_outlier_clean_row", + "prepare_outlier_clean_column", + "inner_outlier_extract_left_cdna4", + "inner_outlier_extract_right_cdna4", +] diff --git a/alto/components/__init__.py b/alto/components/__init__.py index eb0ad187..96870f51 100644 --- a/alto/components/__init__.py +++ b/alto/components/__init__.py @@ -3,12 +3,15 @@ # SPDX-License-Identifier: MIT from .converter import ModelOptConverter +from .m_adam import MAdamOptimizersContainer, m_adam from .optimizer import DeOscillationConfig, enable_de_oscillation from .state_dict_adapter_mixin import StateDictAdapterMixin __all__ = [ "DeOscillationConfig", "enable_de_oscillation", + "MAdamOptimizersContainer", + "m_adam", "ModelOptConverter", "StateDictAdapterMixin", ] diff --git a/alto/components/m_adam.py b/alto/components/m_adam.py new file mode 100644 index 00000000..06f8df11 --- /dev/null +++ b/alto/components/m_adam.py @@ -0,0 +1,440 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT + +"""``m_adam``: a hybrid additive/multiplicative optimizer and its +torchtitan :class:`OptimizersContainer` wiring. + +``m_adam`` decomposes every weight via ``torch.frexp`` into a mantissa and +an exponent (``w = m * 2**e``) and applies two coupled updates each step: + +* an **AdamW** (additive) update on the mantissa, and +* a **Madam-style** (multiplicative, RMSProp-normalized) update on the + exponent -- i.e. gradient descent on the log2-magnitude. + +Net effect (with ``weight_decay_e == 0``):: + + w_new ~= w * 2**(delta_e) + delta_w_AdamW + \\_ multiplicative (Madam) \\_ additive (AdamW) + +This explicit mantissa/exponent split lets the exponent (dynamic range) +and mantissa (precision) be controlled with independent learning rates, +weight decays, and -- for the exponent -- an independent schedule, which +is why it is a natural fit for low-precision / quantization-aware training. + +The :class:`MAdamOptimizersContainer` adapts ``m_adam`` to torchtitan's +``config.optimizer.build(...)`` path so it can be selected from a model's +``config_registry`` just like the built-in ``Adam``/``AdamW``. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, Iterable, Optional + +import torch +import torch.nn as nn +from torch.optim.optimizer import Optimizer +from torchtitan.components.optimizer import OptimizersContainer + +__all__ = [ + "m_adam", + "MAdamOptimizersContainer", +] + + +def _rms(x: torch.Tensor) -> torch.Tensor: + return x.pow(2).mean().sqrt() + + +def _sched_value( + mode: str, + *, + base_lr: float, + t: int, + total_steps: Optional[int], + warmup_steps: int, + min_lr_ratio: float, + logcosine_alpha: float, +) -> float: + if warmup_steps > 0 and t < warmup_steps: + return base_lr * (float(t) / float(max(1, warmup_steps))) + if total_steps is None or mode == "constant": + return base_lr + + T = max(1, total_steps - max(0, warmup_steps)) + p = min(1.0, max(0.0, (t - warmup_steps) / T)) + rmin = float(min_lr_ratio) + + if mode == "linear": + return base_lr * (rmin + (1.0 - rmin) * (1.0 - p)) + if mode == "cosine": + c = 0.5 * (1.0 + math.cos(math.pi * p)) + return base_lr * (rmin + (1.0 - rmin) * c) + if mode == "logcosine": + a = float(logcosine_alpha) + g = math.exp(-a * (1.0 - math.cos(math.pi * p))) + g1 = math.exp(-2.0 * a) + s = (g - g1) / (1.0 - g1 + 1e-12) + return base_lr * (rmin + (1.0 - rmin) * s) + + return base_lr + + +class m_adam(Optimizer): + def __init__( + self, + params: Iterable[torch.nn.Parameter], + lr_m: float = 1e-3, + lr_e: float = 1e-2, + beta1: float = 0.9, + beta2: float = 0.999, + eps: float = 1e-8, + p_scale: float = 3.0, + g_bound: float = 20.0, + *, + weight_decay_m: float = 0.0, + weight_decay_e: float = 0.0, + abs_clamp: bool = False, + abs_clamp_floor: float = 0.0, + clip_e_final: bool = False, + e_final_min: float = -60.0, + e_final_max: float = 60.0, + use_de_step_cap: bool = True, + de_step_cap: float = 0.5, + tie_e_to_m: bool = False, + sched_e: str = "constant", + total_steps_e: Optional[int] = None, + warmup_steps_e: int = 0, + min_lr_ratio_e: float = 0.0, + logcosine_alpha_e: float = 6.0, + ): + if not (0.0 <= lr_m and 0.0 <= lr_e): + raise ValueError("Learning rates must be non-negative.") + if not (0.0 <= beta1 < 1.0 and 0.0 <= beta2 < 1.0): + raise ValueError("betas must be in [0,1).") + + ratio_e_init = float(lr_e) / max(1e-20, float(lr_m)) + + defaults = dict( + lr=lr_m, + beta1=beta1, + beta2=beta2, + eps=eps, + p_scale=p_scale, + g_bound=g_bound, + weight_decay_m=weight_decay_m, + weight_decay_e=weight_decay_e, + abs_clamp=abs_clamp, + abs_clamp_floor=abs_clamp_floor, + clip_e_final=clip_e_final, + e_final_min=e_final_min, + e_final_max=e_final_max, + use_de_step_cap=use_de_step_cap, + de_step_cap=de_step_cap, + tie_e_to_m=tie_e_to_m, + ratio_e_init=ratio_e_init, + sched_e=sched_e, + total_steps_e=total_steps_e, + warmup_steps_e=warmup_steps_e, + min_lr_ratio_e=min_lr_ratio_e, + lr_e_base=float(lr_e), + logcosine_alpha_e=float(logcosine_alpha_e), + t=0, + last_lr_e=float(lr_e), + last_lr_m=float(lr_m), + ) + super().__init__(params, defaults) + + @torch.no_grad() + def step(self, closure: Optional[callable] = None): + loss = closure() if closure is not None else None + + for grp in self.param_groups: + t = int(grp.get("t", 0)) + + lr_m = float(grp["lr"]) + grp["last_lr_m"] = lr_m + + if grp["tie_e_to_m"]: + lr_e = lr_m * float(grp["ratio_e_init"]) + else: + lr_e = _sched_value( + mode=grp["sched_e"], + base_lr=float(grp["lr_e_base"]), + t=t, + total_steps=grp["total_steps_e"], + warmup_steps=int(grp["warmup_steps_e"]), + min_lr_ratio=float(grp["min_lr_ratio_e"]), + logcosine_alpha=float(grp["logcosine_alpha_e"]), + ) + + grp["last_lr_e"] = float(lr_e) + + b1 = float(grp["beta1"]) + b2 = float(grp["beta2"]) + eps = float(grp["eps"]) + gmax = float(grp["g_bound"]) + + wd_m = float(grp["weight_decay_m"]) + wd_e = float(grp["weight_decay_e"]) + + clamp_w = bool(grp["abs_clamp"]) + w_floor0 = float(grp["abs_clamp_floor"]) + + clip_e = bool(grp["clip_e_final"]) + emin = float(grp["e_final_min"]) + emax = float(grp["e_final_max"]) + + cap_de = bool(grp["use_de_step_cap"]) + de_cap = float(grp["de_step_cap"]) + + for p in grp["params"]: + if p.grad is None: + continue + + dt = p.data.dtype + w = p.data + g = ( + p.grad.data.to(dt) + if p.grad.data.dtype != dt + else p.grad.data + ) + + st = self.state[p] + + if not st: + st["step_m"] = 0 + st["step_e"] = 0 + st["w_exp_avg"] = torch.zeros_like(w, dtype=dt) + st["w_exp_avg_sq"] = torch.zeros_like(w, dtype=dt) + st["exp_avg_sq"] = torch.zeros_like(w, dtype=dt) + + init = _rms(w.float()).item() + st["max"] = max( + grp["p_scale"] * (init + 1e-12), + w_floor0, + ) + + mw = st["w_exp_avg"] + vw = st["w_exp_avg_sq"] + ve = st["exp_avg_sq"] + + wf = w.float() + gw = g.float() + + m, e = torch.frexp(wf) + e_use = ( + torch.clamp(e, min=emin, max=emax) + if clip_e + else e + ) + + st["step_e"] += 1 + se = st["step_e"] + + e_stats = torch.clamp( + e_use, + min=-60.0, + max=60.0, + ).to(wf.dtype) + + w_cur = m * torch.exp2(e_stats) + ge = gw * w_cur * math.log(2.0) + + ve_f = ve.float() + ge_clip = ge.clamp(-1e19, 1e19) + + ve_f.mul_(b2).addcmul_( + ge_clip, + ge_clip, + value=1.0 - b2, + ) + + den = ( + ve_f / (1.0 - b2**se) + ).sqrt_().clamp_(min=eps) + + ge_n = (ge / den).clamp_(-gmax, gmax) + dw_e = -lr_e * ge_n + + w_floor = max( + 1e-8, + float(_rms(wf)) * 1e-6, + ) + + dabs = w_cur.abs().clamp_min(w_floor) + + r = (dw_e / dabs).clamp( + min=-0.75 + 1e-6, + max=0.75, + ) + + de = torch.log1p(r) / math.log(2.0) + + if cap_de: + de = de.clamp( + min=-de_cap, + max=de_cap, + ) + + e_new = ( + e_use.to(wf.dtype) * (1.0 - lr_e * wd_e) + + de + ) + + if clip_e: + e_new = torch.clamp( + e_new, + min=emin, + max=emax, + ) + + ve.copy_(ve_f.to(dt)) + + st["step_m"] += 1 + sm = st["step_m"] + + exp_scale = torch.exp2(-e_use.to(wf.dtype)) + gm = gw * torch.exp2(e_use.to(wf.dtype)) + gm_dt = gm.to(dt) + gm_scaled = gm_dt * exp_scale + + mw.mul_(b1).add_( + gm_scaled, + alpha=1.0 - b1, + ) + + vw.mul_(b2).addcmul_( + gm_scaled, + gm_scaled, + value=1.0 - b2, + ) + + mh = mw.float() / (1.0 - b1**sm) + vh = vw.float() / (1.0 - b2**sm) + + dw = -lr_m * mh / (vh.sqrt() + eps) + + if wd_m != 0.0: + w_adamw = ( + wf * (1.0 - lr_m * wd_m) + + dw + ) + else: + w_adamw = wf + dw + + d_w = w_adamw - wf + d_m = d_w * torch.exp2(-e_new) + m_new = m + d_m + w_new = m_new * torch.exp2(e_new) + + if clamp_w: + w_new.clamp_( + -st["max"], + st["max"], + ) + + w.copy_(w_new.to(dt)) + + grp["t"] = t + 1 + + return loss + + +class MAdamOptimizersContainer(OptimizersContainer): + """:class:`OptimizersContainer` that builds :class:`m_adam`. + + Selectable from a model ``config_registry`` via:: + + config.optimizer = MAdamOptimizersContainer.Config(lr=1e-3, lr_e=1e-2) + + The standard ``lr`` field maps to ``m_adam``'s ``lr_m`` (the additive + AdamW branch), so the usual torchtitan LR scheduler drives ``lr_m`` + for free; the exponent learning rate ``lr_e`` is scheduled + independently via ``sched_e`` (or tied to ``lr_m`` with + ``tie_e_to_m``). + + Note: the inherited ``weight_decay`` and ``implementation`` fields are + unused -- ``m_adam`` has decoupled ``weight_decay_m`` / ``weight_decay_e`` + and does not support the fused/foreach implementations. + """ + + @dataclass(kw_only=True, slots=True) + class Config(OptimizersContainer.Config): + name: str = "m_adam" + lr: float = 1e-3 + """Learning rate for the additive (AdamW / mantissa) branch (m_adam ``lr_m``).""" + + lr_e: float = 1e-2 + """Learning rate for the multiplicative (Madam / exponent) branch.""" + + beta1: float = 0.9 + beta2: float = 0.999 + eps: float = 1e-8 + + p_scale: float = 3.0 + """Magnitude-clamp bound as a multiple of the initial weight RMS (used only when ``abs_clamp``).""" + + g_bound: float = 20.0 + """Clamp on the RMS-normalized exponent gradient.""" + + weight_decay_m: float = 0.0 + weight_decay_e: float = 0.0 + + abs_clamp: bool = False + """Clamp weights to +/- p_scale * rms(init) (the original Madam safety net; off by default).""" + abs_clamp_floor: float = 0.0 + + clip_e_final: bool = False + e_final_min: float = -60.0 + e_final_max: float = 60.0 + + use_de_step_cap: bool = True + de_step_cap: float = 0.5 + + tie_e_to_m: bool = False + """If set, lr_e = lr_m * (lr_e / lr_m at init), so lr_e tracks the lr_m schedule.""" + + sched_e: str = "constant" + """Schedule for lr_e: 'constant' | 'linear' | 'cosine' | 'logcosine'.""" + total_steps_e: Optional[int] = None + warmup_steps_e: int = 0 + min_lr_ratio_e: float = 0.0 + logcosine_alpha_e: float = 6.0 + + @staticmethod + def _resolve_optimizer_cls(name: str) -> type: + if name != "m_adam": + raise NotImplementedError( + f"MAdamOptimizersContainer only builds 'm_adam', got {name!r}." + ) + return m_adam + + @staticmethod + def _build_optimizer_kwargs(config: "MAdamOptimizersContainer.Config") -> dict[str, Any]: + return { + "lr_m": config.lr, + "lr_e": config.lr_e, + "beta1": config.beta1, + "beta2": config.beta2, + "eps": config.eps, + "p_scale": config.p_scale, + "g_bound": config.g_bound, + "weight_decay_m": config.weight_decay_m, + "weight_decay_e": config.weight_decay_e, + "abs_clamp": config.abs_clamp, + "abs_clamp_floor": config.abs_clamp_floor, + "clip_e_final": config.clip_e_final, + "e_final_min": config.e_final_min, + "e_final_max": config.e_final_max, + "use_de_step_cap": config.use_de_step_cap, + "de_step_cap": config.de_step_cap, + "tie_e_to_m": config.tie_e_to_m, + "sched_e": config.sched_e, + "total_steps_e": config.total_steps_e, + "warmup_steps_e": config.warmup_steps_e, + "min_lr_ratio_e": config.min_lr_ratio_e, + "logcosine_alpha_e": config.logcosine_alpha_e, + } diff --git a/alto/kernels/dispatch/adahop_tensor.py b/alto/kernels/dispatch/adahop_tensor.py new file mode 100644 index 00000000..321f01a3 --- /dev/null +++ b/alto/kernels/dispatch/adahop_tensor.py @@ -0,0 +1,213 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""AdaHOP-aware MXFP4 wrapper tensor subclass. + +FSDP-execution fix: there is a SINGLE wrapper class. + +* :class:`MXFP4AdaHOPWrapper` — used for ``scheme="mxfp4_adahop"`` from the + moment the model is converted (so FSDP captures it as the param). It carries + per-slot transform modes ``(forward_y_mode, backward_gx_mode, + backward_gw_mode)`` and an optional ``HadamardTransform``. During calibration + all three modes are ``"none"`` and the linear dispatch defers to plain MXFP4. + When calibration finishes the modifier sets the chosen modes IN PLACE via + :meth:`set_modes` (no new ``nn.Parameter``, no wrapper-type swap), so the + modes take effect on the FSDP-owned tensor and the AdaHOP linear function + actually runs in the forward. + +Why single-wrapper + in-place modes (vs. the old two-wrapper swap): under +``fully_shard`` the param is captured at convert time. Replacing +``module.weight`` with a new Parameter at Phase B does NOT reach the tensor FSDP +all-gathers into ``F.linear`` (it keeps using the original), so AdaHOP never +executes. Keeping one wrapper and mutating its modes in place updates the +canonical FSDP-owned tensor; ``fsdp_post_all_gather`` propagates the modes to +the unsharded instance used in the forward. + +Calibration observation uses transient *module* hooks (see +``calibration_hooks.py``); nothing unpicklable is stored on the tensor. +""" + +from typing import Any, Optional, Tuple + +import torch +from torchtitan.tools.logging import logger + +from alto.modifiers.lpt.adahop_internals.mxfp4_linear_function import MXFP4AdaHOPLinearFunction +from alto.modifiers.lpt.adahop_internals.transform_mode import TransformMode, assert_mode_supported +from .tensor import MXFP4TrainingWeightWrapperTensor, gemm_ops + + +class MXFP4AdaHOPWrapper(MXFP4TrainingWeightWrapperTensor): + """Single AdaHOP wrapper: per-slot modes (default ``"none"`` == plain MXFP4). + + During calibration all modes are ``"none"``; the modifier later flips them + in place via :meth:`set_modes`. + """ + + @staticmethod + def __new__( + cls, + tensor, + config, + *, + hadamard_transform: Optional[Any] = None, + forward_y_mode: TransformMode = "none", + backward_gx_mode: TransformMode = "none", + backward_gw_mode: TransformMode = "none", + ): + return super().__new__(cls, tensor, config) + + def __init__( + self, + tensor, + config, + *, + hadamard_transform: Optional[Any] = None, + forward_y_mode: TransformMode = "none", + backward_gx_mode: TransformMode = "none", + backward_gw_mode: TransformMode = "none", + ): + super().__init__(tensor, config) + assert_mode_supported(forward_y_mode, "forward_y") + assert_mode_supported(backward_gx_mode, "backward_gx") + assert_mode_supported(backward_gw_mode, "backward_gw") + self._hadamard_transform = hadamard_transform + self._forward_y_mode = forward_y_mode + self._backward_gx_mode = backward_gx_mode + self._backward_gw_mode = backward_gw_mode + + def set_modes( + self, + forward_y_mode: TransformMode, + backward_gx_mode: TransformMode, + backward_gw_mode: TransformMode, + hadamard_transform: Optional[Any] = None, + ) -> None: + """Set per-slot modes (and optionally the Hadamard transform) IN PLACE. + + Used at the Phase-A -> Phase-B transition, so no new ``nn.Parameter`` is + created and the modes take effect on the FSDP-owned tensor. + """ + assert_mode_supported(forward_y_mode, "forward_y") + assert_mode_supported(backward_gx_mode, "backward_gx") + assert_mode_supported(backward_gw_mode, "backward_gw") + self._forward_y_mode = forward_y_mode + self._backward_gx_mode = backward_gx_mode + self._backward_gw_mode = backward_gw_mode + if hadamard_transform is not None: + self._hadamard_transform = hadamard_transform + + @property + def is_calibrating(self) -> bool: + """True while all slots are ``"none"`` (== plain MXFP4 / Phase A).""" + return ( + self._forward_y_mode == "none" + and self._backward_gx_mode == "none" + and self._backward_gw_mode == "none" + ) + + def _adahop_kwargs(self) -> dict: + return { + "hadamard_transform": self._hadamard_transform, + "forward_y_mode": self._forward_y_mode, + "backward_gx_mode": self._backward_gx_mode, + "backward_gw_mode": self._backward_gw_mode, + } + + def __repr__(self): + ht_state = "attached" if self._hadamard_transform is not None else "none" + return (f"MXFP4AdaHOPWrapper(shape={tuple(self.shape)}, dtype={self.dtype}, " + f"forward_y={self._forward_y_mode}, backward_gx={self._backward_gx_mode}, " + f"backward_gw={self._backward_gw_mode}, hadamard={ht_state})") + + @classmethod + def __torch_function__(cls, func, types, args, kwargs={}): + if func.__name__ in gemm_ops: + x, weight, bias, trans_b = _extract_x_w_with_bias(func, args) + assert isinstance(weight, cls), f"weight should be a {cls.__name__} for {func.__name__}" + # During calibration (all slots "none") behave EXACTLY like plain + # MXFP4: defer to the parent dispatch. Avoids the AdaHOP linear + # function's "none" path (unused by the recipe post-Phase-B). + if weight.is_calibrating: + return super().__torch_function__(func, types, args, kwargs) + # Pass the WRAPPER tensor (grad-tracked) into apply(), NOT weight._data + # (which is requires_grad=False). Passing the raw payload severs the + # autograd link so the weight gradient is silently discarded and the + # Linear never updates. Transpose is subclass- and graph-preserving + # (see _ops_to_preserve_subclass in tensor.py), so weight.T keeps the + # link; the payload is unwrapped inside MXFP4AdaHOPLinearFunction.forward, + # past the autograd boundary. Mirrors the plain MXFP4 / NVFP4 paths. + operand_w = weight if trans_b else weight.T + y = MXFP4AdaHOPLinearFunction.apply( + x, + operand_w, + weight.config.use_sr_grad, + weight._hadamard_transform, + weight._forward_y_mode, + weight._backward_gx_mode, + weight._backward_gw_mode, + ) + if bias is not None: + y = y + bias + return y + # All other ops (grouped_mm, etc.) defer to the standard parent path. + return super().__torch_function__(func, types, args, kwargs) + + def __tensor_flatten__(self): + meta = { + "config": self.config, + "forward_y_mode": self._forward_y_mode, + "backward_gx_mode": self._backward_gx_mode, + "backward_gw_mode": self._backward_gw_mode, + } + return ["_data"], meta + + @classmethod + def __tensor_unflatten__(cls, inner_tensors, flatten_spec, outer_size, outer_stride): + instance = cls( + inner_tensors["_data"], + flatten_spec["config"], + hadamard_transform=None, + forward_y_mode=flatten_spec["forward_y_mode"], + backward_gx_mode=flatten_spec["backward_gx_mode"], + backward_gw_mode=flatten_spec["backward_gw_mode"], + ) + return instance + + def fsdp_post_all_gather( + self, + all_gather_outputs: Tuple[torch.Tensor, ...], + metadata: Any, + param_dtype: torch.dtype, + *, + out: Optional[torch.Tensor] = None, + ): + # Step 1+: `out` is pre-allocated. Preserve our modes onto it. + if out is not None: + if isinstance(out, MXFP4AdaHOPWrapper): + out._hadamard_transform = self._hadamard_transform + out._forward_y_mode = self._forward_y_mode + out._backward_gx_mode = self._backward_gx_mode + out._backward_gw_mode = self._backward_gw_mode + return super().fsdp_post_all_gather(all_gather_outputs, metadata, param_dtype, out=out) + # Step 0: parent creates a fresh wrapper without modes; re-create with modes. + (data,) = all_gather_outputs + output = type(self)(data, self.config, **self._adahop_kwargs()) + return output, (data,) + + +def _extract_x_w_with_bias(func, args): + """Return ``(x, weight, bias, trans_b)`` for a gemm-family op.""" + trans_b = func.__name__ == "linear" + if func.__name__ == "addmm.default": + bias, A, B = args[0], args[1], args[2] + else: + A, B = args[0], args[1] + bias = args[2] if len(args) > 2 else None + return A, B, bias, trans_b + + +# Allowlist the wrapper for DCP checkpoint load (PyTorch >= 2.6 defaults +# torch.load to weights_only=True, which rejects unknown tensor-subclass globals). +# Mirrors the base-wrapper registration in dispatch/tensor.py. +torch.serialization.add_safe_globals([MXFP4AdaHOPWrapper]) diff --git a/alto/kernels/dispatch/config.py b/alto/kernels/dispatch/config.py index 7dd4862c..f9a439b5 100644 --- a/alto/kernels/dispatch/config.py +++ b/alto/kernels/dispatch/config.py @@ -30,6 +30,15 @@ class TrainingOpConfig: use_sr_grad: bool use_dge: bool + full_precision_backward: bool = False + """ + Quantize only the forward GEMM; keep the backward (dgrad + wgrad) in bf16 with + the gradient left unquantized. The backward GEMMs reuse the quantize-then- + dequantize x/w from the forward (QDQ operands). Used to isolate the training + quality impact of forward-only vs. full low-precision. Currently supported only + on the dense MXFP4 linear path. + """ + clip_mode: Literal["none", "static", "dynamic"] = "none" """ clipping mode applied in MXFP4/NVFP4 quantization. @@ -48,5 +57,13 @@ class TrainingOpConfig: * NVFP4: not implemented """ + blockscale_selection: Literal["default", "midmax", "uos"] = "default" + """ + block scale-selection strategy for MXFP4 quantization. + * default: round-even (adaptive) exponent selection (threshold > 7.0) + * midmax: alternative implementation of default with threshold >= 7.0 + * uos: not yet implemented + """ + torch.serialization.add_safe_globals([TrainingOpConfig]) diff --git a/alto/kernels/dispatch/conversion.py b/alto/kernels/dispatch/conversion.py index 325c9e35..cf465d82 100644 --- a/alto/kernels/dispatch/conversion.py +++ b/alto/kernels/dispatch/conversion.py @@ -8,7 +8,7 @@ # # SPDX-License-Identifier: BSD-3-Clause AND MIT -from typing import Callable, Optional, Type +from typing import Any, Callable, Dict, Optional, Type import torch from torch import nn @@ -46,6 +46,8 @@ def swap_params( config: Optional[TrainingOpConfig] = None, target_parameter_name: Optional[str] = None, module_name: Optional[str] = None, + tensor_cls: Optional[Type[torch.Tensor]] = None, + tensor_cls_kwargs: Optional[Dict[str, Any]] = None, ) -> nn.Module: """ Recurses through the nn.Module, recursively swapping the data tensor of @@ -68,13 +70,15 @@ def swap_params( if config is None: raise ValueError("training op config is required") - tensor_cls = _get_tensor_cls_for_config(config) + if tensor_cls is None: + tensor_cls = _get_tensor_cls_for_config(config) + extra_kwargs: Dict[str, Any] = dict(tensor_cls_kwargs or {}) if isinstance(module, nn.Parameter) and (module_filter_fn is None or module_filter_fn(module, "")): if len(list(module.children())) > 0: raise AssertionError(f"Does not support a root nn.Parameter with children: {module}") if not isinstance(module.data, TrainingWeightWrapperBaseTensor): - new_data = tensor_cls(module.data, config) + new_data = tensor_cls(module.data, config, **extra_kwargs) return nn.Parameter(new_data, requires_grad=module.requires_grad) return module @@ -106,7 +110,7 @@ def post_order_traversal( continue if not isinstance(param.data, TrainingWeightWrapperBaseTensor): new_param = nn.Parameter( - tensor_cls(param.data, config), + tensor_cls(param.data, config, **extra_kwargs), requires_grad=param.requires_grad, ) setattr(module, param_name, new_param) diff --git a/alto/kernels/dispatch/tensor.py b/alto/kernels/dispatch/tensor.py index 40868afd..a0c02ff9 100644 --- a/alto/kernels/dispatch/tensor.py +++ b/alto/kernels/dispatch/tensor.py @@ -16,6 +16,7 @@ from torchtitan.tools.logging import logger from alto.kernels.fp4.mxfp4.mxfp_linear import _to_mxfp4_then_scaled_mm +from alto.kernels.fp4.mxfp4.mxfp4_forward_only import _mxfp4_forward_only from alto.kernels.fp4.mxfp4.mxfp_grouped_gemm.functional import _quantize_then_mxfp_scaled_grouped_mm from alto.kernels.fp4.nvfp4.nvfp_linear import _to_nvfp4_then_scaled_mm from alto.kernels.fp4.nvfp4.nvfp_grouped_gemm.functional import ( @@ -215,6 +216,10 @@ class MXFP4TrainingWeightWrapperTensor(TrainingWeightWrapperBaseTensor): based on the training config. """ + # Set by GradientClippingModifier.on_initialize to wire this tensor back to + # its owning module so the backward pass can look up clipping config. + module_id: int | None = None + @classmethod def __torch_function__(cls, func, types, args, kwargs={}): # grouped_mm op override @@ -241,6 +246,10 @@ def __torch_function__(cls, func, types, args, kwargs={}): assert A_is_2d and B_is_3d and offs is not None, "Only 2d x 3d with offsets is supported for now" assert bias is None, "Bias is not supported for now" assert config.precision == "mxfp4", ("expected TrainingOpConfig with precision=mxfp4") + if config.full_precision_backward: + raise NotImplementedError( + "full_precision_backward is not implemented for the MoE grouped-mm path yet " + "(dense Linear only)") # logger.info( # f"[MXFP4GroupedMM]config: {config} A.shape: {A.shape} B.shape: {B.shape} offs.shape: {offs.shape}") @@ -256,6 +265,7 @@ def __torch_function__(cls, func, types, args, kwargs={}): use_hadamard=config.use_hadamard, clip_mode=config.clip_mode, use_macro_block_scaling=config.two_level_scaling == "blockwise", + blockscale_selection=config.blockscale_selection, ) # linear op override @@ -277,16 +287,28 @@ def __torch_function__(cls, func, types, args, kwargs={}): # logger.info(f"[MXFP4Linear]func: {func.__name__} config: {config}" # f"A.shape: {A.shape} B.shape: {B.shape} bias.shape: {bias.shape if bias is not None else None}") - Y = _to_mxfp4_then_scaled_mm( - A, - B if trans_b else B.T, - use_2dblock_x=config.use_2dblock_x, - use_2dblock_w=config.use_2dblock_w, - use_sr_grad=config.use_sr_grad, - use_dge=config.use_dge, - clip_mode=config.clip_mode, - use_hadamard=config.use_hadamard, - use_macro_block_scaling=config.two_level_scaling == "blockwise", + module_id = getattr(B, "module_id", None) + if config.full_precision_backward: + Y = _mxfp4_forward_only( + A, + B if trans_b else B.T, + use_2dblock_x=config.use_2dblock_x, + use_2dblock_w=config.use_2dblock_w, + use_macro_block_scaling=config.two_level_scaling == "blockwise", + ) + else: + Y = _to_mxfp4_then_scaled_mm( + A, + B if trans_b else B.T, + use_2dblock_x=config.use_2dblock_x, + use_2dblock_w=config.use_2dblock_w, + use_sr_grad=config.use_sr_grad, + use_dge=config.use_dge, + clip_mode=config.clip_mode, + use_hadamard=config.use_hadamard, + use_macro_block_scaling=config.two_level_scaling == "blockwise", + module_id=module_id, + blockscale_selection=config.blockscale_selection, ) if bias is not None: Y = Y + bias diff --git a/alto/kernels/fp4/fp4_common/grad_clip_config.py b/alto/kernels/fp4/fp4_common/grad_clip_config.py new file mode 100644 index 00000000..19bbf64c --- /dev/null +++ b/alto/kernels/fp4/fp4_common/grad_clip_config.py @@ -0,0 +1,16 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT + +from dataclasses import dataclass + + +@dataclass +class GradClipConfig: + clip_grad_output: bool = False + grad_output_max_norm: float | None = None + grad_output_clip_value: float | None = None + + clip_grad_weight: bool = False + grad_weight_max_norm: float | None = None + grad_weight_clip_value: float | None = None diff --git a/alto/kernels/fp4/fp4_common/grad_clip_registry.py b/alto/kernels/fp4/fp4_common/grad_clip_registry.py new file mode 100644 index 00000000..b5cd69e0 --- /dev/null +++ b/alto/kernels/fp4/fp4_common/grad_clip_registry.py @@ -0,0 +1,40 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT + +import torch + +from .grad_clip_config import GradClipConfig + +# Thread-local dict: id(module) -> GradClipConfig +_registry: dict[int, GradClipConfig] = {} + + +def register(module_id: int, cfg: GradClipConfig) -> None: + _registry[module_id] = cfg + + +def deregister(module_id: int) -> None: + _registry.pop(module_id, None) + + +def get(module_id: int | None) -> GradClipConfig | None: + if module_id is None: + return None + return _registry.get(module_id) + + +def apply_clip(t: torch.Tensor, max_norm: float | None, clip_value: float | None) -> torch.Tensor: + """Apply L2-norm clipping first (if set), then element-wise value clamp (if set). + + Returns t unchanged if both are None. + """ + if max_norm is None and clip_value is None: + return t + if max_norm is not None: + norm = t.norm() + if norm > max_norm: + t = t * (max_norm / norm) + if clip_value is not None: + t = t.clamp(-clip_value, clip_value) + return t diff --git a/alto/kernels/fp4/mxfp4/mxfp4_forward_only.py b/alto/kernels/fp4/mxfp4/mxfp4_forward_only.py new file mode 100644 index 00000000..1378a78d --- /dev/null +++ b/alto/kernels/fp4/mxfp4/mxfp4_forward_only.py @@ -0,0 +1,89 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT + +import torch + +from alto.kernels.fp4.fp4_common import unwrap_weight_wrapper +from .mxfp_quantization import is_cdna4 + + +class MXFP4ForwardOnlyLinearFunction(torch.autograd.Function): + """Quantize the forward GEMM only; keep the backward in bf16. + + The forward quantizes x and weight to MXFP4 and dequantizes them (QDQ), + exactly as the standard MXFP4 forward does, then runs ``y = x_dq @ w_dq.T``. + The backward is a plain bf16 matmul against the saved QDQ operands, with the + incoming gradient left unquantized: + + grad_x = grad_output @ w_dq + grad_w = grad_output.T @ x_dq + + This isolates the effect of forward-only quantization: nothing in the gradient + path is quantized. Unlike ``MXFP4LinearFunction`` this is an additive, separate + Function (the shared quantized-backward path is left untouched). + """ + + @staticmethod + def forward(ctx, x, weight, use_2dblock_x, use_2dblock_w): + assert not is_cdna4(), ( + "MXFP4ForwardOnlyLinearFunction only supports the non-CDNA4 (QDQ) path") + weight = unwrap_weight_wrapper(weight) + + original_shape = x.shape + original_dtype = x.dtype + x = x.reshape(-1, original_shape[-1]) + + x_mxfp4, x_scale = torch.ops.torchtitan.convert_to_mxfp4( + x, + axis=-1, + is_2d_block=use_2dblock_x, + ) + w_mxfp4, w_scale = torch.ops.torchtitan.convert_to_mxfp4( + weight, + axis=-1, + is_2d_block=use_2dblock_w, + ) + x_dq = torch.ops.torchtitan.convert_from_mxfp4( + x_mxfp4, + x_scale, + original_dtype, + axis=-1, + is_2d_block=use_2dblock_x, + ) + w_dq = torch.ops.torchtitan.convert_from_mxfp4( + w_mxfp4, + w_scale, + original_dtype, + axis=-1, + is_2d_block=use_2dblock_w, + ) + + y = x_dq @ w_dq.T + + ctx.save_for_backward(x_dq, w_dq) + ctx.original_shape = original_shape + return y.view(*original_shape[:-1], -1) + + @staticmethod + def backward(ctx, grad_output): + x_dq, w_dq = ctx.saved_tensors + original_shape = ctx.original_shape + + grad_output = grad_output.reshape(-1, grad_output.shape[-1]) + grad_inputs = grad_output @ w_dq + grad_weights = grad_output.T @ x_dq + + return grad_inputs.view(*original_shape[:-1], -1), grad_weights, None, None + + +def _mxfp4_forward_only( + a: torch.Tensor, + b: torch.Tensor, + use_2dblock_x: bool, + use_2dblock_w: bool, + use_macro_block_scaling: bool = False, +) -> torch.Tensor: + assert not use_macro_block_scaling, ( + "full_precision_backward does not support two_level_scaling / macro-block scaling") + return MXFP4ForwardOnlyLinearFunction.apply(a, b, use_2dblock_x, use_2dblock_w) diff --git a/alto/kernels/fp4/mxfp4/mxfp_grouped_gemm/cg_backward.py b/alto/kernels/fp4/mxfp4/mxfp_grouped_gemm/cg_backward.py index b215a122..a483a6ba 100644 --- a/alto/kernels/fp4/mxfp4/mxfp_grouped_gemm/cg_backward.py +++ b/alto/kernels/fp4/mxfp4/mxfp_grouped_gemm/cg_backward.py @@ -245,7 +245,13 @@ def _kernel_mxfp4_grouped_gemm_backward_dw( stride_asm, stride_ask, # Matrix dimensions - M_TOTAL, # Total M dimension + # constexpr: M_TOTAL is passed as a Python int (expert_indices.shape[0]) and + # is used to derive PACKED_M/Ms, which are declared tl.constexpr below. On the + # CDNA4 (gfx950) Triton backend a non-constexpr M_TOTAL stays a runtime tensor, + # so PACKED_M/Ms become int32[] tensors and tl.minimum() fails to compile + # ("cannot convert int32[] to tensor"). CDNA3 tolerated it; making it constexpr + # fixes CDNA4 and is a no-op on CDNA3 (value is already compile-time known). + M_TOTAL: tl.constexpr, # Total M dimension N: tl.constexpr, # N dimension K: tl.constexpr, # K dimension # Number of experts @@ -614,6 +620,7 @@ def forward( use_sr_grad=False, use_dge=False, clip_mode=False, + blockscale_selection="default", use_macro_block_scaling=False, hadamard_transform: Optional[HadamardTransform] = None, ): @@ -627,25 +634,36 @@ def forward( requant_axis_w = -1 expert_weights = unwrap_weight_wrapper(expert_weights) + + # Forward GEMM rotation: Q(XH) @ Q(WH)^T + if hadamard_transform is not None and not use_2dblock_x and not use_2dblock_w: + inputs_fwd = hadamard_transform(inputs) + expert_weights_fwd = hadamard_transform(expert_weights) + else: + inputs_fwd = inputs + expert_weights_fwd = expert_weights + if use_macro_block_scaling: - inputs_scaled, input_mbs = macro_block_scaling(inputs, axis=-1, use_2d_block=use_2dblock_x) - expert_weights_scaled, expert_weight_mbs = macro_block_scaling(expert_weights, + inputs_scaled, input_mbs = macro_block_scaling(inputs_fwd, axis=-1, use_2d_block=use_2dblock_x) + expert_weights_scaled, expert_weight_mbs = macro_block_scaling(expert_weights_fwd, axis=quant_axis_w, use_2d_block=use_2dblock_w) else: - inputs_scaled = inputs - input_mbs = inputs.new_empty([]) - expert_weights_scaled = expert_weights - expert_weight_mbs = expert_weights.new_empty([]) + inputs_scaled = inputs_fwd + input_mbs = inputs_fwd.new_empty([]) + expert_weights_scaled = expert_weights_fwd + expert_weight_mbs = expert_weights_fwd.new_empty([]) inputs_mxfp4, input_scales = torch.ops.torchtitan.convert_to_mxfp4( inputs_scaled, axis=-1, is_2d_block=use_2dblock_x, + blockscale_selection=blockscale_selection, ) expert_weights_mxfp4, expert_weight_scales = torch.ops.torchtitan.convert_to_mxfp4( expert_weights_scaled, axis=quant_axis_w, is_2d_block=use_2dblock_w, + blockscale_selection=blockscale_selection, ) if is_cdna4(): @@ -686,17 +704,20 @@ def forward( res = cg_grouped_gemm_forward(x_dq, w_dq, expert_indices) if not use_2dblock_w: + # dgrad weight rotation: Q(HW) for dgrad path + w_for_dgrad = hadamard_transform(expert_weights, left_mul=True) if hadamard_transform is not None else expert_weights if use_macro_block_scaling: - expert_weights_scaled, expert_weight_mbs = macro_block_scaling(expert_weights, + expert_weights_scaled, expert_weight_mbs = macro_block_scaling(w_for_dgrad, axis=requant_axis_w, use_2d_block=False) else: - expert_weights_scaled = expert_weights - expert_weight_mbs = expert_weights.new_empty([]) + expert_weights_scaled = w_for_dgrad + expert_weight_mbs = w_for_dgrad.new_empty([]) expert_weights_mxfp4, expert_weight_scales = torch.ops.torchtitan.convert_to_mxfp4( expert_weights_scaled, axis=requant_axis_w, is_2d_block=False, + blockscale_selection=blockscale_selection, ) if not is_cdna4(): w_dq = torch.ops.torchtitan.convert_from_mxfp4( @@ -724,6 +745,7 @@ def forward( axis=0, is_2d_block=False, clip_mode=clip_mode, + blockscale_selection=blockscale_selection, ) if not is_cdna4(): x_dq = torch.ops.torchtitan.convert_from_mxfp4( @@ -751,6 +773,7 @@ def forward( ctx.use_dge = use_dge ctx.hadamard_transform = hadamard_transform ctx.clip_mode = clip_mode + ctx.blockscale_selection = blockscale_selection ctx.use_macro_block_scaling = use_macro_block_scaling return res @@ -785,6 +808,7 @@ def backward(ctx, grad_output): axis=-1, use_sr=ctx.use_sr_grad, is_2d_block=True, + blockscale_selection=ctx.blockscale_selection, ) grad_output_mxfp4_m = grad_output_mxfp4 grad_output_scales_m = grad_output_scales @@ -802,17 +826,24 @@ def backward(ctx, grad_output): #grad_output_dq = grad_output_dq.contiguous() grad_output_m_dq = grad_output_dq else: + # dgrad grad_output rotation: Q(GH^T) for dgrad path (1D quant only) + if ctx.hadamard_transform is not None and not ctx.use_2dblock_w: + grad_output_dgrad = ctx.hadamard_transform(grad_output, inverse=True) + else: + grad_output_dgrad = grad_output if ctx.use_macro_block_scaling: - grad_output_scaled, grad_output_mbs = macro_block_scaling(grad_output, axis=-1, use_2d_block=False) + grad_output_scaled, grad_output_mbs = macro_block_scaling(grad_output_dgrad, axis=-1, use_2d_block=False) else: - grad_output_scaled = grad_output - grad_output_mbs = grad_output.new_empty([]) + grad_output_scaled = grad_output_dgrad + grad_output_mbs = grad_output_dgrad.new_empty([]) grad_output_mxfp4, grad_output_scales = torch.ops.torchtitan.convert_to_mxfp4( grad_output_scaled, axis=-1, use_sr=ctx.use_sr_grad, is_2d_block=False, + blockscale_selection=ctx.blockscale_selection, ) + # wgrad grad_output rotation: Q(HG) for wgrad path (existing) if ctx.hadamard_transform is not None: grad_output = ctx.hadamard_transform(grad_output, left_mul=True) if ctx.use_macro_block_scaling: @@ -826,6 +857,7 @@ def backward(ctx, grad_output): use_sr=ctx.use_sr_grad, is_2d_block=False, clip_mode=ctx.clip_mode, + blockscale_selection=ctx.blockscale_selection, ) if not is_cdna4(): @@ -918,7 +950,7 @@ def backward(ctx, grad_output): ) grad_weights *= dge_bwd(w_fp4_values, torch.float4_e2m1fn_x2) - return grad_inputs, grad_weights, None, None, None, None, None, None, None, None, None + return grad_inputs, grad_weights, None, None, None, None, None, None, None, None, None, None def mxfp4_grouped_gemm( @@ -933,6 +965,7 @@ def mxfp4_grouped_gemm( use_dge: bool = False, use_hadamard: bool = False, clip_mode: str = "none", + blockscale_selection: str = "default", use_macro_block_scaling: bool = False, ) -> torch.Tensor: """ @@ -973,6 +1006,7 @@ def mxfp4_grouped_gemm( use_sr_grad, use_dge, clip_mode, + blockscale_selection, use_macro_block_scaling, hadamard_transform, ) diff --git a/alto/kernels/fp4/mxfp4/mxfp_grouped_gemm/functional.py b/alto/kernels/fp4/mxfp4/mxfp_grouped_gemm/functional.py index 9dee3cf6..46f5bd7c 100644 --- a/alto/kernels/fp4/mxfp4/mxfp_grouped_gemm/functional.py +++ b/alto/kernels/fp4/mxfp4/mxfp_grouped_gemm/functional.py @@ -18,7 +18,8 @@ def _quantize_then_mxfp_scaled_grouped_mm( use_dge: bool, use_hadamard: bool, clip_mode: str, - use_macro_block_scaling: bool, + blockscale_selection: str = "default", + use_macro_block_scaling: bool = False, ) -> torch.Tensor: m_indices = create_indices_from_offsets_nosync(offs) return mxfp4_grouped_gemm( @@ -32,5 +33,6 @@ def _quantize_then_mxfp_scaled_grouped_mm( use_dge=use_dge, use_hadamard=use_hadamard, clip_mode=clip_mode, + blockscale_selection=blockscale_selection, use_macro_block_scaling=use_macro_block_scaling, ) diff --git a/alto/kernels/fp4/mxfp4/mxfp_linear.py b/alto/kernels/fp4/mxfp4/mxfp_linear.py index 4abe7e12..2158b954 100644 --- a/alto/kernels/fp4/mxfp4/mxfp_linear.py +++ b/alto/kernels/fp4/mxfp4/mxfp_linear.py @@ -19,6 +19,7 @@ import triton.language as tl from alto.kernels.fp4.fp4_common import unwrap_weight_wrapper +from alto.kernels.fp4.fp4_common.grad_clip_registry import apply_clip, get as get_grad_clip_cfg from alto.kernels.hadamard_transform import (HadamardTransform, HadamardFactory) from alto.kernels.dge import dge_bwd from .mxfp_quantization import ( @@ -272,8 +273,10 @@ def forward( use_sr_grad, use_dge, clip_mode, + blockscale_selection, use_macro_block_scaling, hadamard_transform: Optional[HadamardTransform] = None, + module_id: Optional[int] = None, ): """ Forward pass for the blockwise FP8 linear operation. @@ -297,25 +300,35 @@ def forward( original_dtype = x.dtype x = x.reshape(-1, original_shape[-1]) # Ensure x is 2D + # Forward GEMM rotation: Q(XH) @ Q(WH)^T + if hadamard_transform is not None and not use_2dblock_x and not use_2dblock_w: + x_fwd = hadamard_transform(x) + w_fwd = hadamard_transform(weight) + else: + x_fwd = x + w_fwd = weight + if use_macro_block_scaling: - x_scaled, x_mbs = macro_block_scaling(x, axis=-1, use_2d_block=use_2dblock_x) - w_scaled, w_mbs = macro_block_scaling(weight, axis=-1, use_2d_block=use_2dblock_w) + x_scaled, x_mbs = macro_block_scaling(x_fwd, axis=-1, use_2d_block=use_2dblock_x) + w_scaled, w_mbs = macro_block_scaling(w_fwd, axis=-1, use_2d_block=use_2dblock_w) else: - x_scaled = x - x_mbs = x.new_empty([]) - w_scaled = weight - w_mbs = weight.new_empty([]) + x_scaled = x_fwd + x_mbs = x_fwd.new_empty([]) + w_scaled = w_fwd + w_mbs = w_fwd.new_empty([]) x_mxfp4, x_scale = torch.ops.torchtitan.convert_to_mxfp4( x_scaled, axis=-1, is_2d_block=use_2dblock_x, + blockscale_selection=blockscale_selection, ) w_mxfp4, w_scale = torch.ops.torchtitan.convert_to_mxfp4( w_scaled, axis=-1, is_2d_block=use_2dblock_w, + blockscale_selection=blockscale_selection, ) if is_cdna4(): @@ -353,16 +366,19 @@ def forward( y = x_dq @ w_dq.T if not use_2dblock_w: + # dgrad weight rotation: Q(HW) for dgrad path + w_for_dgrad = hadamard_transform(weight, left_mul=True) if hadamard_transform is not None else weight if use_macro_block_scaling: - w_scaled, w_mbs = macro_block_scaling(weight, axis=0, use_2d_block=False) + w_scaled, w_mbs = macro_block_scaling(w_for_dgrad, axis=0, use_2d_block=False) else: - w_scaled = weight - w_mbs = weight.new_empty([]) + w_scaled = w_for_dgrad + w_mbs = w_for_dgrad.new_empty([]) w_mxfp4, w_scale = torch.ops.torchtitan.convert_to_mxfp4( w_scaled, axis=0, is_2d_block=False, + blockscale_selection=blockscale_selection, ) if not is_cdna4(): w_dq = torch.ops.torchtitan.convert_from_mxfp4( @@ -388,6 +404,7 @@ def forward( axis=0, is_2d_block=False, clip_mode=clip_mode, + blockscale_selection=blockscale_selection, ) if not is_cdna4(): x_dq = torch.ops.torchtitan.convert_from_mxfp4( @@ -411,7 +428,9 @@ def forward( ctx.hadamard_transform = hadamard_transform ctx.use_dge = use_dge ctx.clip_mode = clip_mode + ctx.blockscale_selection = blockscale_selection ctx.use_macro_block_scaling = use_macro_block_scaling + ctx.module_id = module_id return y.view(*original_shape[:-1], -1) # Reshape back to original @@ -419,8 +438,16 @@ def forward( def backward(ctx, grad_output): original_shape = grad_output.shape grad_output = grad_output.reshape(-1, original_shape[-1]) # Ensure grad_output is 2D + # [A/B #1] Match alto_rad baseline: dequantize grad_output to the incoming + # gradient's own dtype rather than the forward-committed ctx.original_dtype. original_dtype = grad_output.dtype + # Site ①: clip grad_output before it enters the quantizer. + _clip_cfg = get_grad_clip_cfg(ctx.module_id) + if _clip_cfg is not None and _clip_cfg.clip_grad_output: + grad_output = apply_clip(grad_output, _clip_cfg.grad_output_max_norm, + _clip_cfg.grad_output_clip_value) + if is_cdna4(): inputs_mxfp4, input_scales, weight_mxfp4, weight_scales, x_mbs, w_mbs = ctx.saved_tensors else: @@ -440,6 +467,7 @@ def backward(ctx, grad_output): axis=-1, is_2d_block=True, use_sr=ctx.use_sr_grad, + blockscale_selection=ctx.blockscale_selection, ) grad_output_mxfp4_m = grad_output_mxfp4 grad_output_scales_m = grad_output_scales @@ -457,18 +485,25 @@ def backward(ctx, grad_output): grad_output_dq = macro_block_descaling(grad_output_dq, grad_output_mbs, axis=-1, use_2d_block=True) grad_output_m_dq = grad_output_dq else: + # dgrad grad_output rotation: Q(GH^T) for dgrad path (1D quant only) + if ctx.hadamard_transform is not None and not ctx.use_2dblock_w: + grad_output_dgrad = ctx.hadamard_transform(grad_output, inverse=True) + else: + grad_output_dgrad = grad_output if ctx.use_macro_block_scaling: - grad_output_scaled, grad_output_mbs = macro_block_scaling(grad_output, axis=-1, use_2d_block=False) + grad_output_scaled, grad_output_mbs = macro_block_scaling(grad_output_dgrad, axis=-1, use_2d_block=False) else: - grad_output_scaled = grad_output - grad_output_mbs = grad_output.new_empty([]) + grad_output_scaled = grad_output_dgrad + grad_output_mbs = grad_output_dgrad.new_empty([]) grad_output_mxfp4, grad_output_scales = torch.ops.torchtitan.convert_to_mxfp4( grad_output_scaled, axis=-1, use_sr=ctx.use_sr_grad, is_2d_block=False, + blockscale_selection=ctx.blockscale_selection, ) + # wgrad grad_output rotation: Q(HG) for wgrad path (existing) if ctx.hadamard_transform is not None: grad_output = ctx.hadamard_transform(grad_output, left_mul=True) if ctx.use_macro_block_scaling: @@ -482,6 +517,7 @@ def backward(ctx, grad_output): use_sr=ctx.use_sr_grad, is_2d_block=False, clip_mode=ctx.clip_mode, + blockscale_selection=ctx.blockscale_selection, ) if not is_cdna4(): @@ -558,7 +594,7 @@ def backward(ctx, grad_output): ) grad_weights *= dge_bwd(w_fp4_values, torch.float4_e2m1fn_x2) - return grad_inputs.view(*original_shape[:-1], -1), grad_weights, None, None, None, None, None, None, None + return grad_inputs.view(*original_shape[:-1], -1), grad_weights, None, None, None, None, None, None, None, None, None def _to_mxfp4_then_scaled_mm( @@ -571,6 +607,8 @@ def _to_mxfp4_then_scaled_mm( clip_mode: str, use_hadamard: bool, use_macro_block_scaling: bool = False, + blockscale_selection: str = "default", + module_id: Optional[int] = None, ) -> torch.Tensor: if use_hadamard: with torch.no_grad(): @@ -585,7 +623,9 @@ def _to_mxfp4_then_scaled_mm( use_sr_grad, use_dge, clip_mode, + blockscale_selection, use_macro_block_scaling, hadamard_transform, + module_id, ) return y diff --git a/alto/kernels/fp4/mxfp4/mxfp_quantization.py b/alto/kernels/fp4/mxfp4/mxfp_quantization.py index dde4e6ae..50cc0b95 100644 --- a/alto/kernels/fp4/mxfp4/mxfp_quantization.py +++ b/alto/kernels/fp4/mxfp4/mxfp_quantization.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: MIT from typing import Tuple, Optional +from functools import lru_cache import torch from torch.library import triton_op, wrap_triton import triton @@ -20,6 +21,7 @@ _quantize_e2m1 = make_quantize_e2m1() +@lru_cache(maxsize=1) def is_cdna4(): target = triton.runtime.driver.active.get_current_target() return target is not None and target.backend == "hip" and target.arch == "gfx950" @@ -27,12 +29,13 @@ def is_cdna4(): @triton.jit def _calculate_scales( - x, + x, # raw input dtype, not e2m1 yet. likely fp16 or bf16? BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, QUANT_BLOCK_SIZE: tl.constexpr, IS_2D_BLOCK: tl.constexpr = False, USE_DYNAMIC_CLIP: tl.constexpr = False, + SCALE_SELECTION: tl.constexpr = "default", ): if x.type.element_ty == tl.float32: hp_int_dtype = tl.int32 @@ -44,39 +47,92 @@ def _calculate_scales( hp_ebits = 8 mbits = 1 sbits = 1 - target_max_pow2 = 2 + target_max_pow2 = 2 # maximum exponent value ( exp^{target_max_pow2} ) + # Reduce each quantization block to its running stats. The 2D and 1D cases + # differ only in how many intra-block axes we reduce over; the downstream + # math is shared. NEW_BLOCK_N: tl.constexpr = BLOCK_N // QUANT_BLOCK_SIZE if IS_2D_BLOCK: NEW_BLOCK_M: tl.constexpr = BLOCK_M // QUANT_BLOCK_SIZE x = x.reshape(NEW_BLOCK_M, QUANT_BLOCK_SIZE, NEW_BLOCK_N, QUANT_BLOCK_SIZE) + block_numel = QUANT_BLOCK_SIZE * QUANT_BLOCK_SIZE if USE_DYNAMIC_CLIP: - mean_squared = tl.sum(tl.sum(x * x, axis=-1), axis=-2) / (QUANT_BLOCK_SIZE * QUANT_BLOCK_SIZE) - mean = tl.sum(tl.sum(x, axis=-1), axis=-2) / (QUANT_BLOCK_SIZE * QUANT_BLOCK_SIZE) - std = tl.sqrt(mean_squared - mean * mean) - max_abs = (2.92247856 / 6.0) * std + 1e-8 - target_max_pow2 = 0 + sum_x = tl.sum(tl.sum(x, axis=-1), axis=-2) + sum_sq = tl.sum(tl.sum(x * x, axis=-1), axis=-2) else: - max_abs = tl.max(tl.abs(x), axis=-1) - max_abs = tl.max(max_abs, axis=-2) + max_abs = tl.max(tl.max(tl.abs(x), axis=-1), axis=-2) else: x = x.reshape(BLOCK_M, NEW_BLOCK_N, QUANT_BLOCK_SIZE) + block_numel = QUANT_BLOCK_SIZE if USE_DYNAMIC_CLIP: - mean_squared = tl.sum(x * x, axis=-1) / QUANT_BLOCK_SIZE - mean = tl.sum(x, axis=-1) / QUANT_BLOCK_SIZE - std = tl.sqrt(mean_squared - mean * mean) - max_abs = (2.92247856 / 6.0) * std + 1e-8 - target_max_pow2 = 0 + sum_x = tl.sum(x, axis=-1) + sum_sq = tl.sum(x * x, axis=-1) else: max_abs = tl.max(tl.abs(x), axis=-1) - max_abs = max_abs.to(x.type.element_ty) - # round even (adaptive) - max_abs = max_abs.to(hp_int_dtype, bitcast=True) - val_to_add = 1 << (hp_mbits - mbits - 1) - mask = ((1 << (hp_ebits + sbits)) - 1) << hp_mbits - max_abs = ((max_abs + val_to_add) & mask) >> hp_mbits - scales = max_abs - target_max_pow2 + if USE_DYNAMIC_CLIP: + # Estimate absmax from the block's std instead of its true max. + mean = sum_x / block_numel + std = tl.sqrt(sum_sq / block_numel - mean * mean) + max_abs = (2.92247856 / 6.0) * std + 1e-8 + target_max_pow2 = 0 + + # Each branch casts max_abs to the width it needs: the midmax path promotes + # to FP32 for exponent-field surgery, while round-even bitcasts in the input's + # native precision. Casting to element_ty up here would needlessly truncate + # the FP32 dynamic-clip estimate before midmax re-widens it. + if SCALE_SELECTION == "midmax-legacy": + # Pick the scale from absmax's FP32 exponent, then apply E2M1 "midmax" + # rounding: normalize absmax into [2^target_max_pow2, 2^(target_max_pow2+1)) + # by overwriting its exponent field, and bump the scale by 1 if the result + # exceeds midmax = 7.0 (the midpoint between E2M1's maxfloat 6.0 and 8.0). + FP32_MBITS = 23 + FP32_BIAS = 127 + FP32_MANT_MASK = 0x7FFFFF + FP32_EXP_MAX = 0xFF # exponent field of NaN/Inf + MIDMAX = 7.0 + + # collect exponent (in FP32) from max_abs + max_abs_bits = max_abs.to(tl.float32).to(tl.int32, bitcast=True) + f32_exp = (max_abs_bits >> FP32_MBITS) & FP32_EXP_MAX + # NaN/Inf have exponent 0xFF; cap to 0xFE so the scale stays bounded. + f32_exp = tl.where(f32_exp >= FP32_EXP_MAX, FP32_EXP_MAX - 1, f32_exp) + + scales = f32_exp - target_max_pow2 + amax_scaled_bits = (max_abs_bits & FP32_MANT_MASK) | ((FP32_BIAS + target_max_pow2) << FP32_MBITS) + amax_scaled = amax_scaled_bits.to(tl.float32, bitcast=True) + + scales = scales + (amax_scaled > MIDMAX).to(tl.int32) + elif SCALE_SELECTION == "uos6": + max_abs = max_abs.to(x.type.element_ty).to(hp_int_dtype, bitcast=True) + val_to_add = 1 << (hp_mbits - mbits) # 0x00400000 for fp32, 0x0040 for bf16 + mask = ((1 << (hp_ebits + sbits)) - 1) << hp_mbits + # apply carry on mantissa, collect only the carried exponent + max_abs = ((max_abs + val_to_add) & mask) >> hp_mbits + scales = max_abs - target_max_pow2 # e8m0 po2 exponent, so applying po2 arithmetic + elif SCALE_SELECTION == "uos": + # "default": round even (adaptive), incoming max_abs is in FP32 + max_abs = max_abs.to(x.type.element_ty).to(hp_int_dtype, bitcast=True) + val_to_add = 3 << (hp_mbits - mbits - 3) # 0x00180000 for fp32, 0x0018 for bf16 + mask = ((1 << (hp_ebits + sbits)) - 1) << hp_mbits + # apply carry on mantissa, collect only the carried exponent + max_abs = ((max_abs + val_to_add) & mask) >> hp_mbits + scales = max_abs - target_max_pow2 # e8m0 po2 exponent, so applying po2 arithmetic + else: + # This is CHECK7 + # "default": round even (adaptive), incoming max_abs is in FP32 + max_abs = max_abs.to(x.type.element_ty).to(hp_int_dtype, bitcast=True) + # with rounding you apply value_to_add on mantissa, + # i.e. 123.2 + 0.5 --> no carry to 124 + # value_to_add is 0.5 here but below, is actually 0.25 + # so anything <0.25 away from carry will be carried + # 7 in + val_to_add = 1 << (hp_mbits - mbits - 1) # 0x00200000 for fp32, 0x0020 for bf16 + mask = ((1 << (hp_ebits + sbits)) - 1) << hp_mbits + # apply carry on mantissa, collect only the carried exponent + max_abs = ((max_abs + val_to_add) & mask) >> hp_mbits + scales = max_abs - target_max_pow2 # e8m0 po2 exponent, so applying po2 arithmetic # Today, 2**-127 returns 0 in compile+inductor+triton because it is in the # float32 denormal range. For now, manually adjust the fp scale. This is @@ -269,6 +325,7 @@ def _convert_to_mxfp4_kernel( USE_ASM: tl.constexpr, USE_STATIC_CLIP: tl.constexpr, USE_DYNAMIC_CLIP: tl.constexpr, + SCALE_SELECTION: tl.constexpr, ): """ Quantizes the input tensor `x_ptr` and stores the result in `y_ptr` and the scaling factor in `s_ptr`. @@ -308,6 +365,7 @@ def _convert_to_mxfp4_kernel( QUANT_BLOCK_SIZE=QUANT_BLOCK_SIZE, IS_2D_BLOCK=IS_2D_BLOCK, USE_DYNAMIC_CLIP=USE_DYNAMIC_CLIP, + SCALE_SELECTION=SCALE_SELECTION, ) if USE_STATIC_CLIP: @@ -410,6 +468,7 @@ def convert_to_mxfp4( philox_seed: Optional[int] = None, philox_offset: Optional[int] = None, clip_mode: str = "none", + blockscale_selection: str = "default", ) -> Tuple[torch.Tensor, torch.Tensor]: torch._check(data_hp.shape[axis] % block_size == 0) assert not is_2d_block or data_hp.size(-2) % block_size == 0 @@ -469,6 +528,7 @@ def convert_to_mxfp4( USE_ASM=use_asm, USE_STATIC_CLIP=use_static_clip, USE_DYNAMIC_CLIP=use_dynamic_clip, + SCALE_SELECTION=blockscale_selection, ) return data_lp.reshape(new_shape).transpose(axis, -1), scales.reshape(scales_shape).transpose(axis, -1) diff --git a/alto/kernels/hadamard_transform/hadamard.py b/alto/kernels/hadamard_transform/hadamard.py index 8208e410..add1630e 100644 --- a/alto/kernels/hadamard_transform/hadamard.py +++ b/alto/kernels/hadamard_transform/hadamard.py @@ -22,6 +22,8 @@ REPO_PATH = Path(__file__).parent / "hadamards.safetensors" +_HADAMARD_CACHE: dict[tuple[int, torch.dtype, torch.device], torch.Tensor] = {} + __all__ = ["random_hadamard_matrix", "deterministic_hadamard_matrix", "is_pow2"] # note that hadamard matrix multiplication can be accelerated using a library such as @@ -108,21 +110,26 @@ def _fetch_hadamard_divisor( be of of size `k` such that `n / k` is a power of two. Return None if no such matrix exists. - Note: This function reopens the safetensors file every time it is called. - This is technically inefficient, but a very small runtime cost and simpler - than forcing callers to manage the file open context + Results are cached by (n, dtype, device) so the safetensors file is only + opened on the first call for each combination. :param n: size of known hadamard matrix :param dtype: data type to move fetched hadamard to :param device: device to move fetched hadamard to :return: a known hadamard matrix of size `n` if one exists, else None """ + cache_key = (n, dtype, device) + if cache_key in _HADAMARD_CACHE: + return _HADAMARD_CACHE[cache_key] + open_device = torch.device("cpu") if device.type == "meta" else device with safe_open(file_path, framework="pt", device=str(open_device)) as file: divisors = sorted((int(key) for key in file.keys()), reverse=True) for divisor in divisors: if n % divisor == 0 and is_pow2(n // divisor): - return file.get_tensor(str(divisor)).to(dtype=dtype, device=device) + result = file.get_tensor(str(divisor)).to(dtype=dtype, device=device) + _HADAMARD_CACHE[cache_key] = result + return result return None @@ -140,7 +147,7 @@ def _matmul_hadU(X: torch.Tensor) -> torch.Tensor: # Reshape diag matrix with randomized -1/+1 input = X.clone().view(-1, size, 1) - output = input.clone() + output = torch.empty_like(input) while input.shape[1] > K: input = input.view(input.shape[0], input.shape[1] // 2, 2, input.shape[2]) output = output.view(input.shape) diff --git a/alto/kernels/hadamard_transform/transform.py b/alto/kernels/hadamard_transform/transform.py index 490695de..608abb26 100644 --- a/alto/kernels/hadamard_transform/transform.py +++ b/alto/kernels/hadamard_transform/transform.py @@ -30,8 +30,10 @@ class HadamardFactory: block_size: int = 32 randomized: bool = True dtype: torch.dtype = torch.float32 + transform_type: str = "default" seed: Optional[int] = None generator: torch.Generator = torch.Generator() + _cached_transform: Optional['HadamardTransform'] = None @classmethod def configure( @@ -39,6 +41,7 @@ def configure( block_size: Optional[int] = None, randomized: Optional[bool] = None, dtype: Optional[torch.dtype] = None, + transform_type: Optional[str] = None, seed: Optional[int] = None, ) -> None: """ @@ -55,10 +58,17 @@ def configure( cls.randomized = randomized if dtype is not None: cls.dtype = dtype + if transform_type is not None: + cls.transform_type = transform_type if seed is not None: cls.seed = seed cls.generator.manual_seed(seed) + @classmethod + def refresh(cls) -> None: + """Clear the cached transform so the next create_transform generates a fresh one.""" + cls._cached_transform = None + @classmethod def create_transform( cls, @@ -73,10 +83,34 @@ def create_transform( :param device: Device to create the transform on :return: HadamardTransform instance """ + if cls._cached_transform is not None: + return cls._cached_transform + + if cls.transform_type == "default": + weight = cls._create_weight(device) + perm = cls._create_permutation(weight) if cls.randomized else None + t = HadamardTransform(weight, perm) + elif cls.transform_type == "3rht": + n = cls.block_size + w = cls._create_weight(device) + p = cls._create_permutation(w) + combined = w[p][:, p] + + w = cls._create_weight(device) + p = cls._create_permutation(w) + combined = combined @ (w[p][:, p]) + + w = cls._create_weight(device) + p = cls._create_permutation(w) + combined = combined @ (w[p][:, p]) + + combined = combined / n + t = HadamardTransform(combined, perm=None) + else: + raise NotImplementedError("transform_type options are: default and 3rht") - weight = cls._create_weight(device) - perm = cls._create_permutation(weight) if cls.randomized else None - return HadamardTransform(weight, perm) + cls._cached_transform = t + return t @classmethod def _create_weight( @@ -108,9 +142,13 @@ def __init__( weight: Tensor, perm: Optional[Tensor], ): - self.weight = weight - self.perm = perm - self._scale = torch.tensor(weight.size(0), dtype=torch.float64, device=weight.device).sqrt() + scale = weight.size(0) ** 0.5 + if perm is not None: + weight = weight[perm][:, perm] + if isinstance(weight, DTensor): + assert weight.placements[0] == Replicate() + weight = weight.to_local() + self.weight = (weight / scale).contiguous() def __call__(self, value: Tensor, inverse: bool = False, left_mul: bool = False) -> Tensor: """ @@ -124,19 +162,11 @@ def __call__(self, value: Tensor, inverse: bool = False, left_mul: bool = False) """ weight = self.weight - if self.perm is not None: - weight = weight[self.perm][:, self.perm] - if inverse: weight = weight.T - if isinstance(weight, DTensor): - assert weight.placements[0] == Replicate() - weight = weight.to_local() - + w = weight.to(device=value.device, dtype=value.dtype) if left_mul: - return (multihead_matmul(weight.to(device=value.device), value.to(dtype=weight.dtype)) / self._scale).to( - value.dtype) + return multihead_matmul(w, value) else: - return (multihead_matmul(value.to(dtype=weight.dtype), weight.to(device=value.device)) / self._scale).to( - value.dtype) + return multihead_matmul(value, w) diff --git a/alto/models/gpt_oss/config_registry.py b/alto/models/gpt_oss/config_registry.py index fe51292b..803b0b3d 100644 --- a/alto/models/gpt_oss/config_registry.py +++ b/alto/models/gpt_oss/config_registry.py @@ -10,13 +10,45 @@ ) from alto.components.converter import ModelOptConverter +from alto.components.m_adam import MAdamOptimizersContainer __all__ = [ "gpt_oss_debugmodel", "gpt_oss_debugmodel_lpt", + "gpt_oss_debugmodel_obs_lpt", + "gpt_oss_debugmodel_obs_bf16", "gpt_oss_20b", "gpt_oss_20b_pretrain", + "gpt_oss_20b_pretrain_c4", "gpt_oss_20b_lpt", + "gpt_oss_20b_lpt_fresh", + "gpt_oss_20b_lpt_1dw", + "gpt_oss_20b_lpt_no2dw", + "gpt_oss_20b_adahop", + "gpt_oss_20b_adahop_hadamard", + "gpt_oss_20b_pretrain_c4_megatron", + "gpt_oss_20b_lpt_c4", + "gpt_oss_20b_grad_clip_lpt", + "gpt_oss_debugmodel_grad_clip_lpt", + "gpt_oss_debugmodel_grad_clip_obs_lpt", + "gpt_oss_debugmodel_grad_clip_lpt_no_fsdp", + "gpt_oss_debugmodel_lpt_1gpu_ckpt", + "gpt_oss_debugmodel_obs_lpt_no_fsdp", + "gpt_oss_debugmodel_moe_pattern_obs", + "gpt_oss_debugmodel_moe_pattern_obs_no_fsdp", + "gpt_oss_20b_moe_pattern_obs", + "gpt_oss_20b_lpt_midmax", + "gpt_oss_20b_lpt_uos", + "gpt_oss_20b_lpt_uos6", + "gpt_oss_20b_lpt_lowrank", + "gpt_oss_20b_lpt_deosc", + "gpt_oss_20b_lpt_madam", + "gpt_oss_20b_lpt_madam_stable", + "gpt_oss_20b_mxfp4_base", + "gpt_oss_20b_mxfp4_had_2dw_sr", + "gpt_oss_20b_mxfp4_3rht_2dw_sr", + "gpt_oss_20b_mxfp4_3rht_2dw_sr_uos", + "gpt_oss_20b_mxfp4_3rht_2dw_sr_1d" ] @@ -40,9 +72,74 @@ def gpt_oss_debugmodel_lpt() -> Trainer.Config: return config +def gpt_oss_debugmodel_obs_lpt() -> Trainer.Config: + """gpt_oss debugmodel + MXFP4 + DebugObserver. Produces a per-step + tensor dump under the path configured in + ``configs/debug_observer_lpt_recipe.yaml``.""" + config = gpt_oss_debugmodel() + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config( + recipe="./alto/models/gpt_oss/configs/debug_observer_lpt_recipe.yaml", + ), + ],) + return config + + +def gpt_oss_debugmodel_lpt_1gpu_ckpt() -> Trainer.Config: + """Single-GPU (no FSDP/EP/TP) MXFP4 debugmodel training that writes full, + resumable checkpoints. Pairs with debug_train_gpt_oss_1gpu_ckpt.sh: run this + on a login-node GPU to produce checkpoints we later load for per-expert + stats. Model quality is irrelevant — we only need loadable checkpoints.""" + config = gpt_oss_debugmodel_lpt() + # Force all sharding degrees to 1 (login node can't do FSDP). + config.parallelism.data_parallel_shard_degree = 1 + config.parallelism.data_parallel_replicate_degree = 1 + config.parallelism.expert_parallel_degree = 1 + config.parallelism.expert_tensor_parallel_degree = 1 + config.parallelism.tensor_parallel_degree = 1 + config.compile.enable = False + # Full (resumable) checkpoints, keep them all so any step can be loaded later. + config.checkpoint.enable = True + config.checkpoint.interval = 50 + config.checkpoint.keep_latest_k = 0 + return config + + +def gpt_oss_debugmodel_obs_lpt_no_fsdp() -> Trainer.Config: + """Single-GPU (no FSDP/EP/TP) MXFP4 + DebugObserver. Loads a checkpoint + produced by gpt_oss_debugmodel_lpt_1gpu_ckpt, runs a few steps, and dumps + per-expert inputs/weights/grads to ./outputs/debug_obs_lpt.pt. Pairs with + debug_stats_gpt_oss_1gpu.sh, which sets --checkpoint.initial_load_path.""" + config = gpt_oss_debugmodel_obs_lpt() + config.parallelism.data_parallel_shard_degree = 1 + config.parallelism.data_parallel_replicate_degree = 1 + config.parallelism.expert_parallel_degree = 1 + config.parallelism.expert_tensor_parallel_degree = 1 + config.parallelism.tensor_parallel_degree = 1 + config.compile.enable = False + # Load pretrained weights only; the stats script points initial_load_path at + # a step-N checkpoint and clears any resumable state so the step counter + # starts at 0 and the observer actually fires. + config.checkpoint.enable = True + return config + + +def gpt_oss_debugmodel_obs_bf16() -> Trainer.Config: + """gpt_oss debugmodel in plain BF16 + DebugObserver (no LPT). Used to + capture the reference dump that visualizer can diff against the + quantized run.""" + config = gpt_oss_debugmodel() + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config( + recipe="./alto/models/gpt_oss/configs/debug_observer_bf16_recipe.yaml", + ), + ],) + return config + + def gpt_oss_20b() -> Trainer.Config: config = gpt_oss_20b_orig() - config.hf_assets_path = "/huggingface/hub/models--openai--gpt-oss-20b/snapshots/6cee5e81ee83917806bbde320786a8fb61efebee/" + config.hf_assets_path = "/hf_home/hub/models--openai--gpt-oss-20b/snapshots/6cee5e81ee83917806bbde320786a8fb61efebee/" config.dump_folder = "gpt_oss_20b-outputs" config.profiling.enable_profiling = False config.training.steps = 0 @@ -55,7 +152,7 @@ def gpt_oss_20b() -> Trainer.Config: config.parallelism.expert_tensor_parallel_degree = 1 config.parallelism.tensor_parallel_degree = 1 config.checkpoint.enable = True - config.checkpoint.initial_load_path = "/huggingface/hub/models--openai--gpt-oss-20b/snapshots/6cee5e81ee83917806bbde320786a8fb61efebee/" + config.checkpoint.initial_load_path = "/hf_home/hub/models--openai--gpt-oss-20b/snapshots/6cee5e81ee83917806bbde320786a8fb61efebee/" config.checkpoint.initial_load_in_hf = True config.checkpoint.initial_load_in_hf_quantized = True config.checkpoint.interval = 100 @@ -71,25 +168,26 @@ def gpt_oss_20b() -> Trainer.Config: def gpt_oss_20b_pretrain() -> Trainer.Config: config = gpt_oss_20b_orig() config.hf_assets_path = "/huggingface/hub/models--openai--gpt-oss-20b/snapshots/6cee5e81ee83917806bbde320786a8fb61efebee/" - config.dump_folder = "gpt_oss_20b-pretrain-subset-lr4e-4-outputs" + config.dump_folder = "gpt_oss_20b-mi300-pretrain-subset-lr4e-4-outputs" config.profiling.enable_profiling = False - config.training.steps = 1200000 + config.training.steps = 1200000 # set by mlperf config.training.local_batch_size = 1 - config.training.global_batch_size = 16 - config.training.seq_len = 8192 - config.optimizer.lr = 4e-4 - config.optimizer.weight_decay = 0.1 - config.optimizer.beta1 = 0.9 - config.optimizer.beta2 = 0.95 - config.optimizer.eps = 1e-5 - config.lr_scheduler.min_lr_factor = 0.1 - config.lr_scheduler.warmup_steps = 128 - config.lr_scheduler.decay_ratio = 1 - 128 / config.training.steps + config.training.global_batch_size = 16 # can be edited for mlperf submission + config.training.seq_len = 8192 # set by mlperf + config.optimizer.lr = 4e-4 # can be edited for mlperf + config.optimizer.weight_decay = 0.1 # set by mlperf + config.optimizer.beta1 = 0.9 # set by mlperf + config.optimizer.beta2 = 0.95 # set by mlperf + config.optimizer.eps = 1e-5 # set by mlperf + config.lr_scheduler.min_lr_factor = 0.1 # set by mlperf + config.lr_scheduler.warmup_steps = 128 # can be edited for mlperf submission + config.lr_scheduler.total_steps = 1200000 + config.lr_scheduler.decay_ratio = 1 - 128 / config.lr_scheduler.total_steps config.lr_scheduler.decay_type = "cosine" config.metrics.log_freq = 1 config.metrics.enable_tensorboard = True config.dataloader.dataset = "megatron" - config.dataloader.dataset_path = "/workspace/workspace/megatron_dataset/data/c4-train.en_6_text_document.idx" + config.dataloader.dataset_path = "/data/c4-train.en_6_text_document.idx" config.parallelism.expert_parallel_degree = 8 config.parallelism.expert_tensor_parallel_degree = 1 config.parallelism.tensor_parallel_degree = 1 @@ -97,18 +195,349 @@ def gpt_oss_20b_pretrain() -> Trainer.Config: config.checkpoint.interval = 1000 config.checkpoint.keep_latest_k = 2 config.validator.enable = True - config.validator.dataloader.dataset = "megatron" - config.validator.dataloader.dataset_path = "/workspace/workspace/megatron_dataset/data/c4-validation-91205-samples.en_text_document.idx" + config.validator.dataloader.dataset = "wikitext_test" + config.validator.dataloader.dataset_path = "" config.validator.freq = 768 config.validator.steps = 64 config.activation_checkpoint.mode = "none" config.debug.seed = 1234 return config +def gpt_oss_20b_pretrain_c4_megatron() -> Trainer.Config: + """gpt_oss_20b_pretrain using HuggingFace C4 dataset (bf16 baseline).""" + config = gpt_oss_20b_pretrain() + config.dump_folder = "gpt_oss_20b-pretrain-subset-bf16-c4-outputs" + config.dataloader.dataset = "megatron" + config.dataloader.dataset_path = "/data/c4-train.en_6_text_document.idx" + config.validator.dataloader.dataset = "megatron" + config.validator.dataloader.dataset_path = "/data/c4-validation-91205-samples.en_text_document.idx" + config.checkpoint.enable = True + config.checkpoint.initial_load_path = None # fresh run: do NOT load any checkpoint + config.checkpoint.initial_load_in_hf = False + config.checkpoint.initial_load_in_hf_quantized = False + config.checkpoint.interval = 1000 # Save at step interval + config.checkpoint.last_save_model_only = False # save full ckpt at final step (model+optim+dataloader) so training can resume + return config + +def gpt_oss_20b_adahop() -> Trainer.Config: + config = gpt_oss_20b_pretrain() + config.training.global_batch_size = 16 + config.parallelism.expert_tensor_parallel_degree = 1 + config.parallelism.tensor_parallel_degree = 1 + config.parallelism.expert_parallel_degree = 8 + config.training.local_batch_size = 1 + config.activation_checkpoint.mode = "none" + config.dataloader.dataset = "c4" + config.dataloader.dataset_path = None + config.validator.dataloader.dataset = "c4_validation" + config.validator.dataloader.dataset_path = None + config.checkpoint.enable = True # save checkpoints so we can resume later + config.checkpoint.initial_load_path = None # fresh run: do NOT load any checkpoint + config.checkpoint.initial_load_in_hf = False + config.checkpoint.initial_load_in_hf_quantized = False + config.checkpoint.interval = 500 # Save at step interval + config.checkpoint.keep_latest_k = 2 # keep only the 2 latest (each ~234G) + # Distinct from gpt_oss_20b_lpt's dump_folder so the adahop and nolora runs + # never share/overwrite each other's checkpoints. + config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4-adahop-outputs" + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/lpt_adahop.yaml",), + ],) + return config +def gpt_oss_20b_adahop_hadamard() -> Trainer.Config: + """Phase-2 Run A: AdaHOP with every slot forced to `hadamard` (no outlier + extraction, no full_precision). Identical to gpt_oss_20b_adahop except the + recipe's layer_transform_config maps all pattern-pairs to "hadamard". + Isolates whether the AdaHOP training regression comes from mode SELECTION + (S3) rather than the transform math. Distinct dump_folder so it never + collides with the calibrated adahop or the nolora runs.""" + config = gpt_oss_20b_adahop() + config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4-adahop-allhadamard-randomized-outputs" + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/lpt_adahop_all_hadamard.yaml",), + ],) + return config + def gpt_oss_20b_lpt() -> Trainer.Config: + config = gpt_oss_20b_pretrain_c4_megatron() + config.dump_folder = "gpt_oss_20b-mi300-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-rank32-lr4e-4-outputs" + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/lpt_recipe.yaml",), + ],) + return config + +def gpt_oss_20b_lpt_madam() -> Trainer.Config: + """Debug-model smoke test for the custom ``m_adam`` optimizer. + + ``lr`` maps to m_adam's additive (AdamW) branch ``lr_m`` and is driven + by the usual LR scheduler; ``lr_e`` is the multiplicative (exponent) + branch learning rate. + """ + config = gpt_oss_20b_lpt() + config.optimizer = MAdamOptimizersContainer.Config( + lr=1e-4, + lr_e=1e-4, + tie_e_to_m=True, + beta1=0.9, + beta2=0.95, + eps=1e-5, + weight_decay_m=0.1, + ) + return config + + +def gpt_oss_20b_lpt_madam_stable() -> Trainer.Config: + """``gpt_oss_lpt_madam`` with a damped exponent branch. + + Same baseline-matched additive branch as ``gpt_oss_lpt_madam``, but the + multiplicative (exponent) branch is turned down to suppress the loss + spikes seen when it runs as loud as the mantissa branch: + + * ``lr_e`` dropped to 0.3x ``lr_m`` so the additive branch leads and the + exponent only fine-tunes magnitudes, and + * ``de_step_cap`` tightened from 0.5 to 0.2 (max per-step magnitude change + ~2**0.2 ~= 1.15x instead of ~1.41x). + """ + config = gpt_oss_20b_lpt_madam() + config.optimizer.lr_e = 3e-5 + config.optimizer.de_step_cap = 0.2 + return config + + +def gpt_oss_20b_lpt_deosc() -> Trainer.Config: + """weight deoscillation config.""" + config = gpt_oss_20b_lpt() + config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-lr4e-4-deosc" + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/lpt_recipe_deosc.yaml",), + ],) + return config + +def gpt_oss_20b_lpt_midmax() -> Trainer.Config: + """gpt_oss_20b_lpt_c4 with midmax scale selection for MXFP4 quantization.""" + config = gpt_oss_20b_lpt() + config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-lr4e-4-midmax-outputs" + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/lpt_recipe_midmax.yaml",), + ],) + return config + +def gpt_oss_20b_lpt_uos() -> Trainer.Config: + """gpt_oss_20b_lpt_c4 with uos scale selection for MXFP4 quantization.""" + config = gpt_oss_20b_lpt() + config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-lr4e-4-uos" + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/lpt_recipe_uos.yaml",), + ],) + return config + +def gpt_oss_20b_lpt_uos6() -> Trainer.Config: + """gpt_oss_20b_lpt_c4 with uos6 scale selection for MXFP4 quantization.""" + config = gpt_oss_20b_lpt() + config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-lr4e-4-uos6" + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/lpt_recipe_uos6.yaml",), + ],) + return config + +def gpt_oss_20b_mxfp4_base() -> Trainer.Config: + """baseline MXFP4 quantization.""" + config = gpt_oss_20b_lpt() + config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-lr4e-4-mxfp4-base" + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/mxfp4_base.yaml",), + ],) + return config + +def gpt_oss_20b_mxfp4_had_2dw_sr() -> Trainer.Config: + """baseline MXFP4 quantization.""" + config = gpt_oss_20b_lpt() + config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-lr4e-4-mxfp4-base" + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/mxfp4_had_2dw_sr.yaml",), + ],) + return config + +def gpt_oss_20b_mxfp4_3rht_2dw_sr() -> Trainer.Config: + """baseline MXFP4 quantization.""" + config = gpt_oss_20b_lpt() + config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-lr4e-4-mxfp4-base" + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/mxfp4_3rht_2dw_sr.yaml",), + ],) + return config + +def gpt_oss_20b_mxfp4_3rht_2dw_sr_1d() -> Trainer.Config: + """baseline MXFP4 quantization.""" + config = gpt_oss_20b_lpt() + config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-lr4e-4-mxfp4-base" + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/mxfp4_3rht_2dw_sr_1d.yaml",), + ],) + return config + +def gpt_oss_20b_mxfp4_3rht_2dw_sr_uos() -> Trainer.Config: + """baseline MXFP4 quantization.""" + config = gpt_oss_20b_lpt() + config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-lr4e-4-mxfp4-base" + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/mxfp4_3rht_2dw_sr_uos.yaml",), + ],) + return config + +def gpt_oss_20b_lpt_lowrank() -> Trainer.Config: + """gpt_oss_20b_lpt_c4 with low-rank (lora_rank=32) correction for MXFP4 quantization.""" + config = gpt_oss_20b_lpt() + config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-lr4e-4-lowrank-outputs" + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/lpt_recipe_lowrank.yaml",), + ],) + return config + +def gpt_oss_20b_lpt_1dw() -> Trainer.Config: + """Plain MXFP4 with 1D-block weight quantization (use_2dblock_w: false): + weights are scaled per-axis (1D macro blocks) instead of the baseline's 2D + blocks. Identical to gpt_oss_20b_lpt otherwise. Own dump folder so it never + collides with the 2D-weight baseline's checkpoints.""" + config = gpt_oss_20b_lpt() + config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4-1dw-outputs" + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/lpt_recipe_1dw.yaml",), + ],) + return config + +def gpt_oss_20b_lpt_fresh() -> Trainer.Config: + """Plain MXFP4 baseline, fresh from step 1, own dump folder — the control for + the weight-identity experiment (2D weight scaling ON). Distinct dump_folder + from gpt_oss_20b_lpt so it never resumes an existing checkpoint.""" + config = gpt_oss_20b_lpt() + config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4-2dw-fresh-outputs" + return config + +def gpt_oss_20b_lpt_no2dw() -> Trainer.Config: + """Plain MXFP4 with 2D weight scaling DISABLED (use_2dblock_w: false). Tests + the AdaHOP root-cause hypothesis: breaking the baseline's single-2D-Q(W) + identity (W re-quantized per-axis instead) should degrade plain MXFP4 toward + AdaHOP's ~5.9 val@768. Fresh from step 1, own dump folder.""" + config = gpt_oss_20b_lpt() + config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4-no2dw-fresh-outputs" + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/lpt_recipe_no2dw.yaml",), + ],) + return config + +def gpt_oss_20b_grad_clip_lpt() -> Trainer.Config: + """20b pretrain + MXFP4 + gradient clipping at the quantizer boundary.""" config = gpt_oss_20b_pretrain() - config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-lr4e-4-outputs" + config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4-grad-clip-lr4e-4-outputs" + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/grad_clip_lpt_recipe.yaml",), + ],) + return config + + +def gpt_oss_debugmodel_grad_clip_lpt() -> Trainer.Config: + """Debugmodel + MXFP4 + gradient clipping. Use for config validation.""" + config = gpt_oss_debugmodel() + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/grad_clip_lpt_recipe.yaml",), + ],) + return config + + +def gpt_oss_debugmodel_grad_clip_lpt_no_fsdp() -> Trainer.Config: + """Debugmodel + MXFP4 + gradient clipping, FSDP disabled. Single GPU, + no weight sharding — used to isolate whether FSDP is preventing the clip + from firing.""" + config = gpt_oss_debugmodel_grad_clip_lpt() + config.parallelism.data_parallel_shard_degree = 1 + return config + + +def gpt_oss_debugmodel_grad_clip_obs_lpt() -> Trainer.Config: + """Debugmodel + MXFP4 + gradient clipping + DebugObserver. Produces tensor + dumps under outputs/debug_obs_grad_clip_lpt.pt for comparison against the + unclipped baseline.""" + config = gpt_oss_debugmodel() + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config( + recipe="./alto/models/gpt_oss/configs/grad_clip_obs_lpt_recipe.yaml", + ), + ],) + return config + + +def gpt_oss_debugmodel_moe_pattern_obs() -> Trainer.Config: + """Debugmodel + MXFP4 + per-expert MoE matmul outlier-pattern observer. + Accumulates patterns over several steps (max_captures in the recipe) and + dumps ./outputs/moe_patterns_rank*.pt with a per-expert majority vote. + Use with COMM_MODE=local_tensor for a single-GPU smoke; compile stays off.""" + config = gpt_oss_debugmodel() + config.training.steps = 10 + config.compile.enable = False + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config( + recipe="./alto/models/gpt_oss/configs/moe_pattern_observer_recipe.yaml", + ), + ],) + return config + + +def gpt_oss_debugmodel_moe_pattern_obs_no_fsdp() -> Trainer.Config: + """Single-GPU (no FSDP/EP/TP) debugmodel + MXFP4 + per-expert MoE matmul + outlier-pattern observer. Loads a checkpoint produced by + gpt_oss_debugmodel_lpt_1gpu_ckpt, runs ONE training step, and dumps + ./outputs/moe_patterns_rank*.pt. Pairs with debug_stats_gpt_oss_1gpu.sh, + which sets --checkpoint.initial_load_path to a step-N checkpoint.""" + config = gpt_oss_debugmodel_moe_pattern_obs() + config.parallelism.data_parallel_shard_degree = 1 + config.parallelism.data_parallel_replicate_degree = 1 + config.parallelism.expert_parallel_degree = 1 + config.parallelism.expert_tensor_parallel_degree = 1 + config.parallelism.tensor_parallel_degree = 1 + # Load pretrained weights; the stats script clears any resumable state so the + # step counter starts at 0 and the single observed step actually runs. + config.checkpoint.enable = True + return config + + +def gpt_oss_20b_moe_pattern_obs() -> Trainer.Config: + """gpt_oss_20b: load a checkpoint, run ONE training step, and dump the + per-expert MoE matmul outlier patterns. Intended flow: pretrain without the + observer, point checkpoint.initial_load_path at a checkpoint, run this.""" + config = gpt_oss_20b_lpt() + config.training.steps = 10 + config.compile.enable = False + config.validator.enable = False + config.checkpoint.enable = True # load the pretrained checkpoint + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config( + recipe="./alto/models/gpt_oss/configs/moe_pattern_observer_recipe.yaml", + ), + ],) + return config + +# legacy dataset ingestion using un-tokenized C4 +def gpt_oss_20b_pretrain_c4() -> Trainer.Config: + """gpt_oss_20b_pretrain using HuggingFace C4 dataset (bf16 baseline, no Megatron files required).""" + config = gpt_oss_20b_pretrain() + config.dump_folder = "gpt_oss_20b-pretrain-subset-bf16-c4-outputs" + config.training.global_batch_size = 64 + config.optimizer.lr = 4e-4 + config.lr_scheduler.min_lr_factor = 0.04 + config.dataloader.dataset = "c4" + config.dataloader.dataset_path = None + config.validator.dataloader.dataset = "c4_validation" + config.validator.dataloader.dataset_path = None + config.checkpoint.initial_load_in_hf = True + config.checkpoint.initial_load_in_hf_quantized = True + return config + +# legacy dataset ingestion using un-tokenized C4 +def gpt_oss_20b_lpt_c4() -> Trainer.Config: + config = gpt_oss_20b_pretrain_c4() + config.dump_folder = "gpt_oss_20b-mi300-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-rank32-c4-outputs" config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/lpt_recipe.yaml",), ],) diff --git a/alto/models/gpt_oss/configs/debug_observer_bf16_recipe.yaml b/alto/models/gpt_oss/configs/debug_observer_bf16_recipe.yaml new file mode 100644 index 00000000..0906bfa7 --- /dev/null +++ b/alto/models/gpt_oss/configs/debug_observer_bf16_recipe.yaml @@ -0,0 +1,12 @@ +training_stage: + debug_modifiers: + DebugObserverModifier: + targets: ["Linear", "GptOssGroupedExperts"] + ignore: ["output", "re:.*\\.router\\.gate"] + capture_every: 1 + max_captures: 10 + output_path: "./outputs/debug_obs_bf16.pt" + capture_input: true + capture_weight: true + capture_grad_output: true + capture_grad_weight: true diff --git a/alto/models/gpt_oss/configs/debug_observer_lpt_recipe.yaml b/alto/models/gpt_oss/configs/debug_observer_lpt_recipe.yaml new file mode 100644 index 00000000..236a363e --- /dev/null +++ b/alto/models/gpt_oss/configs/debug_observer_lpt_recipe.yaml @@ -0,0 +1,24 @@ +training_stage: + lpt_modifiers: + LowPrecisionTrainingModifier: + scheme: "mxfp4" + targets: ["Linear", "GptOssGroupedExperts"] + ignore: ["output", "re:.*\\.router\\.gate"] + use_2dblock_x: false + use_2dblock_w: true + use_hadamard: true + use_sr_grad: true + use_dge: false + clip_mode: none + two_level_scaling: none + debug_modifiers: + DebugObserverModifier: + targets: ["Linear", "GptOssGroupedExperts"] + ignore: ["output", "re:.*\\.router\\.gate"] + capture_every: 1 + max_captures: 10 + output_path: "./outputs/debug_obs_lpt.pt" + capture_input: true + capture_weight: true + capture_grad_output: true + capture_grad_weight: true diff --git a/alto/models/gpt_oss/configs/grad_clip_lpt_recipe.yaml b/alto/models/gpt_oss/configs/grad_clip_lpt_recipe.yaml new file mode 100644 index 00000000..e94674a0 --- /dev/null +++ b/alto/models/gpt_oss/configs/grad_clip_lpt_recipe.yaml @@ -0,0 +1,22 @@ +training_stage: + lpt_modifiers: + LowPrecisionTrainingModifier: + scheme: "mxfp4" + targets: ["Linear", "GptOssGroupedExperts"] + ignore: ["output", "re:.*\\.router\\.gate"] + use_2dblock_x: false + use_2dblock_w: true + use_hadamard: true + use_sr_grad: true + use_dge: false + clip_mode: none + two_level_scaling: none + GradientClippingModifier: + targets: ["Linear"] + ignore: ["output", "re:.*\\.router\\.gate"] + clip_grad_output: true + grad_output_max_norm: 1.0 + grad_output_clip_value: null + clip_grad_weight: true + grad_weight_max_norm: 1.0 + grad_weight_clip_value: null diff --git a/alto/models/gpt_oss/configs/grad_clip_obs_lpt_recipe.yaml b/alto/models/gpt_oss/configs/grad_clip_obs_lpt_recipe.yaml new file mode 100644 index 00000000..dc735d73 --- /dev/null +++ b/alto/models/gpt_oss/configs/grad_clip_obs_lpt_recipe.yaml @@ -0,0 +1,33 @@ +training_stage: + lpt_modifiers: + LowPrecisionTrainingModifier: + scheme: "mxfp4" + targets: ["Linear", "GptOssGroupedExperts"] + ignore: ["output", "re:.*\\.router\\.gate"] + use_2dblock_x: false + use_2dblock_w: true + use_hadamard: true + use_sr_grad: true + use_dge: false + clip_mode: none + two_level_scaling: none + GradientClippingModifier: + targets: ["Linear"] + ignore: ["output", "re:.*\\.router\\.gate"] + clip_grad_output: true + grad_output_max_norm: null + grad_output_clip_value: 0.0 + clip_grad_weight: true + grad_weight_max_norm: null + grad_weight_clip_value: 0.0 + debug_modifiers: + DebugObserverModifier: + targets: ["Linear", "GptOssGroupedExperts"] + ignore: ["output", "re:.*\\.router\\.gate"] + capture_every: 1 + max_captures: 10 + output_path: "./outputs/debug_obs_grad_clip_lpt.pt" + capture_input: true + capture_weight: true + capture_grad_output: true + capture_grad_weight: true diff --git a/alto/models/gpt_oss/configs/lpt_adahop.yaml b/alto/models/gpt_oss/configs/lpt_adahop.yaml new file mode 100644 index 00000000..ce21cb85 --- /dev/null +++ b/alto/models/gpt_oss/configs/lpt_adahop.yaml @@ -0,0 +1,50 @@ +training_stage: + lpt_modifiers: + LowPrecisionTrainingModifier: + scheme: "mxfp4_adahop" + targets: ["Linear"] + # Exclude the MoE router gate: quantizing it leaves top_scores in a dtype + # that mismatches routed_output at the bmm in moe.py (RuntimeError: expected + # BFloat16 but found Float). The nolora recipe already excludes it. AdaHOP + # only transforms LPT-wrapped Linears, so this exclusion also keeps the gate + # out of the AdaHOP pass — no separate ignore needed there. + ignore: ["output", "re:.*\\.router\\.gate"] + use_2dblock_x: false + use_2dblock_w: true + use_hadamard: true + use_sr_grad: true + use_dge: false + clip_mode: none + two_level_scaling: none + AdaHOPModifier: + enabled: true + use_hadamard: true + # Randomized Hadamard (random sign/permutation) to match the plain-MXFP4 + # baseline, which keeps HadamardFactory's randomized=True default. A fixed + # (deterministic) Hadamard decorrelates outliers less well; the small + # per-step difference compounds over training into a widening val-loss gap + # vs plain MXFP4 (S5). See plan keen-questing-ember. + use_randomized_hadamard: true + calibration_steps: 30 + layer_transform_config: + "row-row": "hadamard" + "row-none": "inner_outlier_extract_left" + "row-col": "inner_outlier_extract_right" + "col-row": "hadamard" + "col-none": "hadamard" + "col-col": "full_precision" + "none-row": "hadamard" + "none-none": "hadamard" + "none-col": "inner_outlier_extract_right" + moe_lpt_modifiers: + LowPrecisionTrainingModifier: + scheme: "mxfp4" + targets: ["GptOssGroupedExperts"] + ignore: ["output", "re:.*\\.router\\.gate"] + use_2dblock_x: false + use_2dblock_w: true + use_hadamard: true + use_sr_grad: true + use_dge: false + clip_mode: none + two_level_scaling: none \ No newline at end of file diff --git a/alto/models/gpt_oss/configs/lpt_recipe.yaml b/alto/models/gpt_oss/configs/lpt_recipe.yaml index ad935b53..fe70903c 100644 --- a/alto/models/gpt_oss/configs/lpt_recipe.yaml +++ b/alto/models/gpt_oss/configs/lpt_recipe.yaml @@ -4,7 +4,9 @@ training_stage: scheme: "mxfp4" targets: ["Linear", "GptOssGroupedExperts"] # targets: ["Linear"] + # ignore: ["output", "re:.*\\.router\\.gate", "re:layers\\.0.*", "re:layers\\.22.*", "re:layers\\.23.*"] ignore: ["output", "re:.*\\.router\\.gate"] + # ignore: ["output", "re:.*\\.router\\.gate", "re:layers\\.\\d+\\.attention.wq", "re:layers\\.\\d+\\.attention.wk", "re:layers\\.\\d+\\.attention.wv", "re:layers\\.\\d+\\.attention.wo"] use_2dblock_x: false use_2dblock_w: true use_hadamard: true diff --git a/alto/models/gpt_oss/configs/lpt_recipe_deosc.yaml b/alto/models/gpt_oss/configs/lpt_recipe_deosc.yaml new file mode 100644 index 00000000..56605f78 --- /dev/null +++ b/alto/models/gpt_oss/configs/lpt_recipe_deosc.yaml @@ -0,0 +1,28 @@ +training_stage: + lpt_modifiers: + LowPrecisionTrainingModifier: + scheme: "mxfp4" + targets: ["Linear", "GptOssGroupedExperts"] + # targets: ["Linear"] + # ignore: ["output", "re:.*\\.router\\.gate", "re:layers\\.0.*", "re:layers\\.22.*", "re:layers\\.23.*"] + ignore: ["output", "re:.*\\.router\\.gate"] + # ignore: ["output", "re:.*\\.router\\.gate", "re:layers\\.\\d+\\.attention.wq", "re:layers\\.\\d+\\.attention.wk", "re:layers\\.\\d+\\.attention.wv", "re:layers\\.\\d+\\.attention.wo"] + use_2dblock_x: false + use_2dblock_w: true + use_hadamard: true + use_sr_grad: true + + # differential gradient estimation, disabled by default + use_dge: false + + # choices: none, static, dynamic + clip_mode: none + + # 2-level scaling, disabled by default + # NVFP4 requires tensorwise scaling or you will get overflow issues + two_level_scaling: none + + # weight de-oscillation, disabled by default + deosc_step: 2000 + deosc_period: 200 + deosc_ratio: 4.0 diff --git a/alto/models/gpt_oss/configs/lpt_recipe_midmax.yaml b/alto/models/gpt_oss/configs/lpt_recipe_midmax.yaml new file mode 100644 index 00000000..5f564a9f --- /dev/null +++ b/alto/models/gpt_oss/configs/lpt_recipe_midmax.yaml @@ -0,0 +1,29 @@ +training_stage: + lpt_modifiers: + LowPrecisionTrainingModifier: + scheme: "mxfp4" + targets: ["Linear", "GptOssGroupedExperts"] + # targets: ["Linear"] + # ignore: ["output", "re:.*\\.router\\.gate", "re:layers\\.0.*", "re:layers\\.22.*", "re:layers\\.23.*"] + ignore: ["output", "re:.*\\.router\\.gate"] + # ignore: ["output", "re:.*\\.router\\.gate", "re:layers\\.\\d+\\.attention.wq", "re:layers\\.\\d+\\.attention.wk", "re:layers\\.\\d+\\.attention.wv", "re:layers\\.\\d+\\.attention.wo"] + use_2dblock_x: false + use_2dblock_w: true + use_hadamard: true + use_sr_grad: true + + # differential gradient estimation, disabled by default + use_dge: false + + # choices: none, static, dynamic + clip_mode: none + + # 2-level scaling, disabled by default + # NVFP4 requires tensorwise scaling or you will get overflow issues + two_level_scaling: none + # block scale selection: "default", "midmax-legacy", or "uos" + blockscale_selection: "midmax-legacy" + # weight de-oscillation, disabled by default + # deosc_step: 2000 + # deosc_period: 200 + # deosc_ratio: 4.0 diff --git a/alto/models/gpt_oss/configs/lpt_recipe_uos.yaml b/alto/models/gpt_oss/configs/lpt_recipe_uos.yaml new file mode 100644 index 00000000..81b9f795 --- /dev/null +++ b/alto/models/gpt_oss/configs/lpt_recipe_uos.yaml @@ -0,0 +1,29 @@ +training_stage: + lpt_modifiers: + LowPrecisionTrainingModifier: + scheme: "mxfp4" + targets: ["Linear", "GptOssGroupedExperts"] + # targets: ["Linear"] + # ignore: ["output", "re:.*\\.router\\.gate", "re:layers\\.0.*", "re:layers\\.22.*", "re:layers\\.23.*"] + ignore: ["output", "re:.*\\.router\\.gate"] + # ignore: ["output", "re:.*\\.router\\.gate", "re:layers\\.\\d+\\.attention.wq", "re:layers\\.\\d+\\.attention.wk", "re:layers\\.\\d+\\.attention.wv", "re:layers\\.\\d+\\.attention.wo"] + use_2dblock_x: false + use_2dblock_w: true + use_hadamard: true + use_sr_grad: true + + # differential gradient estimation, disabled by default + use_dge: false + + # choices: none, static, dynamic + clip_mode: none + + # 2-level scaling, disabled by default + # NVFP4 requires tensorwise scaling or you will get overflow issues + two_level_scaling: none + # block scale selection: "default", "midmax-legacy", or "uos" + blockscale_selection: "uos" + # weight de-oscillation, disabled by default + # deosc_step: 2000 + # deosc_period: 200 + # deosc_ratio: 4.0 diff --git a/alto/models/gpt_oss/configs/lpt_recipe_uos6.yaml b/alto/models/gpt_oss/configs/lpt_recipe_uos6.yaml new file mode 100644 index 00000000..089c3eb3 --- /dev/null +++ b/alto/models/gpt_oss/configs/lpt_recipe_uos6.yaml @@ -0,0 +1,29 @@ +training_stage: + lpt_modifiers: + LowPrecisionTrainingModifier: + scheme: "mxfp4" + targets: ["Linear", "GptOssGroupedExperts"] + # targets: ["Linear"] + # ignore: ["output", "re:.*\\.router\\.gate", "re:layers\\.0.*", "re:layers\\.22.*", "re:layers\\.23.*"] + ignore: ["output", "re:.*\\.router\\.gate"] + # ignore: ["output", "re:.*\\.router\\.gate", "re:layers\\.\\d+\\.attention.wq", "re:layers\\.\\d+\\.attention.wk", "re:layers\\.\\d+\\.attention.wv", "re:layers\\.\\d+\\.attention.wo"] + use_2dblock_x: false + use_2dblock_w: true + use_hadamard: true + use_sr_grad: true + + # differential gradient estimation, disabled by default + use_dge: false + + # choices: none, static, dynamic + clip_mode: none + + # 2-level scaling, disabled by default + # NVFP4 requires tensorwise scaling or you will get overflow issues + two_level_scaling: none + # block scale selection: "default", "midmax-legacy", or "uos" + blockscale_selection: "uos6" + # weight de-oscillation, disabled by default + # deosc_step: 2000 + # deosc_period: 200 + # deosc_ratio: 4.0 diff --git a/alto/models/gpt_oss/configs/mxfp4_base.yaml b/alto/models/gpt_oss/configs/mxfp4_base.yaml new file mode 100644 index 00000000..4d40c402 --- /dev/null +++ b/alto/models/gpt_oss/configs/mxfp4_base.yaml @@ -0,0 +1,29 @@ +training_stage: + lpt_modifiers: + LowPrecisionTrainingModifier: + scheme: "mxfp4" + targets: ["Linear", "GptOssGroupedExperts"] + # targets: ["Linear"] + # ignore: ["output", "re:.*\\.router\\.gate", "re:layers\\.0.*", "re:layers\\.22.*", "re:layers\\.23.*"] + ignore: ["output", "re:.*\\.router\\.gate"] + # ignore: ["output", "re:.*\\.router\\.gate", "re:layers\\.\\d+\\.attention.wq", "re:layers\\.\\d+\\.attention.wk", "re:layers\\.\\d+\\.attention.wv", "re:layers\\.\\d+\\.attention.wo"] + use_2dblock_x: false + use_2dblock_w: false + use_hadamard: false + use_sr_grad: false + + # differential gradient estimation, disabled by default + use_dge: false + + # choices: none, static, dynamic + clip_mode: none + + # 2-level scaling, disabled by default + # NVFP4 requires tensorwise scaling or you will get overflow issues + two_level_scaling: none + # block scale selection: "default" (round-even), "midmax", or "uos" + blockscale_selection: "default" + # weight de-oscillation, disabled by default + # deosc_step: 2000 + # deosc_period: 200 + # deosc_ratio: 4.0 diff --git a/alto/models/llama3/config_registry.py b/alto/models/llama3/config_registry.py index 097d1256..c08c2c6e 100644 --- a/alto/models/llama3/config_registry.py +++ b/alto/models/llama3/config_registry.py @@ -19,13 +19,23 @@ "llama3_debugmodel", "llama3_debugmodel_opt", "llama3_debugmodel_lpt", + "llama3_debugmodel_adahop", + "llama3_debugmodel_adahop_short", "llama3_1b", "llama3_1b_opt", "llama3_1b_lpt", + "llama3_1b_lpt_fwdonly", + "llama3_1b_lpt_hadamard", + "llama3_1b_adahop", "llama3_8b", "llama3_8b_pretrain", + "llama3_8b_random_init", "llama3_8b_opt", + "llama3_8b_bf16", "llama3_8b_lpt", + "llama3_8b_lpt_fwdonly", + "llama3_8b_lpt_hadamard", + "llama3_8b_adahop", "llama3_1b_gptq", "llama3_1b_awq", "llama3_1b_mx9_wa", @@ -75,12 +85,33 @@ def llama3_debugmodel_opt() -> Trainer.Config: def llama3_debugmodel_lpt() -> Trainer.Config: config = llama3_debugmodel() - config.training.steps = 10 + config.training.steps = 100 config.model_converters = ModelConvertersContainer.Config( converters=[ModelOptConverter.Config(recipe="./alto/models/llama3/configs/lpt_recipe.yaml",)],) return config +def llama3_debugmodel_adahop() -> Trainer.Config: + config = llama3_debugmodel() + # Need enough steps for the 30-step calibration plus a handful of post-calibration + # iterations to confirm Phase-B wrappers are actually exercised. + config.training.steps = 130 + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/llama3/configs/lpt_adahop_recipe.yaml",), + ],) + return config + + +def llama3_debugmodel_adahop_short() -> Trainer.Config: + """3-step calibration variant for fast debug iteration on cluster.""" + config = llama3_debugmodel() + config.training.steps = 8 # 3 calibration + 5 post-Phase-B + config.model_converters = ModelConvertersContainer.Config(converters=[ + ModelOptConverter.Config(recipe="./alto/models/llama3/configs/lpt_adahop_debug_recipe.yaml",), + ],) + return config + + def llama3_1b() -> Trainer.Config: config = llama3_1b_orig() config.hf_assets_path = "/group/archive_dataset_6_nobkup/archive_modelzoo/sequence_learning/weights/nlp-pretrained-model/meta-llama/Llama-3.2-1B" @@ -122,6 +153,33 @@ def llama3_1b_lpt() -> Trainer.Config: return config +def llama3_1b_lpt_fwdonly() -> Trainer.Config: + """MXFP4 forward-only: quantize the forward, keep the backward in bf16 + (gradient unquantized). Isolates the training-quality impact of forward-only + vs. full low-precision.""" + config = llama3_1b() + config.training.steps = 1000 + config.model_converters = ModelConvertersContainer.Config( + converters=[ModelOptConverter.Config(recipe="./alto/models/llama3/configs/lpt_recipe_fwdonly.yaml",)],) + return config + + +def llama3_1b_lpt_hadamard() -> Trainer.Config: + config = llama3_1b() + config.training.steps = 1000 + config.model_converters = ModelConvertersContainer.Config( + converters=[ModelOptConverter.Config(recipe="./alto/models/llama3/configs/lpt_hadamard_recipe.yaml",)],) + return config + + +def llama3_1b_adahop() -> Trainer.Config: + config = llama3_1b() + config.training.steps = 1000 + config.model_converters = ModelConvertersContainer.Config( + converters=[ModelOptConverter.Config(recipe="./alto/models/llama3/configs/lpt_adahop_recipe.yaml",)],) + return config + + def llama3_8b_pretrain() -> Trainer.Config: config = llama3_8b_orig() config.hf_assets_path = "/huggingface/hub/models--unsloth--Llama-3.1-8B/snapshots/3f0d51f8e5640f98f1a96ea9044a0e55c0a83814" @@ -166,6 +224,48 @@ def llama3_8b_opt() -> Trainer.Config: return config +def llama3_8b_random_init() -> Trainer.Config: + """Random-init Llama 3.1 8B for a controlled init checkpoint, mirroring + AdaHOP's pretrain-from-scratch recipe. The launcher runs this for 1 step, + relies on torchtitan's checkpoint.interval=1 to save a DCP at step 1, and + all quantization variants then `--checkpoint.initial_load_path` that file + so every variant starts from the SAME random weights (seeded with 1234). + """ + config = llama3_8b_orig() + # Tokenizer (overridden on CLI anyway), no weight initial load. + config.hf_assets_path = LLAMA3_8B_PATH + config.metrics.log_freq = 1 + config.profiling.enable_profiling = False + config.training.local_batch_size = 1 + config.training.global_batch_size = 8 + config.training.seq_len = 2048 + config.dataloader = HuggingFaceTextDataLoader.Config(dataset="c4") + config.activation_checkpoint.mode = "selective" + config.activation_checkpoint.selective_ac_option = "op" + config.checkpoint.enable = True + config.checkpoint.interval = 1 # save IMMEDIATELY so we can branch from this + config.validator.enable = False + config.debug.seed = 1234 + return config + + +def llama3_8b_bf16() -> Trainer.Config: + """Plain bf16 full-precision baseline (no quantization converter). The + reference point the mxfp4 / fwdonly variants compare against. Same + architecture / data / init as llama3_8b_pretrain, but with checkpointing + enabled so a long run is resumable.""" + config = llama3_8b_pretrain() + config.dump_folder = "llama3_8b-pretrain-subset-bf16-outputs" + # Fresh random init at debug.seed (1234), same as the gpt_oss runs — do NOT + # load pretrained weights. Reproducible across launches at fixed parallelism. + config.checkpoint.initial_load_path = None + config.checkpoint.initial_load_in_hf = False + config.checkpoint.enable = True + config.checkpoint.interval = 500 + config.checkpoint.keep_latest_k = 2 + return config + + def llama3_8b_lpt() -> Trainer.Config: config = llama3_8b_pretrain() config.dump_folder = "llama3_8b-mi308-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-gbs384-lr1e-4-outputs" @@ -174,6 +274,40 @@ def llama3_8b_lpt() -> Trainer.Config: return config +def llama3_8b_lpt_fwdonly() -> Trainer.Config: + """MXFP4 forward-only: quantize the forward, keep the backward in bf16 + (gradient unquantized). Isolates the training-quality impact of forward-only + vs. full low-precision. 8B is untied-weights so it runs (unlike 1B).""" + config = llama3_8b_pretrain() + config.dump_folder = "llama3_8b-pretrain-subset-mxfp4-fwdonly-outputs" + # Fresh random init at debug.seed (1234), same as the gpt_oss runs — do NOT + # load pretrained weights. Reproducible across launches at fixed parallelism. + config.checkpoint.initial_load_path = None + config.checkpoint.initial_load_in_hf = False + config.checkpoint.enable = True + config.checkpoint.interval = 500 + config.checkpoint.keep_latest_k = 2 + config.model_converters = ModelConvertersContainer.Config( + converters=[ModelOptConverter.Config(recipe="./alto/models/llama3/configs/lpt_recipe_fwdonly.yaml",)],) + return config + + +def llama3_8b_lpt_hadamard() -> Trainer.Config: + config = llama3_8b_pretrain() + config.training.steps = 1000 + config.model_converters = ModelConvertersContainer.Config( + converters=[ModelOptConverter.Config(recipe="./alto/models/llama3/configs/lpt_hadamard_recipe.yaml",)],) + return config + + +def llama3_8b_adahop() -> Trainer.Config: + config = llama3_8b_pretrain() + config.training.steps = 1000 + config.model_converters = ModelConvertersContainer.Config( + converters=[ModelOptConverter.Config(recipe="./alto/models/llama3/configs/lpt_adahop_recipe.yaml",)],) + return config + + def llama3_1b_gptq() -> Trainer.Config: config = llama3_1b() config.training.steps = 1 diff --git a/alto/models/llama3/configs/lpt_adahop_debug_recipe.yaml b/alto/models/llama3/configs/lpt_adahop_debug_recipe.yaml new file mode 100644 index 00000000..96df8392 --- /dev/null +++ b/alto/models/llama3/configs/lpt_adahop_debug_recipe.yaml @@ -0,0 +1,21 @@ +training_stage: + lpt_modifiers: + LowPrecisionTrainingModifier: + scheme: "mxfp4_adahop" + targets: ["Linear"] + ignore: ["output"] + AdaHOPModifier: + enabled: true + use_hadamard: true + use_randomized_hadamard: false + calibration_steps: 3 + layer_transform_config: + "row-row": "hadamard" + "row-none": "inner_outlier_extract_left" + "row-col": "inner_outlier_extract_right" + "col-row": "hadamard" + "col-none": "hadamard" + "col-col": "full_precision" + "none-row": "hadamard" + "none-none": "hadamard" + "none-col": "inner_outlier_extract_right" diff --git a/alto/models/llama3/configs/lpt_adahop_recipe.yaml b/alto/models/llama3/configs/lpt_adahop_recipe.yaml new file mode 100644 index 00000000..76eaf172 --- /dev/null +++ b/alto/models/llama3/configs/lpt_adahop_recipe.yaml @@ -0,0 +1,23 @@ +training_stage: + lpt_modifiers: + LowPrecisionTrainingModifier: + scheme: "mxfp4_adahop" + targets: ["Linear"] + ignore: ["output"] + AdaHOPModifier: + enabled: true + use_hadamard: true + use_randomized_hadamard: false + calibration_steps: 30 + # Mirrors AdaHOP's reference config + # (3rdparty/adahop/torchtitan/models/llama3/train_configs/llama3_3b_mxfp4_adahop_lv2.toml). + layer_transform_config: + "row-row": "hadamard" + "row-none": "inner_outlier_extract_left" + "row-col": "inner_outlier_extract_right" + "col-row": "hadamard" + "col-none": "hadamard" + "col-col": "full_precision" + "none-row": "hadamard" + "none-none": "hadamard" + "none-col": "inner_outlier_extract_right" diff --git a/alto/models/llama3/configs/lpt_hadamard_recipe.yaml b/alto/models/llama3/configs/lpt_hadamard_recipe.yaml new file mode 100644 index 00000000..d5ff5859 --- /dev/null +++ b/alto/models/llama3/configs/lpt_hadamard_recipe.yaml @@ -0,0 +1,13 @@ +training_stage: + lpt_modifiers: + LowPrecisionTrainingModifier: + scheme: "mxfp4" + targets: ["Linear"] + ignore: ["output"] + use_2dblock_x: false + use_2dblock_w: true + use_hadamard: true + use_sr_grad: true + use_dge: false + clip_mode: none + two_level_scaling: none diff --git a/alto/models/llama3/configs/lpt_recipe_fwdonly.yaml b/alto/models/llama3/configs/lpt_recipe_fwdonly.yaml new file mode 100644 index 00000000..a32b0d03 --- /dev/null +++ b/alto/models/llama3/configs/lpt_recipe_fwdonly.yaml @@ -0,0 +1,21 @@ +training_stage: + lpt_modifiers: + LowPrecisionTrainingModifier: + scheme: "mxfp4" + targets: ["Linear"] + ignore: ["output"] + use_2dblock_x: false + use_2dblock_w: true + # This mirrors lpt_hadamard_recipe.yaml, plus full_precision_backward. + # full_precision_backward keeps the backward (dgrad + wgrad) in bf16 with the + # gradient unquantized — only the forward is MXFP4. use_hadamard / use_sr_grad + # act ONLY on the (now-skipped) quantized backward in this codebase, so they + # are inert here; left true to show this is the hadamard recipe with the + # backward ignored. + full_precision_backward: true + use_hadamard: true + use_sr_grad: true + use_dge: false + clip_mode: none + two_level_scaling: none + lora_rank: 0 diff --git a/alto/modifiers/debug/__init__.py b/alto/modifiers/debug/__init__.py new file mode 100644 index 00000000..85eac522 --- /dev/null +++ b/alto/modifiers/debug/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT diff --git a/alto/modifiers/debug/debug_observer.py b/alto/modifiers/debug/debug_observer.py new file mode 100644 index 00000000..8f7e02ad --- /dev/null +++ b/alto/modifiers/debug/debug_observer.py @@ -0,0 +1,247 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""DebugObserverModifier — per-step tensor capture for low-precision MoE debugging. + +Hooks into a normal training run, captures the raw tensors that flow through +targeted layers at every captured iteration, and dumps everything to a single +torch.save'd file at finalization. + +See DEBUG_OBSERVER.md for the full output schema and quickstart. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +import torch +from pydantic import Field, PrivateAttr +from torch.nn import Module +from compressed_tensors.utils import match_named_modules +from torchtitan.tools.logging import logger + +from alto.modifiers import Modifier +from alto.modifiers.debug.observer_hooks import ( + make_linear_fwd_pre_hook, + make_linear_bwd_hook, + make_grouped_experts_fwd_pre_hook, + make_grouped_experts_bwd_hook, + make_grouped_experts_grad_weight_hook, +) + +__all__ = ["DebugObserverModifier"] + + +class DebugObserverModifier(Modifier): + """Captures input, weight, grad_output, and grad_weight tensors for targeted + layers at configurable step intervals. Dumps to a .pt file on finalization. + + Composes with any other modifier (LPT, AdaHOP, BF16 baseline). Drop it into + any recipe YAML under a ``debug_modifiers`` section. + """ + + targets: list[str] = Field(default_factory=lambda: ["Linear", "GptOssGroupedExperts"]) + ignore: list[str] = Field(default_factory=lambda: ["output", "re:.*\\.router\\.gate"]) + + capture_every: int = 1 + max_captures: int = 10 + output_path: str = "./outputs/debug_obs.pt" + + capture_input: bool = True + capture_weight: bool = True + capture_grad_output: bool = True + capture_grad_weight: bool = True + + # --- private mutable state --- + _step_idx: int = PrivateAttr(default=0) + # captures[fqn] = {"active": bool, step_idx: {"input": T, ...}, ...} + _captures: dict = PrivateAttr(default_factory=dict) + # number of steps for which we have already stored data + _n_captured: int = PrivateAttr(default=0) + # set to True once max_captures is reached and hooks are removed early + _detached: bool = PrivateAttr(default=False) + # one-element list shared with all hook closures so they always see the + # current step without needing a reference to self (avoids Pydantic issues) + _step_ref: list = PrivateAttr(default_factory=lambda: [0]) + # handles for parameter-level grad hooks (not managed by HooksMixin) + _param_hook_handles: list = PrivateAttr(default_factory=list) + + @property + def requires_training_mode(self) -> bool: + # Gradient tensors are only captured during backward passes, which only + # run when the trainer is in training mode. Force training mode so the + # BF16 baseline recipe doesn't silently degrade to forward-only. + return True + + def on_convert(self, model: Module, **kwargs) -> bool: + return True + + def on_initialize(self, model_parts: list[Module], **kwargs) -> bool: + for model_part in model_parts: + for fqn, module in match_named_modules(model_part, self.targets, self.ignore): + self._captures[fqn] = {"active": False} + cls_name = module.__class__.__name__ + + if cls_name == "Linear" or isinstance(module, torch.nn.Linear): + self.register_hook( + module, + make_linear_fwd_pre_hook( + self._captures, fqn, self._step_ref, + self.capture_input, self.capture_weight, + ), + "forward_pre", + ) + self.register_hook( + module, + make_linear_bwd_hook( + self._captures, fqn, self._step_ref, + self.capture_grad_output, self.capture_grad_weight, + ), + "full_backward", + ) + + elif cls_name.endswith("GroupedExperts"): + self.register_hook( + module, + make_grouped_experts_fwd_pre_hook( + self._captures, fqn, self._step_ref, + self.capture_input, self.capture_weight, + ), + "forward_pre", + ) + self.register_hook( + module, + make_grouped_experts_bwd_hook( + self._captures, fqn, self._step_ref, self.capture_grad_output, + ), + "full_backward", + ) + if self.capture_grad_weight: + for attr, key in [ + ("mlp1_weight", "grad_mlp1_weight"), + ("mlp2_weight", "grad_mlp2_weight"), + ]: + param = getattr(module, attr, None) + if param is not None: + handle = param.register_hook( + make_grouped_experts_grad_weight_hook( + self._captures, fqn, self._step_ref, + key, self.capture_grad_weight, + ) + ) + self._param_hook_handles.append(handle) + + logger.info( + f"DebugObserverModifier: monitoring {len(self._captures)} layers, " + f"capture_every={self.capture_every}, max_captures={self.max_captures}" + ) + return True + + def on_pre_step(self, model_parts: list[Module], **kwargs) -> bool: + self._step_idx += 1 + self._step_ref[0] = self._step_idx + + if self._detached: + return True + + should_capture = ( + (self._step_idx % self.capture_every == 0) and + (self._n_captured < self.max_captures) + ) + for fqn in self._captures: + self._captures[fqn]["active"] = should_capture + + return True + + def on_post_step(self, model_parts: list[Module], **kwargs) -> bool: + if self._detached: + return True + + # Only count a step as captured if a gradient was actually written. + # Calibration runs only forward passes — input/weight are stored but + # grad_output/grad_weight are not. Counting those would exhaust + # max_captures before any training backward fires. + step_idx = self._step_idx + grad_keys = {"grad_output", "grad_weight", "grad_mlp1_weight", "grad_mlp2_weight"} + captured_this_step = any( + any(k in grad_keys for k in self._captures[fqn].get(step_idx, {}).keys()) + for fqn in self._captures + ) + if captured_this_step: + self._n_captured += 1 + logger.debug(f"DebugObserverModifier: captured step {step_idx} ({self._n_captured}/{self.max_captures})") + else: + # Forward-only step (calibration): discard the partial capture so + # it doesn't pollute the dump with gradient-free entries. + for fqn in self._captures: + self._captures[fqn].pop(step_idx, None) + + # Deactivate all gates + for fqn in self._captures: + self._captures[fqn]["active"] = False + + # Detach once max reached + if self._n_captured >= self.max_captures: + logger.info(f"DebugObserverModifier: reached max_captures={self.max_captures}, detaching hooks") + self._detach() + + return True + + def on_finalize(self, model_parts: list[Module], **kwargs) -> bool: + if not self._detached: + self._detach() + self._dump() + return True + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + def _detach(self) -> None: + self.remove_hooks() + for h in self._param_hook_handles: + h.remove() + self._param_hook_handles.clear() + self._detached = True + + def _dump(self) -> None: + rank = int(os.environ.get("RANK", 0)) + path = Path(self.output_path) + if rank != 0: + # Insert rank suffix: debug_obs.pt -> debug_obs_rank1.pt + path = path.with_stem(f"{path.stem}_rank{rank}") + + path.parent.mkdir(parents=True, exist_ok=True) + + captured_steps = sorted({ + step + for fqn, data in self._captures.items() + for step in data + if step != "active" + }) + + layer_shapes: dict[str, Any] = {} + for fqn, data in self._captures.items(): + for step, tensors in data.items(): + if step == "active": + continue + layer_shapes[fqn] = {k: list(t.shape) for k, t in tensors.items()} + break + + blob: dict[str, Any] = { + fqn: {step: tensors for step, tensors in data.items() if step != "active"} + for fqn, data in self._captures.items() + } + blob["_meta"] = { + "rank": rank, + "iterations_captured": captured_steps, + "layer_shapes": layer_shapes, + "capture_every": self.capture_every, + "max_captures": self.max_captures, + "mlp2_input_captured": False, # v1 limitation + } + + torch.save(blob, path) + logger.info(f"DebugObserverModifier: saved {len(captured_steps)} captures to {path}") diff --git a/alto/modifiers/debug/moe_pattern_hooks.py b/alto/modifiers/debug/moe_pattern_hooks.py new file mode 100644 index 00000000..bb023b26 --- /dev/null +++ b/alto/modifiers/debug/moe_pattern_hooks.py @@ -0,0 +1,222 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Pure-Python helpers for ``MoEMatmulPatternObserverModifier``. + +torch-only (no alto / triton imports) so it can be exec'd in isolation on a +CPU-only box, following the same testing pattern as ``observer_hooks.py``. + +A grouped-GEMM expert layer runs two ``torch._grouped_mm`` calls (MLP1 then +MLP2). Each grouped GEMM's forward+backward is three matmuls whose inputs are +drawn from ``{x, w, grad_output}``: + + forward_y : x @ w (operands x, w) + backward_gx : grad_output @ wᵀ (operands grad_output, w) + backward_gw : grad_outputᵀ @ x (operands grad_output, x) + +For each expert we detect the AdaHOP outlier pattern (row/col/none) of each +operand and combine them into the AdaHOP "T-pair" per matmul path, matching +``alto/modifiers/lpt/adahop_internals/pattern_aggregation.py``. + +Weight-orientation note: the model calls ``_grouped_mm(x, W.transpose(-2,-1))`` +so the ``w`` operand reaching the patched op is ``[E, K, N]`` (in, out). AdaHOP +classifies weights in their ``[out, in]`` orientation, so we transpose each +expert slice back to ``[N, K]`` before ``detect`` — keeping the T-pairs directly +comparable to AdaHOP calibration. +""" + +from __future__ import annotations + +from collections import Counter +from typing import Any, Callable, Dict, List, Optional + +import torch + +# The three matmul paths and which operands feed each (for stats bookkeeping). +MATMUL_PATHS = ("forward_y", "backward_gx", "backward_gw") + + +def _opposite_pattern(pattern: str) -> str: + """row<->col, none->none. Mirrors pattern_aggregation._opposite_pattern.""" + if pattern == "row": + return "col" + if pattern == "col": + return "row" + return "none" + + +def _to_2d_float(t: torch.Tensor) -> torch.Tensor: + """Detach, upcast to float32 (bf16 var/std are imprecise), flatten to 2D.""" + t = t.detach().float() + if t.dim() > 2: + t = t.reshape(-1, t.shape[-1]) + return t + + +def _cv_row_col(t2d: torch.Tensor) -> tuple[float, float]: + """Shape-normalized coefficient of variation of row/col variances — the + quantity ``detect_outlier_pattern`` thresholds on. Returned for insight.""" + import math + + if t2d.numel() == 0 or t2d.dim() != 2: + return float("nan"), float("nan") + num_rows, num_cols = t2d.shape + row_var = t2d.var(dim=1) + col_var = t2d.var(dim=0) + cv_row = (row_var.std() / (row_var.mean() + 1e-8)).item() + cv_col = (col_var.std() / (col_var.mean() + 1e-8)).item() + cv_row /= math.sqrt(2.0 / max(num_cols - 1, 1)) + cv_col /= math.sqrt(2.0 / max(num_rows - 1, 1)) + return cv_row, cv_col + + +def _operand_stats(t2d: torch.Tensor) -> Dict[str, float]: + """Lightweight per-operand summary — plain Python floats only.""" + if t2d.numel() == 0: + return {"absmax": float("nan"), "std": float("nan"), + "cv_row": float("nan"), "cv_col": float("nan")} + cv_row, cv_col = _cv_row_col(t2d) + return { + "absmax": t2d.abs().max().item(), + "std": t2d.std().item(), + "cv_row": cv_row, + "cv_col": cv_col, + } + + +def _offs_to_bounds(offs: List[int], num_experts: int) -> List[tuple[int, int]]: + """Turn a cumulative-sum offsets list into [start, end) row bounds per + expert. offs[i] is the running token count through expert i; tokens beyond + offs[num_experts-1] are grouped-GEMM tail padding and are dropped.""" + bounds = [] + prev = 0 + for e in range(num_experts): + end = int(offs[e]) if e < len(offs) else prev + bounds.append((prev, end)) + prev = end + return bounds + + +def build_expert_records( + x: torch.Tensor, + w: torch.Tensor, + grad_output: torch.Tensor, + offs: List[int], + detect: Callable[[torch.Tensor], str], +) -> Dict[int, Dict[str, Any]]: + """Per (local) expert, detect operand patterns and build the T-pair records. + + Args: + x: activation into this GEMM, ``[T, K]`` (as fed to the matmul). + w: weight operand as fed to ``_grouped_mm``, ``[E, K, N]``. + grad_output: gradient of this GEMM's output, ``[T, N]``. + offs: cumulative per-expert token counts (Python ints). + detect: ``detect_outlier_pattern``-style callable returning + ``"row"|"col"|"none"`` for a 2D tensor. + + Returns ``{local_expert_id: {"n_tokens": int, + "forward_y":{"pair":str}, "backward_gx":{"pair":str}, + "backward_gw":{"pair":str}, + "stats":{"x":{...},"w":{...},"grad_output":{...}}}}``. + """ + x2 = _to_2d_float(x) + go2 = _to_2d_float(grad_output) + num_experts = int(w.shape[0]) + bounds = _offs_to_bounds(offs, num_experts) + + records: Dict[int, Dict[str, Any]] = {} + for e in range(num_experts): + start, end = bounds[e] + n_tokens = max(0, end - start) + # weight back to [out, in] = [N, K] to match AdaHOP's convention. + w_e = _to_2d_float(w[e].transpose(-2, -1)) + + if n_tokens == 0: + # No tokens routed to this expert this step: activation/grad patterns + # are undefined. Weight pattern is still meaningful. + w_pat = detect(w_e) + records[e] = { + "n_tokens": 0, + "forward_y": {"pair": f"none-{_opposite_pattern(w_pat)}"}, + "backward_gx": {"pair": f"none-{w_pat}"}, + "backward_gw": {"pair": "none-none"}, + "stats": {"x": _operand_stats(x2[0:0]), + "w": _operand_stats(w_e), + "grad_output": _operand_stats(go2[0:0])}, + } + continue + + x_e = x2[start:end] + go_e = go2[start:end] + x_pat = detect(x_e) + w_pat = detect(w_e) + g_pat = detect(go_e) + + records[e] = { + "n_tokens": n_tokens, + # forward_y : T1=x_pat, T2=opposite(w_pat) + "forward_y": {"pair": f"{x_pat}-{_opposite_pattern(w_pat)}"}, + # backward_gx : T7=grad_pat, T8=w_pat + "backward_gx": {"pair": f"{g_pat}-{w_pat}"}, + # backward_gw : T4=opposite(grad_pat), T5=x_pat + "backward_gw": {"pair": f"{_opposite_pattern(g_pat)}-{x_pat}"}, + "stats": {"x": _operand_stats(x_e), + "w": _operand_stats(w_e), + "grad_output": _operand_stats(go_e)}, + } + return records + + +def _majority_pair(pairs: List[str]) -> Dict[str, Any]: + """Majority-vote a list of per-step T-pair strings for one matmul path. + + Ties break deterministically by the pair string (so re-runs agree). Returns + the winning pair, the full per-pair vote counts, and the vote total. + """ + counter = Counter(pairs) + if not counter: + return {"pair": "none-none", "votes": {}, "n": 0} + top = max(counter.values()) + winner = sorted(p for p, c in counter.items() if c == top)[0] + return {"pair": winner, "votes": dict(counter), "n": sum(counter.values())} + + +def accumulate_majority( + per_step_records: List[Dict[int, Dict[str, Any]]], +) -> Dict[int, Dict[str, Any]]: + """Collapse a list of per-step ``{expert_id: record}`` maps (one per observed + step, same GEMM) into a single ``{expert_id: majority_record}``. + + For each expert and each of the three matmul paths, the reported ``pair`` is + the majority across the steps in which that expert was routed at least one + token (empty-expert steps contribute a ``none-*`` vote, same as the raw + records, so a rarely-routed expert honestly shows up as mostly ``none``). + + Returns ``{expert_id: {"n_steps": int, "n_tokens_total": int, + "forward_y": {"pair","votes","n"}, "backward_gx": {...}, + "backward_gw": {...}}}``. + """ + expert_ids = sorted({eid for step in per_step_records for eid in step}) + out: Dict[int, Dict[str, Any]] = {} + for eid in expert_ids: + step_recs = [step[eid] for step in per_step_records if eid in step] + rec: Dict[str, Any] = { + "n_steps": len(step_recs), + "n_tokens_total": int(sum(r.get("n_tokens", 0) for r in step_recs)), + } + for path in MATMUL_PATHS: + pairs = [r[path]["pair"] for r in step_recs if path in r] + rec[path] = _majority_pair(pairs) + out[eid] = rec + return out + + +def extract_offs(args: tuple, kwargs: dict) -> Optional[torch.Tensor]: + """Pull the ``offs`` tensor from a ``torch._grouped_mm`` call. It is passed + as a keyword in the model, but accept a trailing positional too.""" + if "offs" in kwargs and kwargs["offs"] is not None: + return kwargs["offs"] + for a in args: + if isinstance(a, torch.Tensor) and a.dim() == 1 and a.dtype in (torch.int32, torch.int64): + return a + return None diff --git a/alto/modifiers/debug/moe_pattern_observer.py b/alto/modifiers/debug/moe_pattern_observer.py new file mode 100644 index 00000000..49c3080b --- /dev/null +++ b/alto/modifiers/debug/moe_pattern_observer.py @@ -0,0 +1,351 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""MoEMatmulPatternObserverModifier — per-expert AdaHOP outlier-pattern capture +for the two Grouped GEMMs (MLP1, MLP2) of gpt_oss MoE blocks. + +Intended flow: pretrain gpt_oss WITHOUT this modifier, load a checkpoint, apply +this modifier via a recipe, run a single training iteration, dump per-rank +results, then visualize with ``scripts/moe_pattern_viz.py``. + +Why not the module-boundary hooks of ``DebugObserverModifier``? +A ``forward_pre``/``full_backward`` hook on ``GptOssGroupedExperts`` only sees +MLP1's input and the module-output grad — it cannot see MLP2's input, MLP1's +output grad, or split operands per expert. Instead we intercept the two +``torch._grouped_mm`` calls INSIDE ``_run_experts_grouped_mm`` (scoped to the +targeted module via a contextvar) so we get all three matmul operands of each +GEMM plus the per-expert token offsets. Capture happens after the module has +already ``.to_local()``'d its DTensor weights and after token permute, so the +tensors are plain local tensors on the current rank — robust under FSDP/EP. + +Constraints: the observed step must run in EAGER mode (the monkeypatch is +invisible to a compiled graph) and the low-precision path is backend-agnostic +(works on the CDNA3 loop fallback and the CDNA4 kernels alike). +""" + +from __future__ import annotations + +import contextvars +import os +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional + +import torch +from pydantic import Field, PrivateAttr +from torch.nn import Module +from compressed_tensors.utils import match_named_modules +from torchtitan.tools.logging import logger + +from alto.modifiers import Modifier +from alto.modifiers.debug.moe_pattern_hooks import ( + accumulate_majority, + build_expert_records, + extract_offs, +) + +__all__ = ["MoEMatmulPatternObserverModifier"] + +# Contextvar carrying the currently-executing experts module's capture context, +# or None when no targeted module is running. Shared by the patched _grouped_mm. +_ACTIVE_CTX: contextvars.ContextVar[Optional[dict]] = contextvars.ContextVar( + "_moe_pattern_active_ctx", default=None +) + +# gemm counter index -> logical name (order of _grouped_mm calls in the experts fwd) +_GEMM_NAMES = ("mlp1", "mlp2") + + +class MoEMatmulPatternObserverModifier(Modifier): + """Captures per-expert AdaHOP outlier-pattern T-pairs for the three matmuls + of each MoE Grouped GEMM (MLP1, MLP2), at one observed step, and dumps them. + """ + + targets: List[str] = Field(default_factory=lambda: ["GptOssGroupedExperts"]) + ignore: List[str] = Field(default_factory=list) + + capture_every: int = 1 + max_captures: int = 1 + output_path: str = "./outputs/moe_patterns.pt" + + # forwarded to detect_outlier_pattern + threshold_ratio: float = 2.0 + kurtosis_threshold: float = 0.0 + + @property + def requires_training_mode(self) -> bool: + # grad_output patterns require a real backward pass. + return True + + # --- private state --- + _step_idx: int = PrivateAttr(default=0) + _n_captured: int = PrivateAttr(default=0) + _detached: bool = PrivateAttr(default=False) + _active: bool = PrivateAttr(default=False) + # results[fqn][step_idx][gemm] = {global_expert_id: record} + _results: dict = PrivateAttr(default_factory=dict) + # per-fqn mesh info captured at initialize: {fqn: {ep_rank, ep_size, ...}} + _mesh_info: dict = PrivateAttr(default_factory=dict) + # saved originals for teardown + _orig_grouped_mm: Any = PrivateAttr(default=None) + _orig_forwards: dict = PrivateAttr(default_factory=dict) + _detect: Any = PrivateAttr(default=None) + _fqns: list = PrivateAttr(default_factory=list) + + # ------------------------------------------------------------------ + # lifecycle + # ------------------------------------------------------------------ + + def on_convert(self, model: Module, **kwargs) -> bool: + return True + + def on_initialize(self, model_parts: List[Module], **kwargs) -> bool: + # Bind detector once (imports the adahop bridge lazily so CPU unit tests + # of the hooks module don't drag in triton). + from alto._adahop_bridge import detect_outlier_pattern + + def _detect(t: torch.Tensor) -> str: + return detect_outlier_pattern( + t, + threshold_ratio=self.threshold_ratio, + kurtosis_threshold=self.kurtosis_threshold, + ) + + self._detect = _detect + + for model_part in model_parts: + for fqn, module in match_named_modules(model_part, self.targets, self.ignore): + self._results[fqn] = {} + self._mesh_info[fqn] = self._read_mesh_info(module) + self._fqns.append(fqn) + self._wrap_forward(module, fqn) + + # Install the grouped_mm patch once (no-op until a wrapped forward runs + # AND the capture gate is active). + self._orig_grouped_mm = torch._grouped_mm + torch._grouped_mm = self._make_patched_grouped_mm(self._orig_grouped_mm) + + logger.info( + f"MoEMatmulPatternObserverModifier: monitoring {len(self._fqns)} expert " + f"layers, max_captures={self.max_captures}" + ) + return True + + def on_pre_step(self, model_parts: List[Module], **kwargs) -> bool: + self._step_idx += 1 + if self._detached: + return True + self._active = ( + (self._step_idx % self.capture_every == 0) + and (self._n_captured < self.max_captures) + ) + return True + + def on_post_step(self, model_parts: List[Module], **kwargs) -> bool: + if self._detached: + return True + # A step counts as captured only if a grad path actually populated a + # T-pair (i.e. a backward ran), mirroring DebugObserverModifier. + step = self._step_idx + captured = any( + step in self._results[fqn] and self._results[fqn][step] + for fqn in self._fqns + ) + if captured: + self._n_captured += 1 + logger.debug( + f"MoEMatmulPatternObserverModifier: captured step {step} " + f"({self._n_captured}/{self.max_captures})" + ) + self._active = False + if self._n_captured >= self.max_captures: + logger.info("MoEMatmulPatternObserverModifier: reached max_captures, detaching") + self._detach() + return True + + def on_finalize(self, model_parts: List[Module], **kwargs) -> bool: + if not self._detached: + self._detach() + self._dump() + return True + + # ------------------------------------------------------------------ + # forward wrapping + grouped_mm patch + # ------------------------------------------------------------------ + + def _wrap_forward(self, module: Module, fqn: str) -> None: + """Replace the module's bound forward with a wrapper that sets the + contextvar (fqn + fresh gemm counter) while the original runs.""" + orig_forward = module.forward + self._orig_forwards[fqn] = orig_forward + modifier = self + + def _wrapped(*args, **kwargs): + if modifier._detached or not modifier._active: + return orig_forward(*args, **kwargs) + ctx = {"fqn": fqn, "gemm_idx": 0} + token = _ACTIVE_CTX.set(ctx) + try: + return orig_forward(*args, **kwargs) + finally: + _ACTIVE_CTX.reset(token) + + module.forward = _wrapped # type: ignore[method-assign] + + def _make_patched_grouped_mm(self, orig: Callable) -> Callable: + modifier = self + + def _patched(*args, **kwargs): + out = orig(*args, **kwargs) + ctx = _ACTIVE_CTX.get() + if ctx is None or modifier._detached or not modifier._active: + return out + try: + modifier._on_grouped_mm(ctx, args, kwargs, out) + except Exception as exc: # never break training on a debug tool + logger.warning(f"MoEMatmulPatternObserverModifier: capture skipped ({exc})") + return out + + return _patched + + def _on_grouped_mm(self, ctx: dict, args: tuple, kwargs: dict, out: torch.Tensor) -> None: + gemm_idx = ctx["gemm_idx"] + ctx["gemm_idx"] = gemm_idx + 1 + if gemm_idx >= len(_GEMM_NAMES): + return # unexpected extra grouped_mm; ignore + gemm_name = _GEMM_NAMES[gemm_idx] + fqn = ctx["fqn"] + + # operands: _grouped_mm(x, w, offs=...) — w is already transposed to [E,K,N] + x = args[0] + w = args[1] + offs_t = extract_offs(args, kwargs) + if offs_t is None or w.dim() != 3: + return + offs = offs_t.detach().cpu().tolist() + + # Snapshot forward operands now; pair with grad_output on backward. + x_snap = x.detach() + w_snap = w.detach() + + step = self._step_idx + + def _grad_hook(grad_output: torch.Tensor): + if grad_output is None: + return + try: + records = build_expert_records( + x_snap, w_snap, grad_output, offs, self._detect + ) + records = self._relabel_global(fqn, records) + bucket = self._results[fqn].setdefault(step, {}) + bucket[gemm_name] = records + except Exception as exc: + logger.warning( + f"MoEMatmulPatternObserverModifier: grad capture skipped for " + f"{fqn}/{gemm_name} ({exc})" + ) + + if out.requires_grad: + out.register_hook(_grad_hook) + + # ------------------------------------------------------------------ + # sharding / ids + # ------------------------------------------------------------------ + + def _read_mesh_info(self, module: Module) -> dict: + """Derive EP/FSDP rank+size from the experts' DTensor weight mesh.""" + info = {"ep_rank": 0, "ep_size": 1, "fsdp_rank": 0, "fsdp_size": 1} + w = getattr(module, "mlp1_weight", None) + try: + from torch.distributed.tensor import DTensor + if isinstance(w, DTensor): + mesh = w.device_mesh + names = mesh.mesh_dim_names or () + for dim_name, keys in (("ep", ("ep_rank", "ep_size")), + ("dp_shard", ("fsdp_rank", "fsdp_size")), + ("efsdp", ("fsdp_rank", "fsdp_size"))): + if dim_name in names: + idx = names.index(dim_name) + info[keys[0]] = mesh.get_local_rank(dim_name) + info[keys[1]] = mesh.size(idx) + except Exception as exc: + logger.debug(f"MoEMatmulPatternObserverModifier: no mesh info ({exc})") + return info + + def _relabel_global(self, fqn: str, records: dict) -> dict: + """Map local expert ids to global ids using EP rank/size.""" + mi = self._mesh_info.get(fqn, {}) + ep_rank = mi.get("ep_rank", 0) + n_local = len(records) + base = ep_rank * n_local + return {base + local_id: rec for local_id, rec in records.items()} + + # ------------------------------------------------------------------ + # teardown + dump + # ------------------------------------------------------------------ + + def _detach(self) -> None: + if self._orig_grouped_mm is not None: + torch._grouped_mm = self._orig_grouped_mm + self._orig_grouped_mm = None + for fqn, orig in self._orig_forwards.items(): + # best-effort restore; the module object still lives in model_parts + pass + self._detached = True + + def _dump(self) -> None: + rank = int(os.environ.get("RANK", 0)) + path = Path(self.output_path) + # always rank-suffix so multi-rank runs don't clobber a single file. + path = path.with_stem(f"{path.stem}_rank{rank}") + path.parent.mkdir(parents=True, exist_ok=True) + + captured_steps = sorted({ + step for data in self._results.values() for step in data + }) + + layer_shapes = {fqn: self._mesh_info.get(fqn, {}) for fqn in self._fqns} + + blob: Dict[str, Any] = dict(self._results) + num_local = 0 + for fqn in self._fqns: + for step_data in self._results[fqn].values(): + for gemm_records in step_data.values(): + num_local = max(num_local, len(gemm_records)) + + # Majority-vote each expert's per-path T-pair across all captured steps: + # _majority[fqn][gemm][global_expert_id] = {n_steps, n_tokens_total, + # forward_y/backward_gx/backward_gw: {pair, votes, n}} + majority: Dict[str, Dict[str, Any]] = {} + for fqn in self._fqns: + per_gemm_steps: Dict[str, List[dict]] = {} + for step_data in self._results[fqn].values(): + for gemm, records in step_data.items(): + per_gemm_steps.setdefault(gemm, []).append(records) + majority[fqn] = { + gemm: accumulate_majority(step_list) + for gemm, step_list in per_gemm_steps.items() + } + blob["_majority"] = majority + + # global expert count = local * ep_size (uniform across ranks) + any_mi = next(iter(self._mesh_info.values()), {}) if self._mesh_info else {} + ep_size = any_mi.get("ep_size", 1) + blob["_meta"] = { + "rank": rank, + "ep_rank": any_mi.get("ep_rank", 0), + "ep_size": ep_size, + "fsdp_rank": any_mi.get("fsdp_rank", 0), + "fsdp_size": any_mi.get("fsdp_size", 1), + "num_local_experts": num_local, + "num_global_experts": num_local * ep_size, + "iterations_captured": captured_steps, + "mesh_info": layer_shapes, + "gemm_order": list(_GEMM_NAMES), + "matmul_paths": ["forward_y", "backward_gx", "backward_gw"], + } + torch.save(blob, path) + logger.info( + f"MoEMatmulPatternObserverModifier: saved {len(captured_steps)} step(s) " + f"for {len(self._fqns)} layer(s) to {path}" + ) diff --git a/alto/modifiers/debug/observer_hooks.py b/alto/modifiers/debug/observer_hooks.py new file mode 100644 index 00000000..5276d625 --- /dev/null +++ b/alto/modifiers/debug/observer_hooks.py @@ -0,0 +1,190 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Pure-Python hook factories for DebugObserverModifier. + +No triton imports — safe to run on a CPU-only box. + +Each factory returns a closure that: +- Checks the per-fqn ``active`` gate before doing anything. +- Detaches and moves tensors to CPU before storing them. +- Writes into ``captures[fqn][step_idx][key]``. + +The captures dict has the shape:: + + { + fqn: { + "active": bool, + step_idx: { + "input": Tensor, + "weight": Tensor, + "grad_output": Tensor, + "grad_weight": Tensor, # nn.Linear + "grad_mlp1_weight": Tensor, # GptOssGroupedExperts + "grad_mlp2_weight": Tensor, # GptOssGroupedExperts + }, + } + } +""" + +from __future__ import annotations + +from typing import Any + + +def _unwrap_weight(weight: Any) -> Any: + """Strip DTensor and TrainingWeightWrapperTensor wrappers to reach the + underlying storage tensor, mirroring the pattern in + alto/modifiers/lpt/adahop_internals/calibration_hooks.py.""" + w = weight.data + try: + from torch.distributed.tensor import DTensor + if isinstance(w, DTensor): + w = w._local_tensor + except Exception: + pass + # TrainingWeightWrapperTensor stores the raw data in ._data + w = getattr(w, "_data", w) + return w + + +def _store(captures: dict, fqn: str, step_idx: int, key: str, tensor: Any) -> None: + """Detach, move to CPU, and store tensor in captures.""" + captures[fqn].setdefault(step_idx, {})[key] = tensor.detach().cpu() + + +# --------------------------------------------------------------------------- +# nn.Linear hooks +# --------------------------------------------------------------------------- + +def make_linear_fwd_pre_hook( + captures: dict, + fqn: str, + step_ref: list, + capture_input: bool, + capture_weight: bool, +): + """Forward pre-hook for nn.Linear. + + Captures the raw (pre-quant) activation and unwrapped weight storage. + step_ref is a one-element list so the closure always sees the current step. + """ + + def _hook(module, args): + if not captures[fqn]["active"]: + return + step_idx = step_ref[0] + if capture_input and args: + _store(captures, fqn, step_idx, "input", args[0]) + if capture_weight: + try: + w = _unwrap_weight(module.weight) + _store(captures, fqn, step_idx, "weight", w) + except Exception: + pass + + return _hook + + +def make_linear_bwd_hook( + captures: dict, fqn: str, step_ref: list, capture_grad_output: bool, capture_grad_weight: bool = False +): + """Full-backward hook for nn.Linear. + + grad_input = (grad_x, grad_weight, grad_bias) — exactly what + MXFP4LinearFunction.backward returns, so grad_weight here is the + post-clip value, not the value accumulated into param.grad later. + grad_output = upstream gradient flowing into this layer's output. + """ + + def _hook(module, grad_input, grad_output): + if not captures[fqn]["active"]: + return + step_idx = step_ref[0] + if capture_grad_output: + go = grad_output[0] if isinstance(grad_output, (list, tuple)) else grad_output + if go is not None: + _store(captures, fqn, step_idx, "grad_output", go) + if capture_grad_weight and isinstance(grad_input, (list, tuple)) and len(grad_input) > 1: + gw = grad_input[1] + if gw is not None: + _store(captures, fqn, step_idx, "grad_weight", gw) + + return _hook + + +# --------------------------------------------------------------------------- +# GptOssGroupedExperts hooks +# --------------------------------------------------------------------------- + +def make_grouped_experts_fwd_pre_hook( + captures: dict, + fqn: str, + step_ref: list, + capture_input: bool, + capture_weight: bool, +): + """Forward pre-hook for GptOssGroupedExperts. + + Captures the activation entering mlp1 (args[0]) and the mlp1_weight storage. + The mlp2 activation is internal to _run_experts_grouped_mm and cannot be + captured at the module-hook level (v1 limitation, noted in _meta). + """ + + def _hook(module, args): + if not captures[fqn]["active"]: + return + step_idx = step_ref[0] + if capture_input and args: + _store(captures, fqn, step_idx, "input", args[0]) + if capture_weight: + for attr in ("mlp1_weight", "mlp2_weight"): + param = getattr(module, attr, None) + if param is not None: + try: + w = _unwrap_weight(param) + _store(captures, fqn, step_idx, attr, w) + except Exception: + pass + + return _hook + + +def make_grouped_experts_bwd_hook( + captures: dict, fqn: str, step_ref: list, capture_grad_output: bool +): + """Full-backward hook for GptOssGroupedExperts.""" + + def _hook(module, grad_input, grad_output): + if not captures[fqn]["active"]: + return + if not capture_grad_output: + return + go = grad_output[0] if isinstance(grad_output, (list, tuple)) else grad_output + if go is None: + return + _store(captures, fqn, step_ref[0], "grad_output", go) + + return _hook + + +def make_grouped_experts_grad_weight_hook( + captures: dict, + fqn: str, + step_ref: list, + key: str, + capture_grad_weight: bool, +): + """Parameter grad hook for mlp1_weight or mlp2_weight of GptOssGroupedExperts. + + key should be "grad_mlp1_weight" or "grad_mlp2_weight". + """ + + def _hook(grad): + if not captures[fqn]["active"]: + return + if not capture_grad_weight: + return + _store(captures, fqn, step_ref[0], key, grad) + + return _hook diff --git a/alto/modifiers/lpt/__init__.py b/alto/modifiers/lpt/__init__.py index 22cdef1e..71d00e68 100644 --- a/alto/modifiers/lpt/__init__.py +++ b/alto/modifiers/lpt/__init__.py @@ -3,5 +3,7 @@ # SPDX-License-Identifier: MIT from .base import LowPrecisionTrainingModifier +from .adahop import AdaHOPModifier +from .grad_clip import GradientClippingModifier -__all__ = ["LowPrecisionTrainingModifier"] +__all__ = ["LowPrecisionTrainingModifier", "AdaHOPModifier", "GradientClippingModifier"] diff --git a/alto/modifiers/lpt/adahop.py b/alto/modifiers/lpt/adahop.py new file mode 100644 index 00000000..cd167c42 --- /dev/null +++ b/alto/modifiers/lpt/adahop.py @@ -0,0 +1,357 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""AdaHOP calibration + per-slot Hadamard mode selection modifier. + +Strategy-2 design (see /home/alirezak/han_branch/porting_notes.txt). A SINGLE +wrapper class (:class:`MXFP4AdaHOPWrapper`) is used from model-conversion time. +During calibration its modes are all ``"none"`` (== plain MXFP4); when +calibration finishes the modes are set IN PLACE (no wrapper-type swap, no new +``nn.Parameter``), so optimizer state references stay valid. + +Lifecycle: + +1. ``LowPrecisionTrainingModifier`` (scheme=``"mxfp4_adahop"``) wraps every + targeted weight in :class:`MXFP4AdaHOPWrapper` (modes ``"none"``). +2. This modifier owns the Phase-A -> Phase-B transition: + + * ``on_convert`` configures the HadamardFactory. + * ``on_initialize`` discovers the wrapped linears. (If ``transform_config_path`` + is set, pre-baked modes are applied here and calibration is skipped.) + * ``on_pre_step`` — on the FIRST call (which happens AFTER + ``checkpointer.load`` in the forge trainer), it inspects the checkpointed + :class:`CalibrationStateManager`: + - calibration already completed -> apply the restored modes in place and + skip calibration (clean resume); + - otherwise -> arm transient *module* observation hooks and begin + calibration. It then opens a fresh per-step pattern bucket. + * ``on_post_step`` aggregates after ``calibration_steps``, maps patterns to + per-slot modes, applies them in place, records them in the + CalibrationStateManager (so they ride the next checkpoint), writes the + JSON artifact, and removes the observation hooks. + +Calibration observation uses transient forward-pre / full-backward hooks on the +``nn.Linear`` modules — no closures are stored on tensors, so checkpoints remain +picklable. +""" + +from typing import Any, Callable, Dict, List, Optional + +import torch +from pydantic import Field, PrivateAttr +from torch import nn +from torch.nn import Module +from torchtitan.tools.logging import logger + +from alto.modifiers import Modifier +from alto.modifiers.lpt.adahop_internals.calibration_hooks import ( + load_modes_from_json, + make_backward_hook, + make_forward_pre_hook, + write_modes_json, +) +from alto.modifiers.lpt.adahop_internals.calibration_state import ( + get_calibration_modes, + is_calibration_completed, + set_calibration_result, +) + +__all__ = ["AdaHOPModifier"] + +_NONE_MODES = {"forward_y": "none", "backward_gx": "none", "backward_gw": "none"} + + +class AdaHOPModifier(Modifier): + """Outlier-pattern-aware Hadamard calibration over MXFP4-wrapped linears.""" + + enabled: bool = True + + use_hadamard: bool = True + use_randomized_hadamard: bool = False + + calibration_steps: int = 30 + """Number of pre-step observations before transitioning to Phase B.""" + + layer_transform_config: Dict[str, str] = Field(default_factory=dict) + """Pattern-pair → ``TransformMode`` lookup. Keys are ``"row-row"``, + ``"col-col"``, etc.; values are members of ``TransformMode``.""" + + transform_config_path: Optional[str] = None + """If set, load a pre-baked per-layer mode JSON and skip calibration. + Format: ``{layer_fqn: {forward_y, backward_gx, backward_gw}}``.""" + + dump_json_path: Optional[str] = None + """If set, write the aggregated per-layer modes here after Phase B.""" + + _step_idx: int = PrivateAttr(default=0) + _per_step_patterns: List[Dict[str, Dict[str, str]]] = PrivateAttr(default_factory=list) + _fqn_to_wrapper: Dict[str, Any] = PrivateAttr(default_factory=dict) + _fqn_to_module: Dict[str, nn.Module] = PrivateAttr(default_factory=dict) + _handles: List[Any] = PrivateAttr(default_factory=list) + _phase_b_done: bool = PrivateAttr(default=False) + _resume_decided: bool = PrivateAttr(default=False) + _calibration_armed: bool = PrivateAttr(default=False) + + @property + def requires_training_mode(self) -> bool: + return True + + def on_convert(self, model: Module, **kwargs) -> bool: + # The wrapper swap (to MXFP4AdaHOPWrapper, modes "none") is performed by + # LowPrecisionTrainingModifier when scheme="mxfp4_adahop". Configure the + # HadamardFactory here so the transform object is available later. + if not self.enabled: + logger.info("[AdaHOP] Modifier disabled (enabled=False); pass-through only.") + return True + logger.info(f"[AdaHOP] Modifier active: use_hadamard={self.use_hadamard}, " + f"use_randomized_hadamard={self.use_randomized_hadamard}, " + f"calibration_steps={self.calibration_steps}, " + f"transform_config_path={self.transform_config_path}, " + f"dump_json_path={self.dump_json_path or ''}, " + f"layer_transform_config entries={len(self.layer_transform_config)}") + if not self.use_hadamard: + logger.info("[AdaHOP] use_hadamard=False; skipping HadamardFactory configuration.") + return True + from alto._adahop_bridge import HadamardFactory + HadamardFactory.configure(randomized=self.use_randomized_hadamard) + logger.info(f"[AdaHOP] HadamardFactory configured (randomized={self.use_randomized_hadamard}).") + return True + + def on_initialize(self, model_parts: list[Module], **kwargs) -> bool: + if not self.enabled: + return True + from alto.kernels.dispatch.adahop_tensor import MXFP4AdaHOPWrapper + + self._collect_wrappers(model_parts, MXFP4AdaHOPWrapper) + n_wrappers = len(self._fqn_to_wrapper) + logger.info(f"[AdaHOP] Discovered {n_wrappers} MXFP4AdaHOPWrapper-wrapped linears.") + if n_wrappers > 0: + sample = list(self._fqn_to_wrapper.keys())[:5] + logger.info(f"[AdaHOP] First {len(sample)} FQNs: {', '.join(sample)}" + f"{' ...' if n_wrappers > 5 else ''}") + elif not self._fqn_to_wrapper: + logger.warning("[AdaHOP] No MXFP4AdaHOPWrapper found; " + "is LowPrecisionTrainingModifier(scheme='mxfp4_adahop') in the recipe?") + + # Manual pre-baked modes override: applied now (before checkpointer.load), + # in place. Calibration is then skipped entirely. + if self.transform_config_path is not None: + modes_by_fqn = load_modes_from_json(self.transform_config_path) + logger.info(f"[AdaHOP] Pre-baked modes: {len(modes_by_fqn)} loaded from " + f"{self.transform_config_path}; applying in place (calibration skipped).") + self._log_mode_table(modes_by_fqn, aggregated=None) + self._apply_modes_in_place(modes_by_fqn) + self._resume_decided = True + # Otherwise: defer the calibrate-vs-resume decision to the first + # on_pre_step, which runs AFTER checkpointer.load() so the restored + # CalibrationStateManager is visible. + return True + + def on_pre_step(self, model_parts: list[Module], **kwargs) -> bool: + if not self.enabled or self._phase_b_done: + return True + + # First step after checkpointer.load(): decide resume vs. calibrate. + if not self._resume_decided: + self._resume_decided = True + if is_calibration_completed(): + modes_by_fqn = get_calibration_modes() + logger.info(f"[AdaHOP] Calibration restored from checkpoint " + f"({len(modes_by_fqn)} layers). Applying modes in place; " + f"skipping re-calibration.") + self._log_mode_table(modes_by_fqn, aggregated=None) + self._apply_modes_in_place(modes_by_fqn) + return True + # Fresh calibration. + self._arm_calibration() + + if self._step_idx >= self.calibration_steps: + return True + # Open a fresh dict for this step's pattern observations. + self._per_step_patterns.append({}) + logger.info(f"[AdaHOP] Calibration step {self._step_idx + 1}/{self.calibration_steps} opening bucket.") + return True + + def on_post_step(self, model_parts: list[Module], **kwargs) -> bool: + if not self.enabled or self._phase_b_done: + return True + if not self._calibration_armed: + return True + self._step_idx += 1 + + # Per-step observation summary + total_layers = len(self._fqn_to_wrapper) + bucket = self._per_step_patterns[-1] if self._per_step_patterns else {} + layers_with_data = len(bucket) + nx = sum(1 for v in bucket.values() if "x" in v) + nw = sum(1 for v in bucket.values() if "w" in v) + ng = sum(1 for v in bucket.values() if "grad_output" in v) + logger.info(f"[AdaHOP] Calibration step {self._step_idx}/{self.calibration_steps} observed: " + f"{layers_with_data}/{total_layers} layers reported patterns. " + f"(x, w, grad_output) presence: {nx}/{nw}/{ng}.") + + if self._step_idx < self.calibration_steps: + return True + + from alto.modifiers.lpt.adahop_internals.pattern_aggregation import (aggregate_patterns, patterns_to_modes) + logger.info(f"[AdaHOP] Calibration window complete. Aggregating {self._step_idx} steps " + f"of observations across {total_layers} layers.") + aggregated = aggregate_patterns(self._per_step_patterns) + logger.info(f"[AdaHOP] Aggregation done. Mapping pattern pairs → TransformModes " + f"via recipe table ({len(self.layer_transform_config)} entries).") + modes_by_fqn = patterns_to_modes(aggregated, self.layer_transform_config) + self._log_mode_table(modes_by_fqn, aggregated=aggregated) + + effective_dump_path = self.dump_json_path or "./outputs/adahop_calibration.json" + write_modes_json(effective_dump_path, aggregated, modes_by_fqn) + logger.info(f"[AdaHOP] JSON dump: wrote {len(modes_by_fqn)} layer modes to {effective_dump_path}") + + logger.info("[AdaHOP] Phase B (in-place mode application) starting...") + self._apply_modes_in_place(modes_by_fqn) + self._detach_observation_hooks() + # Record into the checkpointable calibration state so the next checkpoint + # carries the modes and a resume can skip calibration. + set_calibration_result(modes_by_fqn, self._step_idx) + logger.info(f"[AdaHOP] Phase B complete. Modes set in place for {len(modes_by_fqn)} layers; " + f"recorded in CalibrationStateManager. Observation hooks detached.") + return True + + def on_finalize(self, model_parts: list[Module], **kwargs) -> bool: + logger.info(f"[AdaHOP] Modifier finalize: phase_b_done={self._phase_b_done}, " + f"total_calibration_steps_seen={self._step_idx}.") + self._detach_observation_hooks() + return True + + def _log_mode_table( + self, + modes_by_fqn: Dict[str, Dict[str, str]], + aggregated: Optional[Dict[str, Dict[str, str]]] = None, + ) -> None: + """Emit a compact per-layer ``x w g -> forward_y backward_gx backward_gw`` table to stdout.""" + logger.info("[AdaHOP] Per-layer mode assignments:") + for fqn in sorted(modes_by_fqn.keys()): + slots = modes_by_fqn[fqn] + if aggregated and fqn in aggregated: + pats = aggregated[fqn] + prefix = (f"x={pats.get('x', '?')} " + f"w={pats.get('w', '?')} " + f"g={pats.get('grad_output', '?')} -> ") + else: + prefix = "" + logger.info(f" {fqn}: {prefix}" + f"forward_y={slots.get('forward_y', 'none')} " + f"backward_gx={slots.get('backward_gx', 'none')} " + f"backward_gw={slots.get('backward_gw', 'none')}") + + # ------------------------------------------------------------------ helpers + + def _arm_calibration(self) -> None: + """Register transient module observation hooks and begin Phase A.""" + if self._calibration_armed: + return + from alto._adahop_bridge import detect_outlier_pattern + + for fqn, module in self._fqn_to_module.items(): + h_fwd = module.register_forward_pre_hook( + make_forward_pre_hook(self, fqn, detect_outlier_pattern) + ) + h_bwd = module.register_full_backward_hook( + make_backward_hook(self, fqn, detect_outlier_pattern) + ) + self._handles.append(h_fwd) + self._handles.append(h_bwd) + self._calibration_armed = True + logger.info(f"[AdaHOP] Calibration armed: {len(self._fqn_to_module)} layers, " + f"{self.calibration_steps} steps. Module forward-pre + backward hooks registered.") + + def _make_ht_resolver(self) -> Callable[[torch.device], Any]: + if not self.use_hadamard: + return lambda _dev: None + from alto._adahop_bridge import HadamardFactory + ht_cache: Dict[torch.device, Any] = {} + + def _get_ht(device): + if device not in ht_cache: + ht_cache[device] = HadamardFactory.create_transform(device=device) + return ht_cache[device] + + return _get_ht + + def _collect_wrappers(self, model_parts, wrapper_cls) -> None: + # Attach to the canonical instance F.linear will see (see long comment + # below). Same discovery logic as before; only the target class changed + # to the single MXFP4AdaHOPWrapper. + # + # No-FSDP path: `module.weight` IS the subclass. FSDP path: the + # canonical instance lives at fsdp_param._sharded_local_tensor and is + # the one FSDP later passes into fsdp_post_all_gather; in-place mode + # updates there propagate to the unsharded param used in forward. + from torch.distributed.tensor import DTensor + try: + from torch.distributed.fsdp._fully_shard._fsdp_state import _get_module_fsdp_state + except Exception: + _get_module_fsdp_state = lambda _m: None # noqa: E731 + + self._fqn_to_wrapper.clear() + self._fqn_to_module.clear() + + module_param_to_fsdp_param = {} + for part in model_parts: + for _m in part.modules(): + state = _get_module_fsdp_state(_m) + if state is None or state._fsdp_param_group is None: + continue + for fp in state._fsdp_param_group.fsdp_params: + mi = getattr(fp, "_module_info", None) + if mi is None: + continue + module_param_to_fsdp_param[(id(mi.module), mi.param_name)] = fp + + for part in model_parts: + for module_fqn, module in part.named_modules(): + if not isinstance(module, nn.Linear): + continue + weight = getattr(module, "weight", None) + if weight is None: + continue + + fp = module_param_to_fsdp_param.get((id(module), "weight")) + if fp is not None: + inner = fp._sharded_local_tensor + elif isinstance(weight, wrapper_cls): + inner = weight + else: + raw = weight.data + if isinstance(raw, DTensor): + inner = raw._local_tensor + else: + inner = raw + + if not isinstance(inner, wrapper_cls): + continue + self._fqn_to_wrapper[module_fqn] = inner + self._fqn_to_module[module_fqn] = module + + def _apply_modes_in_place(self, modes_by_fqn: Dict[str, Dict[str, str]]) -> None: + """Set per-slot modes on every wrapper IN PLACE (no Parameter swap).""" + ht_resolver = self._make_ht_resolver() + n = 0 + for fqn, wrapper in self._fqn_to_wrapper.items(): + modes = modes_by_fqn.get(fqn, _NONE_MODES) + ht = ht_resolver(wrapper._data.device) + wrapper.set_modes( + forward_y_mode=modes.get("forward_y", "none"), + backward_gx_mode=modes.get("backward_gx", "none"), + backward_gw_mode=modes.get("backward_gw", "none"), + hadamard_transform=ht, + ) + n += 1 + ht_state = "attached" if self.use_hadamard else "none" + logger.info(f"[AdaHOP] Applied modes in place for {n} layers (hadamard={ht_state}).") + self._phase_b_done = True + + def _detach_observation_hooks(self) -> None: + for handle in self._handles: + handle.remove() + self._handles.clear() + self._calibration_armed = False diff --git a/alto/modifiers/lpt/adahop_internals/__init__.py b/alto/modifiers/lpt/adahop_internals/__init__.py new file mode 100644 index 00000000..f3e33abb --- /dev/null +++ b/alto/modifiers/lpt/adahop_internals/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Internal building blocks for the AdaHOP port. + +ALTO-side implementations of AdaHOP's per-slot Hadamard + MXFP4 math. +Reference: 3rdparty/adahop/torchtitan/experiments/kernels/mxfp4/mxfp_linear.py +""" diff --git a/alto/modifiers/lpt/adahop_internals/calibration_hooks.py b/alto/modifiers/lpt/adahop_internals/calibration_hooks.py new file mode 100644 index 00000000..1bff5a75 --- /dev/null +++ b/alto/modifiers/lpt/adahop_internals/calibration_hooks.py @@ -0,0 +1,102 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Pure-Python hook factories + JSON I/O for ``AdaHOPModifier``. + +Split out of ``adahop.py`` so it can be unit-tested without dragging in +ALTO's triton-backed kernel imports. +""" + +import json +from pathlib import Path +from typing import Any, Callable, Dict + + +def make_forward_callback(modifier: Any, fqn: str, detect: Callable) -> Callable: + """Build a closure that writes ``(x, w)`` patterns for ``fqn`` into the + modifier's current per-step bucket.""" + + def _cb(x, w) -> None: + if not modifier._per_step_patterns: + return + per_step = modifier._per_step_patterns[-1] + per_layer = per_step.setdefault(fqn, {}) + per_layer["x"] = detect(x) + per_layer["w"] = detect(w) + + return _cb + + +def make_forward_pre_hook(modifier: Any, fqn: str, detect: Callable) -> Callable: + """Forward pre-hook on the nn.Linear module. Captures (x, weight._data) directly + from the module — avoids the tensor-subclass dispatch path entirely, which is + fragile under FSDP because all the rewrap paths (__torch_dispatch__, + fsdp_post_all_gather, and any intervening detach/view) drop wrapper instance + state. By the time __torch_function__("linear") fires the canonical wrapper + is gone; the module reference is stable across all of that.""" + + def _pre(module, args): + if not modifier._per_step_patterns: + return + if not args: + return + x = args[0] + weight = module.weight + # Under FSDP weight is a DTensor wrapping our subclass; under no-FSDP + # it's the subclass directly. Either way ._data lives on the subclass. + w = weight.data + # Unwrap DTensor if present. + try: + from torch.distributed.tensor import DTensor + if isinstance(w, DTensor): + w = w._local_tensor + except Exception: + pass + w_data = getattr(w, "_data", w) + per_step = modifier._per_step_patterns[-1] + per_layer = per_step.setdefault(fqn, {}) + per_layer["x"] = detect(x.detach()) + per_layer["w"] = detect(w_data.detach()) + + return _pre + + +def make_backward_hook(modifier: Any, fqn: str, detect: Callable) -> Callable: + """Build a backward hook that writes ``grad_output`` pattern for ``fqn``.""" + + def _hook(_module, _grad_input, grad_output) -> None: + if not modifier._per_step_patterns: + return + go = grad_output[0] if isinstance(grad_output, (list, tuple)) else grad_output + if go is None: + return + per_step = modifier._per_step_patterns[-1] + per_layer = per_step.setdefault(fqn, {}) + per_layer["grad_output"] = detect(go.detach()) + + return _hook + + +def load_modes_from_json(path: str) -> Dict[str, Dict[str, str]]: + """Load per-layer modes from JSON, tolerating both schemas: + + * Bare ``{fqn: {slot: mode}}`` — what users hand-author or extract. + * Wrapped ``{"aggregated_patterns": {...}, "per_layer_modes": {...}}`` + — what :func:`write_modes_json` produces. + """ + with open(path) as f: + blob = json.load(f) + if isinstance(blob, dict) and "per_layer_modes" in blob and isinstance(blob["per_layer_modes"], dict): + return blob["per_layer_modes"] + return blob + + +def write_modes_json( + path: str, + aggregated: Dict[str, Dict[str, str]], + modes: Dict[str, Dict[str, str]], +) -> None: + Path(path).parent.mkdir(parents=True, exist_ok=True) + blob = {"aggregated_patterns": aggregated, "per_layer_modes": modes} + with open(path, "w") as f: + json.dump(blob, f, indent=2, sort_keys=True) diff --git a/alto/modifiers/lpt/adahop_internals/calibration_state.py b/alto/modifiers/lpt/adahop_internals/calibration_state.py new file mode 100644 index 00000000..fd6aa9dc --- /dev/null +++ b/alto/modifiers/lpt/adahop_internals/calibration_state.py @@ -0,0 +1,112 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Checkpointable AdaHOP calibration state. + +Strategy-2 port (see /home/alirezak/han_branch/porting_notes.txt). Mirrors the +upstream adahop design +(3rdparty/adahop/torchtitan/experiments/kernels/mxfp4/transform_config.py: +``CalibrationStateManager`` + global state dict), adapted to ALTO. + +The per-layer transform modes and the "calibration completed" flag live in a +module-level dict here -- NOT on any tensor subclass. ``CalibrationStateManager`` +implements the PyTorch DCP ``Stateful`` interface so this state is saved and +restored alongside model/optimizer state. The bytes are produced with +``pickle.dumps`` (same approach upstream uses) so DCP treats the value as an +opaque non-tensor blob. + +Lifecycle: + * Fresh run: state starts empty (completed=False). The AdaHOPModifier observes + patterns for ``calibration_steps`` and then calls :func:`set_calibration_result` + which records the modes and flips completed=True. + * Resume: CheckpointManager.load() calls ``load_state_dict`` which repopulates + the module-level dict BEFORE the modifier inspects it on the first training + step, so the modifier applies the stored modes and skips re-calibration. +""" + +import pickle +from typing import Any, Dict + +from torch.distributed.checkpoint.stateful import Stateful + +__all__ = [ + "CalibrationStateManager", + "get_adahop_calibration_state", + "is_calibration_completed", + "get_calibration_modes", + "set_calibration_result", + "reset_calibration_state", +] + + +# Module-level calibration state. Plain, fully-picklable Python objects only. +_STATE: Dict[str, Any] = { + "completed": False, # bool: has Phase B (mode selection) finished? + "step_idx": 0, # int: calibration steps observed so far + "modes_by_fqn": {}, # {fqn: {forward_y, backward_gx, backward_gw}} +} + + +def get_calibration_state_dict() -> Dict[str, Any]: + """Return the live module-level calibration state dict (mutable).""" + return _STATE + + +def is_calibration_completed() -> bool: + return bool(_STATE.get("completed", False)) + + +def get_calibration_modes() -> Dict[str, Dict[str, str]]: + return _STATE.get("modes_by_fqn", {}) + + +def set_calibration_result(modes_by_fqn: Dict[str, Dict[str, str]], step_idx: int) -> None: + """Record the aggregated per-layer modes and mark calibration complete.""" + _STATE["modes_by_fqn"] = dict(modes_by_fqn) + _STATE["step_idx"] = int(step_idx) + _STATE["completed"] = True + + +def set_step_idx(step_idx: int) -> None: + _STATE["step_idx"] = int(step_idx) + + +def reset_calibration_state() -> None: + """Reset to the fresh-run defaults (used by tests).""" + _STATE["completed"] = False + _STATE["step_idx"] = 0 + _STATE["modes_by_fqn"] = {} + + +class CalibrationStateManager(Stateful): + """DCP ``Stateful`` wrapper around the module-level calibration state. + + Registered into ``CheckpointManager.states`` (from ``alto/train.py``) under + the key ``"adahop_calibration"`` so calibration results round-trip with the + checkpoint. The payload is pickled to keep DCP handling simple (opaque bytes, + same pattern as the upstream reference). + """ + + _STATE_KEY = "adahop_calibration_data" + + def state_dict(self) -> Dict[str, Any]: + return {self._STATE_KEY: pickle.dumps(_STATE)} + + def load_state_dict(self, state_dict: Dict[str, Any]) -> None: + if not state_dict or self._STATE_KEY not in state_dict: + return + loaded = pickle.loads(state_dict[self._STATE_KEY]) + if isinstance(loaded, dict): + _STATE.update(loaded) + + +# Process-wide singleton so the modifier and the Trainer (alto/train.py) share +# one instance/state. +_singleton: CalibrationStateManager | None = None + + +def get_adahop_calibration_state() -> CalibrationStateManager: + global _singleton + if _singleton is None: + _singleton = CalibrationStateManager() + return _singleton diff --git a/alto/modifiers/lpt/adahop_internals/mxfp4_linear_function.py b/alto/modifiers/lpt/adahop_internals/mxfp4_linear_function.py new file mode 100644 index 00000000..b8dae603 --- /dev/null +++ b/alto/modifiers/lpt/adahop_internals/mxfp4_linear_function.py @@ -0,0 +1,553 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""ALTO-side MXFP4 + per-slot Hadamard autograd Function. + +Port of AdaHOP's ``MXFP4LinearFunction`` adapted for ALTO's tensor-subclass +wrapper design: modes are passed as ``apply()`` arguments instead of looked +up from a global FQN-keyed registry. + +Reference (AdaHOP, BSD-3): + 3rdparty/adahop/torchtitan/experiments/kernels/mxfp4/mxfp_linear.py:206-635 + +Supported modes per slot: + * ``none`` + * ``hadamard`` / ``outer_hadamard`` + * ``inner_outlier_extract_left`` / ``inner_outlier_extract_right`` + * ``full_precision`` + +Notes vs AdaHOP's version: + * ``layer_name`` removed; modes are explicit ``apply()`` args. + * Calibration/tensor-capture globals removed; calibration runs via the + wrapper's ``__torch_function__`` callback installed by the modifier. + * ``hadamard`` mode uses AdaHOP's fused ``iht_quantization`` kernel + (loaded via the bridge); when SR is requested it composes through. + * Outlier-extract paths use ``inner_outlier_extract_{left,right}_cdna4`` + from AdaHOP's submodule — those kernels self-gate cdna3 (gfx942) via + PyTorch fallbacks (see outlier_extract.py:165, :323). + * Forward ``inner_outlier_extract_left`` saves the BF16 weight to + ``ctx`` for backward (the kernel takes B as BF16 internally). + * Backward ``inner_outlier_extract_left`` for ``backward_gw`` falls + back to plain ``hadamard`` semantics (matches AdaHOP's note at + mxfp_linear.py:560-575). +""" + +from typing import Optional, Tuple + +import torch + +from alto.kernels.fp4.mxfp4.mxfp_quantization import is_cdna4 +from alto.kernels.fp4.fp4_common.tensor_wrappers import unwrap_weight_wrapper +from .transform_mode import TransformMode, assert_mode_supported + +HadamardTransformType = "HadamardTransform" # type-hint placeholder; avoid hard import + + +def _blockwise_mxfp4_gemm_or_dequant( + a_mxfp4: torch.Tensor, + a_scale: torch.Tensor, + b_mxfp4: torch.Tensor, + b_scale: torch.Tensor, + *, + output_dtype: torch.dtype, + trans_a: bool = False, + trans_b: bool = False, + a_axis: int = -1, + b_axis: int = -1, +) -> torch.Tensor: + """Triton MXFP4 GEMM on cdna4, dequant-then-bf16-matmul fallback on cdna3. + Mirrors ``alto/kernels/fp4/mxfp4/mxfp_linear.py:321-402``. + + ``a_axis`` / ``b_axis`` give the axis each operand was PACKED along by + ``convert_to_mxfp4`` (== the GEMM contraction axis K). The cdna3 dequant must + unpack along that same axis; the cdna4 kernel infers it from the scale layout. + Defaults (-1) are correct for the forward; the backward GEMMs pack the weight + / activation along axis 0, so those operands need axis=0. Mismatching this + leaves the packed dim half-size, e.g. ``(8192x32) @ (16x5760)``.""" + if is_cdna4(): + return torch.ops.torchtitan.blockwise_mxfp4_gemm( + a_mxfp4, + a_scale, + b_mxfp4, + b_scale, + trans_a=trans_a, + trans_b=trans_b, + output_dtype=output_dtype, + ) + a_dq = torch.ops.torchtitan.convert_from_mxfp4( + a_mxfp4, a_scale, output_dtype, axis=a_axis, is_2d_block=False, + ) + b_dq = torch.ops.torchtitan.convert_from_mxfp4( + b_mxfp4, b_scale, output_dtype, axis=b_axis, is_2d_block=False, + ) + if trans_a: + a_dq = a_dq.T + if trans_b: + b_dq = b_dq.T + return a_dq @ b_dq + + +@torch.compiler.allow_in_graph +class MXFP4AdaHOPLinearFunction(torch.autograd.Function): + """MXFP4 linear with per-slot Hadamard / outlier-extract, modes baked in. + + Slots: + * ``forward_y`` — ``y = x @ wᵀ`` + * ``backward_gx`` — ``grad_inputs = grad_output @ w`` + * ``backward_gw`` — ``grad_weights = grad_outputᵀ @ x`` + """ + + @staticmethod + def forward( + ctx, + x: torch.Tensor, + weight: torch.Tensor, + use_sr_grad: bool, + hadamard_transform: Optional["HadamardTransformType"], + forward_y_mode: TransformMode, + backward_gx_mode: TransformMode, + backward_gw_mode: TransformMode, + ) -> torch.Tensor: + assert_mode_supported(forward_y_mode, "forward_y") + assert_mode_supported(backward_gx_mode, "backward_gx") + assert_mode_supported(backward_gw_mode, "backward_gw") + + # The wrapper (grad-tracked) is passed into apply() so the weight gradient + # reaches the Parameter; unwrap to the plain payload HERE, inside forward + # (past the autograd boundary), for the quant math. Mirrors the plain + # MXFP4 path (mxfp_linear.py). Safe: forward runs in no-grad, so the graph + # input recorded by apply() is still the wrapper. + weight = unwrap_weight_wrapper(weight) + + original_shape = x.shape + original_dtype = x.dtype + x = x.reshape(-1, original_shape[-1]) + + # ---- forward: y = x @ wᵀ ------------------------------------------ + y = _forward_y( + x=x, + weight=weight, + mode=forward_y_mode, + hadamard_transform=hadamard_transform, + original_dtype=original_dtype, + ) + + # ---- prep operands for the two backward GEMMs --------------------- + w_bx, w_scale_bx, w_outlier_compact_bx, w_outlier_indices_bx = _prep_w_for_gx( + weight=weight, + mode=backward_gx_mode, + hadamard_transform=hadamard_transform, + ) + x_bw, x_scale_bw, x_outlier_compact_bw, x_outlier_indices_bw = _prep_x_for_gw( + x=x, + mode=backward_gw_mode, + hadamard_transform=hadamard_transform, + use_sr_grad=use_sr_grad, + ) + + ctx.save_for_backward( + x_bw, x_scale_bw, w_bx, w_scale_bx, + x_outlier_compact_bw, x_outlier_indices_bw, + w_outlier_compact_bx, w_outlier_indices_bx, + ) + ctx.original_dtype = original_dtype + ctx.use_sr_grad = use_sr_grad + ctx.hadamard_transform = hadamard_transform + ctx.forward_y_mode = forward_y_mode + ctx.backward_gx_mode = backward_gx_mode + ctx.backward_gw_mode = backward_gw_mode + + return y.view(*original_shape[:-1], -1) + + @staticmethod + def backward(ctx, grad_output): + original_shape = grad_output.shape + grad_output = grad_output.reshape(-1, original_shape[-1]) + + (x_bw, x_scale_bw, w_bx, w_scale_bx, + x_outlier_compact_bw, x_outlier_indices_bw, + w_outlier_compact_bx, w_outlier_indices_bx) = ctx.saved_tensors + + grad_inputs = _backward_gx( + grad_output=grad_output, + w=w_bx, + w_scale=w_scale_bx, + w_outlier_compact=w_outlier_compact_bx, + w_outlier_indices=w_outlier_indices_bx, + mode=ctx.backward_gx_mode, + hadamard_transform=ctx.hadamard_transform, + use_sr_grad=ctx.use_sr_grad, + original_dtype=ctx.original_dtype, + ) + + grad_weights = _backward_gw( + grad_output=grad_output, + x=x_bw, + x_scale=x_scale_bw, + x_outlier_compact=x_outlier_compact_bw, + x_outlier_indices=x_outlier_indices_bw, + mode=ctx.backward_gw_mode, + hadamard_transform=ctx.hadamard_transform, + use_sr_grad=ctx.use_sr_grad, + original_dtype=ctx.original_dtype, + ) + + # forward has 7 inputs; only the first two get grads. + return ( + grad_inputs.view(*original_shape[:-1], -1), + grad_weights, + None, None, None, None, None, + ) + + +# --------------------------------------------------------------------------- +# Forward y +# --------------------------------------------------------------------------- + +def _forward_y( + *, + x: torch.Tensor, + weight: torch.Tensor, + mode: TransformMode, + hadamard_transform, + original_dtype: torch.dtype, +) -> torch.Tensor: + """Compute ``y = x @ wᵀ`` under ``forward_y_mode``.""" + if mode == "full_precision": + return x @ weight.T + + if mode == "inner_outlier_extract_left": + # Extract row outliers from x, quantize the clean part, run the + # AdaHOP outlier-aware GEMM. Weight stays BF16; the kernel applies + # its own HT+quant to it internally. + from alto._adahop_bridge import ( + iht_quantization, + inner_outlier_extract_left_cdna4, + prepare_outlier_clean_row, + OUTLIER_K, + ) + assert hadamard_transform is not None + x_clean, x_outlier_compact, x_outlier_indices = prepare_outlier_clean_row(x, k=OUTLIER_K) + x_clean_mxfp4, x_clean_scale = iht_quantization(x_clean, left_mul=False) + return inner_outlier_extract_left_cdna4( + x_clean_mxfp4, + x_clean_scale, + x_outlier_compact, + x_outlier_indices, + weight, + hadamard_transform, + trans_a=False, + trans_b=True, + k=OUTLIER_K, + original_dtype=original_dtype, + ) + + if mode == "inner_outlier_extract_right": + # Extract row outliers from weight (== col outliers in wᵀ), + # quantize the clean weight, run the AdaHOP outlier-aware GEMM. + # x stays BF16; the kernel handles it. + from alto._adahop_bridge import ( + iht_quantization, + inner_outlier_extract_right_cdna4, + prepare_outlier_clean_row, + OUTLIER_K, + ) + assert hadamard_transform is not None + w_clean, w_outlier_compact, w_outlier_indices = prepare_outlier_clean_row(weight, k=OUTLIER_K) + w_clean_mxfp4, w_clean_scale = iht_quantization(w_clean, left_mul=False) + return inner_outlier_extract_right_cdna4( + x, + w_clean_mxfp4, + w_clean_scale, + w_outlier_compact, + w_outlier_indices, + hadamard_transform, + trans_a=False, + trans_b=True, + k=OUTLIER_K, + original_dtype=original_dtype, + ) + + # hadamard / outer_hadamard / none — two-step convert+GEMM. + if mode == "hadamard": + assert hadamard_transform is not None + x_for_y = hadamard_transform(x, left_mul=False) + w_for_y = hadamard_transform(weight, left_mul=False) + elif mode == "outer_hadamard": + assert hadamard_transform is not None + x_for_y = hadamard_transform(x, left_mul=True) + w_for_y = hadamard_transform(weight, left_mul=True) + else: # "none" + x_for_y = x + w_for_y = weight + + x_mxfp4, x_scale = torch.ops.torchtitan.convert_to_mxfp4( + x_for_y, axis=-1, is_2d_block=False, + ) + w_mxfp4, w_scale = torch.ops.torchtitan.convert_to_mxfp4( + w_for_y, axis=-1, is_2d_block=False, + ) + y = _blockwise_mxfp4_gemm_or_dequant( + x_mxfp4, x_scale, w_mxfp4, w_scale, + trans_b=True, output_dtype=original_dtype, + ) + if mode == "outer_hadamard": + y = hadamard_transform(hadamard_transform(y, left_mul=True)) + return y + + +# --------------------------------------------------------------------------- +# Backward operand preparation (runs in forward, saves into ctx) +# --------------------------------------------------------------------------- + +def _prep_w_for_gx( + *, + weight: torch.Tensor, + mode: TransformMode, + hadamard_transform, +) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: + """Prepare weight for ``backward_gx`` (``grad_output @ w``). + + Returns ``(w_data, w_scale, w_outlier_compact, w_outlier_indices)``. + Last three are ``None`` unless mode needs them. + """ + if mode == "full_precision": + return weight, None, None, None + + if mode == "inner_outlier_extract_left": + # Outliers will be extracted from grad_output at backward time; + # weight is saved as BF16 (kernel handles HT+quant internally). + return weight, None, None, None + + if mode == "inner_outlier_extract_right": + # Extract column outliers from weight now; quantize the clean part. + from alto._adahop_bridge import ( + iht_quantization, + prepare_outlier_clean_column, + OUTLIER_K, + ) + assert hadamard_transform is not None + w_clean, w_outlier_compact, w_outlier_indices = prepare_outlier_clean_column(weight, k=OUTLIER_K) + w_mxfp4, w_scale = iht_quantization(w_clean, left_mul=True) + return w_mxfp4, w_scale, w_outlier_compact, w_outlier_indices + + # hadamard / outer_hadamard / none — single-tensor convert path. + if mode == "hadamard": + from alto._adahop_bridge import iht_quantization + assert hadamard_transform is not None + w_mxfp4, w_scale = iht_quantization(weight, left_mul=True) + return w_mxfp4, w_scale, None, None + if mode == "outer_hadamard": + assert hadamard_transform is not None + prepared = hadamard_transform(weight) + w_mxfp4, w_scale = torch.ops.torchtitan.convert_to_mxfp4(prepared, axis=0, is_2d_block=False) + return w_mxfp4, w_scale, None, None + # "none" + w_mxfp4, w_scale = torch.ops.torchtitan.convert_to_mxfp4(weight, axis=0, is_2d_block=False) + return w_mxfp4, w_scale, None, None + + +def _prep_x_for_gw( + *, + x: torch.Tensor, + mode: TransformMode, + hadamard_transform, + use_sr_grad: bool, +) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: + """Prepare x for ``backward_gw`` (``grad_outputᵀ @ x``). + + Returns ``(x_data, x_scale, x_outlier_compact, x_outlier_indices)``. + Last three are ``None`` unless mode needs them. + + Note: ``inner_outlier_extract_left`` for ``backward_gw`` is identical + to plain ``hadamard`` (per AdaHOP mxfp_linear.py:560-575) — left OE + can't compose with backward_gw because x needs to be saved as MXFP4 + for memory, conflicting with the BF16 requirement of the left kernel. + """ + if mode == "full_precision": + return x, None, None, None + + if mode in ("hadamard", "inner_outlier_extract_left"): + from alto._adahop_bridge import iht_quantization + assert hadamard_transform is not None + x_mxfp4, x_scale = iht_quantization(x, left_mul=True) + return x_mxfp4, x_scale, None, None + + if mode == "inner_outlier_extract_right": + from alto._adahop_bridge import ( + iht_quantization, + prepare_outlier_clean_column, + OUTLIER_K, + ) + assert hadamard_transform is not None + x_clean, x_outlier_compact, x_outlier_indices = prepare_outlier_clean_column(x, k=OUTLIER_K) + x_mxfp4, x_scale = iht_quantization(x_clean, left_mul=True) + return x_mxfp4, x_scale, x_outlier_compact, x_outlier_indices + + if mode == "outer_hadamard": + assert hadamard_transform is not None + prepared = hadamard_transform(x) + x_mxfp4, x_scale = torch.ops.torchtitan.convert_to_mxfp4(prepared, axis=0, is_2d_block=False) + return x_mxfp4, x_scale, None, None + + # "none" + x_mxfp4, x_scale = torch.ops.torchtitan.convert_to_mxfp4(x, axis=0, is_2d_block=False) + return x_mxfp4, x_scale, None, None + + +# --------------------------------------------------------------------------- +# Backward compute +# --------------------------------------------------------------------------- + +def _backward_gx( + *, + grad_output: torch.Tensor, + w: torch.Tensor, + w_scale: Optional[torch.Tensor], + w_outlier_compact: Optional[torch.Tensor], + w_outlier_indices: Optional[torch.Tensor], + mode: TransformMode, + hadamard_transform, + use_sr_grad: bool, + original_dtype: torch.dtype, +) -> torch.Tensor: + """``grad_inputs = grad_output @ w`` under ``backward_gx_mode``.""" + if mode == "full_precision": + # w is the unquantized weight here. + return grad_output @ w + + if mode == "inner_outlier_extract_left": + # Extract row outliers from grad_output; w is BF16 in `w`. + from alto._adahop_bridge import ( + iht_quantization, + inner_outlier_extract_left_cdna4, + prepare_outlier_clean_row, + OUTLIER_K, + ) + assert hadamard_transform is not None + g_clean, g_outlier_compact, g_outlier_indices = prepare_outlier_clean_row(grad_output, k=OUTLIER_K) + g_clean_mxfp4, g_clean_scale = iht_quantization( + g_clean, left_mul=False, use_sr=use_sr_grad, + ) + return inner_outlier_extract_left_cdna4( + g_clean_mxfp4, + g_clean_scale, + g_outlier_compact, + g_outlier_indices, + w, # BF16 weight + hadamard_transform, + trans_a=False, + trans_b=False, + k=OUTLIER_K, + original_dtype=original_dtype, + ) + + if mode == "inner_outlier_extract_right": + # Pre-extracted column outliers from weight (saved at forward). + from alto._adahop_bridge import ( + inner_outlier_extract_right_cdna4, + OUTLIER_K, + ) + assert hadamard_transform is not None + return inner_outlier_extract_right_cdna4( + grad_output, + w, + w_scale, + w_outlier_compact, + w_outlier_indices, + hadamard_transform, + trans_a=False, + trans_b=False, + k=OUTLIER_K, + original_dtype=original_dtype, + use_sr=use_sr_grad, # grad_output is A -> stochastic-round its quant + ) + + # hadamard / outer_hadamard / none + if mode == "hadamard": + from alto._adahop_bridge import iht_quantization + assert hadamard_transform is not None + g_mxfp4, g_scale = iht_quantization(grad_output, left_mul=False, use_sr=use_sr_grad) + elif mode == "outer_hadamard": + assert hadamard_transform is not None + g_transformed = hadamard_transform(grad_output, left_mul=True) + g_mxfp4, g_scale = torch.ops.torchtitan.convert_to_mxfp4( + g_transformed, axis=-1, use_sr=use_sr_grad, is_2d_block=False, + ) + else: # "none" + g_mxfp4, g_scale = torch.ops.torchtitan.convert_to_mxfp4( + grad_output, axis=-1, use_sr=use_sr_grad, is_2d_block=False, + ) + + grad_inputs = _blockwise_mxfp4_gemm_or_dequant( + g_mxfp4, g_scale, w, w_scale, + output_dtype=original_dtype, + b_axis=0, # weight packed along axis 0 (= contraction axis "out") + ) + if mode == "outer_hadamard": + grad_inputs = hadamard_transform(hadamard_transform(grad_inputs, left_mul=True)) + return grad_inputs + + +def _backward_gw( + *, + grad_output: torch.Tensor, + x: torch.Tensor, + x_scale: Optional[torch.Tensor], + x_outlier_compact: Optional[torch.Tensor], + x_outlier_indices: Optional[torch.Tensor], + mode: TransformMode, + hadamard_transform, + use_sr_grad: bool, + original_dtype: torch.dtype, +) -> torch.Tensor: + """``grad_weights = grad_outputᵀ @ x`` under ``backward_gw_mode``.""" + if mode == "full_precision": + # x is the unquantized input here. + return grad_output.T @ x + + if mode == "inner_outlier_extract_right": + from alto._adahop_bridge import ( + inner_outlier_extract_right_cdna4, + OUTLIER_K, + ) + assert hadamard_transform is not None + return inner_outlier_extract_right_cdna4( + grad_output, + x, + x_scale, + x_outlier_compact, + x_outlier_indices, + hadamard_transform, + trans_a=True, + trans_b=False, + k=OUTLIER_K, + original_dtype=original_dtype, + use_sr=use_sr_grad, # grad_output is A -> stochastic-round its quant + ) + + # hadamard / inner_outlier_extract_left (same path per AdaHOP) / + # outer_hadamard / none + if mode in ("hadamard", "inner_outlier_extract_left"): + from alto._adahop_bridge import iht_quantization + assert hadamard_transform is not None + g_mxfp4, g_scale = iht_quantization(grad_output, left_mul=True, use_sr=use_sr_grad) + elif mode == "outer_hadamard": + assert hadamard_transform is not None + g_transformed = hadamard_transform(grad_output) + g_mxfp4, g_scale = torch.ops.torchtitan.convert_to_mxfp4( + g_transformed, axis=0, use_sr=use_sr_grad, is_2d_block=False, + ) + else: # "none" + g_mxfp4, g_scale = torch.ops.torchtitan.convert_to_mxfp4( + grad_output, axis=0, use_sr=use_sr_grad, is_2d_block=False, + ) + + grad_weights = _blockwise_mxfp4_gemm_or_dequant( + g_mxfp4, g_scale, x, x_scale, + trans_a=True, output_dtype=original_dtype, + a_axis=0, b_axis=0, # grad_output and x both packed along axis 0 (= "M") + ) + if mode == "outer_hadamard": + grad_weights = hadamard_transform(hadamard_transform(grad_weights, left_mul=True)) + return grad_weights diff --git a/alto/modifiers/lpt/adahop_internals/pattern_aggregation.py b/alto/modifiers/lpt/adahop_internals/pattern_aggregation.py new file mode 100644 index 00000000..9c77d2ff --- /dev/null +++ b/alto/modifiers/lpt/adahop_internals/pattern_aggregation.py @@ -0,0 +1,92 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Per-layer outlier-pattern aggregation and pattern-pair → TransformMode mapping. + +Reference (AdaHOP, BSD-3): + 3rdparty/adahop/torchtitan/components/quantization/mx_calibration.py:158-265 (_aggregate_patterns) + 3rdparty/adahop/torchtitan/components/quantization/mx_calibration.py:412-481 (_apply_calibrated_transforms) + +The aggregation logic is verbatim (Counter.most_common per tensor type). The +"T1..T8" pattern-pair construction is identical. The pair→mode lookup is also +identical; the only difference is that we return a dict of per-layer +(forward_y, backward_gx, backward_gw) modes instead of calling AdaHOP's global +``configure_layer_transforms``. The caller (the modifier) feeds the result +into the Phase-B wrapper constructors. +""" + +from collections import Counter +from typing import Dict, List + +from .transform_mode import OutlierPattern, PerSlotModes, TransformMode + + +def _opposite_pattern(pattern: OutlierPattern) -> OutlierPattern: + if pattern == "row": + return "col" + if pattern == "col": + return "row" + return "none" + + +def aggregate_patterns( + per_step_patterns: List[Dict[str, Dict[str, OutlierPattern]]],) -> Dict[str, Dict[str, OutlierPattern]]: + """Reduce per-step pattern observations to one majority pattern per (layer, tensor). + + ``per_step_patterns[i][layer_name][tensor_name]`` is the pattern observed + at step ``i`` for ``tensor_name`` in ``{"x", "w", "grad_output"}``. Missing + entries are skipped; missing tensors default to ``"none"``. + + Returns ``{layer_name: {"x": ..., "w": ..., "grad_output": ...}}``. + """ + all_layers = set() + for step in per_step_patterns: + all_layers.update(step.keys()) + + final: Dict[str, Dict[str, OutlierPattern]] = {} + for layer in sorted(all_layers): + per_tensor: Dict[str, List[OutlierPattern]] = {"x": [], "w": [], "grad_output": []} + for step in per_step_patterns: + if layer not in step: + continue + for t in per_tensor: + if t in step[layer]: + per_tensor[t].append(step[layer][t]) + + final[layer] = { + t: (Counter(per_tensor[t]).most_common(1)[0][0] if per_tensor[t] else "none") for t in per_tensor + } + return final + + +def patterns_to_modes( + aggregated: Dict[str, Dict[str, OutlierPattern]], + layer_transform_config: Dict[str, TransformMode], +) -> Dict[str, PerSlotModes]: + """Map per-layer aggregated patterns to per-slot ``TransformMode`` triples. + + ``layer_transform_config`` is the recipe-supplied pattern-pair → mode lookup, + e.g. ``{"row-row": "hadamard", "col-col": "full_precision", "none-none": "hadamard"}``. + Unknown pairs fall back to ``"none"``. + + T1..T8 follow AdaHOP exactly: + T1 = x_pat, T2 = opposite(w_pat) → forward_y key "T1-T2" + T4 = opposite(grad_output_pat), T5 = x_pat → backward_gw key "T4-T5" + T7 = grad_output_pat, T8 = w_pat → backward_gx key "T7-T8" + """ + out: Dict[str, PerSlotModes] = {} + for layer, pats in aggregated.items(): + x_pat = pats.get("x", "none") + w_pat = pats.get("w", "none") + g_pat = pats.get("grad_output", "none") + + t1, t2 = x_pat, _opposite_pattern(w_pat) + t4, t5 = _opposite_pattern(g_pat), x_pat + t7, t8 = g_pat, w_pat + + out[layer] = { + "forward_y": layer_transform_config.get(f"{t1}-{t2}", "none"), + "backward_gw": layer_transform_config.get(f"{t4}-{t5}", "none"), + "backward_gx": layer_transform_config.get(f"{t7}-{t8}", "none"), + } + return out diff --git a/alto/modifiers/lpt/adahop_internals/transform_mode.py b/alto/modifiers/lpt/adahop_internals/transform_mode.py new file mode 100644 index 00000000..8c3da211 --- /dev/null +++ b/alto/modifiers/lpt/adahop_internals/transform_mode.py @@ -0,0 +1,62 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Per-slot Hadamard transform modes for the AdaHOP-in-ALTO port. + +Mirrors AdaHOP's ``TransformMode`` Literal (see reference below) so JSON +configs produced by AdaHOP's calibration deserialize without translation. + +Reference: 3rdparty/adahop/torchtitan/experiments/kernels/mxfp4/transform_config.py:21 +""" + +from typing import Dict, Literal, Tuple + +TransformMode = Literal[ + "none", + "hadamard", + "outer_hadamard", + "inner_outlier_extract_left", + "inner_outlier_extract_left_col", + "inner_outlier_extract_right", + "outer_outlier_extract_left", + "outer_outlier_extract_right", + "full_precision", +] + +OutlierPattern = Literal["row", "col", "none"] + +VALID_MODES: frozenset = frozenset([ + "none", + "hadamard", + "outer_hadamard", + "inner_outlier_extract_left", + "inner_outlier_extract_left_col", + "inner_outlier_extract_right", + "outer_outlier_extract_left", + "outer_outlier_extract_right", + "full_precision", +]) + +# Modes implemented in the ALTO-side MXFP4LinearFunction port. +# Other VALID_MODES are accepted by the schema but will raise at runtime if +# encountered in a calibration JSON until their kernel paths are ported. +SUPPORTED_MODES: frozenset = frozenset([ + "none", + "hadamard", + "outer_hadamard", + "full_precision", + "inner_outlier_extract_left", + "inner_outlier_extract_right", +]) + + +def assert_mode_supported(mode: str, slot: str) -> None: + if mode not in VALID_MODES: + raise ValueError(f"Unknown TransformMode {mode!r} for slot {slot!r}") + if mode not in SUPPORTED_MODES: + raise NotImplementedError(f"TransformMode {mode!r} (slot {slot!r}) is not yet implemented in ALTO. " + f"Supported: {sorted(SUPPORTED_MODES)}.") + + +PatternPair = Tuple[OutlierPattern, OutlierPattern] +PerSlotModes = Dict[Literal["forward_y", "backward_gx", "backward_gw"], TransformMode] diff --git a/alto/modifiers/lpt/base.py b/alto/modifiers/lpt/base.py index 319d0835..6546546c 100644 --- a/alto/modifiers/lpt/base.py +++ b/alto/modifiers/lpt/base.py @@ -33,11 +33,19 @@ class LowPrecisionTrainingModifier(Modifier): use_2dblock_x: bool = False use_2dblock_w: bool = True use_hadamard: bool = False + hadamard_type: Literal["default", "3rht"] = "default" use_sr_grad: bool = False use_dge: bool = False + full_precision_backward: bool = False + """ + Quantize the forward only; keep the backward (dgrad + wgrad) in bf16 with the + gradient unquantized. Dense MXFP4 linear path only. + """ two_level_scaling: Literal["none", "tensorwise", "blockwise"] = "none" clip_mode: Literal["none", "static", "dynamic"] = "none" - + blockscale_selection: Literal["default", "midmax-legacy", "uos", "uos6"] = "default" + + lora_rank: int = 0 """ Lora rank for the decomposed linear layer. @@ -81,7 +89,7 @@ def validate_targets(cls, value: str | list[str]) -> list[str]: @field_validator("scheme", mode="before") def validate_scheme(cls, value: str | dict[str, str | list[str]]) -> str | dict[str, list[str]]: - if isinstance(value, str) and value not in ["mxfp4", "mxfp8_e4m3", "mxfp8_e5m2", "nvfp4", "amdfp4"]: + if isinstance(value, str) and value not in ["mxfp4", "mxfp4_adahop", "mxfp8_e4m3", "mxfp8_e5m2", "nvfp4", "amdfp4"]: raise ValueError(f"Unsupported training op scheme: {value}") if isinstance(value, dict): @@ -93,6 +101,19 @@ def validate_scheme(cls, value: str | dict[str, str | list[str]]) -> str | dict[ return value + @model_validator(mode="after") + def validate_full_precision_backward(self): + if self.full_precision_backward: + scheme_names = self.scheme if isinstance(self.scheme, dict) else {self.scheme: None} + for scheme_name in scheme_names: + if scheme_name not in ("mxfp4",): + raise ValueError( + f"full_precision_backward is only supported for scheme 'mxfp4', got '{scheme_name}'") + if self.use_dge: + raise ValueError("full_precision_backward is incompatible with use_dge " + "(DGE is a gradient-quantization estimator; the backward is not quantized here)") + return self + @model_validator(mode="after") def validate_lora_rank_alignment(self): if self.lora_rank <= 0: @@ -107,9 +128,7 @@ def validate_lora_rank_alignment(self): ) elif scheme_name in ("mxfp4", "mxfp8_e4m3", "mxfp8_e5m2"): if self.lora_rank % 32 != 0: - raise ValueError( - f"lora_rank must be divisible by 32 for {scheme_name}, got {self.lora_rank}" - ) + raise ValueError(f"lora_rank must be divisible by 32 for {scheme_name}, got {self.lora_rank}") return self @field_validator("deosc_step", mode="after") @@ -149,21 +168,41 @@ def resolved_config(self) -> dict[TrainingOpConfig, list[str]]: self._resolved_config = {} for scheme_name, targets in self.scheme.items(): + # "mxfp4_adahop" reuses the MXFP4 kernel path; the only + # difference is which wrapper class is used at swap time + # (handled in on_convert). At the TrainingOpConfig level it + # is plain mxfp4. + precision = "mxfp4" if scheme_name == "mxfp4_adahop" else scheme_name scheme_obj = TrainingOpConfig( - precision=scheme_name, + precision=precision, use_2dblock_x=self.use_2dblock_x, use_2dblock_w=self.use_2dblock_w, use_hadamard=self.use_hadamard, use_sr_grad=self.use_sr_grad, use_dge=self.use_dge, + full_precision_backward=self.full_precision_backward, two_level_scaling=self.two_level_scaling, clip_mode=self.clip_mode, + blockscale_selection=self.blockscale_selection, ) + # Tag the underlying scheme so on_convert can pick the right + # wrapper class. Stored on the dict key via a sibling attribute + # of the modifier rather than mutating TrainingOpConfig. self._resolved_config[scheme_obj] = targets + self._scheme_tag = getattr(self, "_scheme_tag", {}) + self._scheme_tag[scheme_obj] = scheme_name return self._resolved_config def on_convert(self, model: Module, **kwargs) -> bool: + if self.use_hadamard and self.hadamard_type != "default": + from alto.kernels.hadamard_transform import HadamardFactory + HadamardFactory.configure(transform_type=self.hadamard_type) for scheme_obj, targets in self.resolved_config.items(): + tensor_cls = self._wrapper_cls_for_scheme(scheme_obj) + scheme_name = getattr(self, "_scheme_tag", {}).get(scheme_obj, scheme_obj.precision) + wrapper_label = tensor_cls.__name__ if tensor_cls is not None else "default-for-precision" + logger.info(f"LowPrecisionTrainingModifier: scheme={scheme_name}, wrapper_cls={wrapper_label}, " + f"targets={targets}, ignore={self.ignore}") for name, module in match_named_modules(model, targets, self.ignore): if isinstance(module, BaseAttention): assert module.attn_backend == "sdpa", "Only SDPA attention is supported for now." @@ -171,20 +210,33 @@ def on_convert(self, model: Module, **kwargs) -> bool: elif isinstance(module, torch.nn.Linear): if self.lora_rank > 0: module = DecomposedLinear.from_linear(module, lora_rank=self.lora_rank) - swap_params(module, config=scheme_obj, target_parameter_name="weight") - swap_params(module, config=scheme_obj, target_parameter_name="u") - swap_params(module, config=scheme_obj, target_parameter_name="v") + swap_params(module, config=scheme_obj, target_parameter_name="weight", tensor_cls=tensor_cls) + swap_params(module, config=scheme_obj, target_parameter_name="u", tensor_cls=tensor_cls) + swap_params(module, config=scheme_obj, target_parameter_name="v", tensor_cls=tensor_cls) model.set_submodule(name, module, strict=True) else: - swap_params(module, config=scheme_obj, module_name=name) + swap_params(module, config=scheme_obj, module_name=name, tensor_cls=tensor_cls) elif module.__class__.__name__.endswith("GroupedExperts"): - swap_params(module, config=scheme_obj, module_name=name) + swap_params(module, config=scheme_obj, module_name=name, tensor_cls=tensor_cls) else: raise ValueError(f"Unsupported module type: {type(module)}") logger.info(f"LowPrecisionTrainingModifier converted model: {model}") return True + def _wrapper_cls_for_scheme(self, scheme_obj): + """Return the wrapper tensor class for ``scheme_obj``. ``None`` falls + back to ``swap_params``' default (looked up from the config precision).""" + scheme_name = getattr(self, "_scheme_tag", {}).get(scheme_obj) + if scheme_name == "mxfp4_adahop": + # Single-wrapper design: wrap with the AdaHOP wrapper from convert + # time (modes default "none" == plain MXFP4 during calibration). The + # AdaHOPModifier flips the modes in place at Phase B so they take + # effect on the FSDP-owned tensor. + from alto.kernels.dispatch.adahop_tensor import MXFP4AdaHOPWrapper + return MXFP4AdaHOPWrapper + return None + def on_initialize(self, model_parts: list[Module], **kwargs) -> bool: for model_part in model_parts: for child in model_part.modules(): @@ -193,6 +245,10 @@ def on_initialize(self, model_parts: list[Module], **kwargs) -> bool: return True def on_pre_step(self, model_parts: list[Module], **kwargs) -> bool: + if self.use_hadamard: + from alto.kernels.hadamard_transform import HadamardFactory + HadamardFactory.refresh() + trainer = kwargs.get("trainer", None) if self.deosc_step > 0: diff --git a/alto/modifiers/lpt/grad_clip.py b/alto/modifiers/lpt/grad_clip.py new file mode 100644 index 00000000..1d16140a --- /dev/null +++ b/alto/modifiers/lpt/grad_clip.py @@ -0,0 +1,90 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT + +from typing import Optional + +from pydantic import Field, PrivateAttr +from torch import nn +from compressed_tensors.utils import match_named_modules + +from alto.modifiers import Modifier +from alto.kernels.fp4.fp4_common.grad_clip_config import GradClipConfig +from alto.kernels.fp4.fp4_common import grad_clip_registry + +__all__ = ["GradientClippingModifier"] + + +class GradientClippingModifier(Modifier): + """Injects per-layer gradient clipping into MXFP4LinearFunction.backward. + + Must be listed **after** LowPrecisionTrainingModifier in the recipe YAML so + that swap_params has already run when on_initialize sets module_id on each + wrapped weight tensor. + + Two clip sites per MXFP4 nn.Linear backward: + - grad_output before convert_to_mxfp4 (controls quantizer input scale) + - grad_weights after the wgrad GEMM (limits per-layer optimizer step size) + """ + + targets: list[str] = Field(default_factory=lambda: ["Linear"]) + ignore: list[str] = Field(default_factory=lambda: ["output", "re:.*\\.router\\.gate"]) + + clip_grad_output: bool = True + grad_output_max_norm: Optional[float] = None + grad_output_clip_value: Optional[float] = None + + clip_grad_weight: bool = True + grad_weight_max_norm: Optional[float] = None + grad_weight_clip_value: Optional[float] = None + + _registered_ids: list[int] = PrivateAttr(default_factory=list) + _hook_handles: list = PrivateAttr(default_factory=list) + + def on_convert(self, model, **kwargs) -> bool: + return True + + def on_initialize(self, model_parts: list[nn.Module], **kwargs) -> bool: + cfg = GradClipConfig( + clip_grad_output=self.clip_grad_output, + grad_output_max_norm=self.grad_output_max_norm, + grad_output_clip_value=self.grad_output_clip_value, + clip_grad_weight=self.clip_grad_weight, + grad_weight_max_norm=self.grad_weight_max_norm, + grad_weight_clip_value=self.grad_weight_clip_value, + ) + for model_part in model_parts: + for _fqn, module in match_named_modules(model_part, self.targets, self.ignore): + mid = id(module) + grad_clip_registry.register(mid, cfg) + self._registered_ids.append(mid) + + # FSDP's fsdp_post_all_gather rebuilds the weight tensor as a + # new object each forward pass, dropping any instance attribute + # set at initialize time. Register a forward_pre_hook instead: + # it fires after the all-gather so the freshly reconstructed + # weight tensor is already in place, and we re-stamp module_id + # before __torch_function__ is called. + def _stamp_module_id(mod, _args, _mid=mid): + w = mod.weight + if hasattr(w, "module_id"): + w.module_id = _mid + + handle = module.register_forward_pre_hook(_stamp_module_id, prepend=True) + self._hook_handles.append(handle) + return True + + def on_finalize(self, model_parts: list[nn.Module], **kwargs) -> bool: + for handle in self._hook_handles: + handle.remove() + self._hook_handles.clear() + for mid in self._registered_ids: + grad_clip_registry.deregister(mid) + self._registered_ids.clear() + return True + + def on_pre_step(self, model_parts: list[nn.Module], **kwargs) -> bool: + return True + + def on_post_step(self, model_parts: list[nn.Module], **kwargs) -> bool: + return True diff --git a/alto/nn/decomposed_linear.py b/alto/nn/decomposed_linear.py index 54e721a2..83664ec3 100644 --- a/alto/nn/decomposed_linear.py +++ b/alto/nn/decomposed_linear.py @@ -31,13 +31,9 @@ def from_linear(cls, linear: nn.Linear, lora_rank: int = 32): new_layer = cls(linear.in_features, linear.out_features, linear.bias is not None, lora_rank) new_layer.weight = linear.weight new_layer.bias = linear.bias - device = linear.weight.device - dtype = linear.weight.dtype - new_layer.u.data = new_layer.u.data.to(device=device, dtype=dtype) - new_layer.v.data = new_layer.v.data.to(device=device, dtype=dtype) - new_layer.sigma.data = new_layer.sigma.data.to(device=device, dtype=dtype) return new_layer - + + def init_lora_weights(self, init_std: float = 0.02): nn.init.normal_(self.u, mean=0.0, std=init_std) nn.init.zeros_(self.v) diff --git a/alto/train.py b/alto/train.py index cec58aec..186705a5 100644 --- a/alto/train.py +++ b/alto/train.py @@ -4,6 +4,7 @@ from typing import Iterable, Any from contextlib import contextmanager +import os import time import torch from torchtitan.components.loss import IGNORE_INDEX @@ -12,6 +13,7 @@ from torchtitan.experiments.forge.example_train import Trainer as ForgeTrainer, main as forge_main from torchtitan.components.metrics import MetricsProcessor from alto.components.converter import ModelOptConverter +from alto.components.optimizer import DeOscillationConfig, enable_de_oscillation from torchtitan.tools.logging import logger @@ -84,6 +86,34 @@ class Trainer(ForgeTrainer): def __init__(self, config: TitanTrainer.Config): super().__init__(config) + # AdaHOP (Strategy-2 port): if an AdaHOPModifier is in the recipe, + # register its checkpointable calibration state with the checkpointer so + # per-layer transform modes are saved/restored alongside model/optimizer + # state. CheckpointManager.states is a live, mutable dict iterated by + # save()/load(); injecting here (after super().__init__, before + # train()->checkpointer.load) keeps the whole change inside alto/. + try: + from alto.modifiers.lpt.adahop import AdaHOPModifier + from alto.modifiers.lpt.adahop_internals.calibration_state import ( + get_adahop_calibration_state, + ) + has_adahop = any( + isinstance(m, AdaHOPModifier) + for conv in self.model_converters.converters + if isinstance(conv, ModelOptConverter) + for m in getattr(conv, "recipe", None).modifiers + ) if not self.model_converters.is_empty() else False + if has_adahop and getattr(self, "checkpointer", None) is not None: + self.checkpointer.states["adahop_calibration"] = get_adahop_calibration_state() + logger.info("[AdaHOP] Registered CalibrationStateManager with checkpointer " + "(states key='adahop_calibration').") + except Exception as e: # never let this break trainer construction + logger.warning(f"[AdaHOP] Could not register calibration state with checkpointer: {e}") + + self.checkpointer.states["dataloader"] = self.dataloader + + self.ntokens_seen = 0 + self.training_mode = True self.enable_data_cache = False @@ -109,6 +139,22 @@ def __init__(self, config: TitanTrainer.Config): logger.info("data replay buffer disabled") self.enable_data_cache = False + def state_dict(self) -> dict[str, Any]: + sd = super().state_dict() + sd["ntokens_seen"] = self.ntokens_seen + return sd + + def load_state_dict(self, state_dict: dict[str, Any]): + super().load_state_dict(state_dict) + self.ntokens_seen = state_dict.get("ntokens_seen", 0) + + def batch_generator( + self, data_iterable: Iterable[tuple[dict[str, torch.Tensor], torch.Tensor]] + ) -> Iterable[tuple[dict[str, torch.Tensor], torch.Tensor]]: + for input_dict, labels in super().batch_generator(data_iterable): + self.ntokens_seen += labels.numel() + yield input_dict, labels + def cache_input(self, microbatches: list[tuple[dict[str, torch.Tensor], torch.Tensor]]): if self.enable_data_cache: self._input_cache = microbatches @@ -190,6 +236,18 @@ def train_step( data_iterator: Iterable[tuple[dict[str, torch.Tensor], torch.Tensor]], ): if self.training_mode: + # FIXME: This is a hack to enable de-oscillation at a specific step. + deosc_step = int(os.environ.get("DEOSC_STEP", "0")) + ratio_threshold = float(os.environ.get("DEOSC_RATIO", "8.0")) + if deosc_step > 0 and self.step == deosc_step: + deosc_config = DeOscillationConfig( + enable=True, + period=200, + ratio_threshold=ratio_threshold, + log_freq=1, + ) + enable_de_oscillation(self.optimizers, deosc_config) + return super().train_step(data_iterator) # Keep these variables local to shorten the code as these are diff --git a/plotting/baselines.toml b/plotting/baselines.toml new file mode 100644 index 00000000..401b0284 --- /dev/null +++ b/plotting/baselines.toml @@ -0,0 +1,20 @@ +# Example config for plot_val_loss.py +# python3 plot_val_loss.py -c runs.example.toml +# +# `output` and `title` are optional (both overridable with -o / -t on the CLI). +# Relative `file` paths are resolved from THIS config file's directory. + +output = "plotting/plots/baselines.png" +title = "Reproducing Low-Precision Baselines on GPT-OSS 20b" + +[[runs]] +file = "../slurm-207989.out" +name = "bf16" + +[[runs]] +file = "../slurm-207639.out" +name = "lpt_recipe" + +[[runs]] +file = "../slurm-208070.out" +name = "deoscillation" \ No newline at end of file diff --git a/plotting/lpt.toml b/plotting/lpt.toml new file mode 100644 index 00000000..764e2916 --- /dev/null +++ b/plotting/lpt.toml @@ -0,0 +1,29 @@ +# Example config for plot_val_loss.py +# python3 plot_val_loss.py -c runs.example.toml +# +# `output` and `title` are optional (both overridable with -o / -t on the CLI). +# Relative `file` paths are resolved from THIS config file's directory. + +output = "plotting/plots/lpt.png" +title = "Low-Precision Training on GPT-OSS 20b" + +[[runs]] +file = "../slurm-207989.out" +name = "bf16" + +[[runs]] +file = "../slurm-207639.out" +name = "lpt_recipe" + +[[runs]] +file = "../slurm-208070.out" +name = "deoscillation" + +[[runs]] +file = ["../slurm-208455.out","../slurm-208634.out"] +name = "midmax" + +[[runs]] +file = ["../slurm-208293.out"] +name = "madam" + diff --git a/plotting/plot_training_stats.py b/plotting/plot_training_stats.py new file mode 100644 index 00000000..0e18adbc --- /dev/null +++ b/plotting/plot_training_stats.py @@ -0,0 +1,477 @@ +#!/usr/bin/env python3 +"""Plot training/validation loss and grad norm from slurm .out logs and/or +TensorBoard event files. + +Two source kinds are supported and may be freely mixed within a single plot +(and even within a single run): + + * slurm .out logs -- parsed by regex from lines like: + ... validate step: 768 loss: 4.9572 memory: ... + (ANSI color codes are stripped before matching.) + + * TensorBoard sources -- either an event file (events.out.tfevents.*) or a + directory containing them (e.g. a run's tb/ dir, whose timestamped subdirs + from resumes are read in order). Scalar tags read: + loss_metrics/global_avg_loss -> training loss + grad_norm -> gradient norm + validation_metrics/loss -> validation loss + +Alongside the plot, a table on the right lists each method's loss at every +step (union of all steps across files; blank where a method has no datapoint). + +Usage: + # inline on the command line + python3 plot_training_stats.py [ ...] [-o output.png] + # or from a .toml config + python3 plot_training_stats.py -c runs.toml + +A custom legend/column name can be given per file with "path=label" syntax, e.g.: + python3 plot_training_stats.py run_a.out=baseline run_b.out="lr 3e-4" +Files given without "=label" fall back to their basename. + +A single run may span several .out files (e.g. a resumed job). List them +comma-separated on the CLI; they are parsed in order and drawn as one curve: + python3 plot_training_stats.py part1.out,part2.out=baseline + +TOML config format (see runs.example.toml): + output = "val_loss.png" # optional, overridden by -o + title = "Validation loss vs. step" # optional + max_step = 5000 # optional, global step cutoff (-m) + + [[runs]] + file = "slurm-207639.out" + name = "baseline" + + [[runs]] + file = "slurm-207504.out" + name = "lr 3e-4" + + # a run split across multiple files: pass a list to "file" (or "files"). + # files are concatenated in the given order into a single curve/column. + [[runs]] + file = ["slurm-208455.out", "slurm-208634.out"] + name = "midmax" + + # a TensorBoard run: point "file" at the tb dir (or a single event file). + # .out logs and tb sources may be mixed in the same list and across runs. + [[runs]] + file = "gpt_oss_20b-pretrain-bf16/tb" + name = "bf16 (tb)" + + # limit how far along the x-axis data is shown. a top-level `max_step` + # applies to every run; a per-run `max_step` overrides it for that run. + # (also settable on the CLI with --max-step.) + max_step = 5000 # optional, global cutoff + + [[runs]] + file = "slurm-209000.out" + name = "short view" + max_step = 2000 # optional, per-run cutoff +""" +import argparse +import os +import re +import sys +import textwrap + +try: + import tomllib # Python >= 3.11 +except ModuleNotFoundError: # Python <= 3.10 + try: + import tomli as tomllib + except ModuleNotFoundError: + tomllib = None + +import matplotlib +matplotlib.use("Agg") # save-to-file only; avoids loading/blocking on a GUI backend +import matplotlib.pyplot as plt + +plt.style.use("dark_background") + +# strip ANSI escape sequences, then pull step + loss +ANSI = re.compile(r"\x1b\[[0-9;]*m") +VAL = re.compile(r"validate step:\s*(\d+)\s+loss:\s*([\d.]+)") +# training lines look like "step: 2 loss: 12.68598 grad_norm: 1.3970 ..." +# (exclude "validate step:" via a negative lookbehind) +TRAIN = re.compile(r"(? {val_step: val_loss} (table data) + colors = {} # label -> line color + total_val = 0 # validation datapoints (table) + + for paths, label, run_max_step in runs: + # per-run max_step overrides the global cutoff; the global applies otherwise + max_step = run_max_step if run_max_step is not None else global_max_step + existing = [p for p in paths if os.path.exists(p)] + missing = [p for p in paths if not os.path.exists(p)] + if missing: + print(f"warning: file(s) not found: {', '.join(missing)}", file=sys.stderr) + if not existing: + # keep a placeholder so the run still shows in the legend + table + (line,) = ax.plot([], [], linewidth=1.2, label=f"{label} (missing)") + ax_grad.plot([], [], linewidth=1.2, color=line.get_color()) + ax_val.plot([], [], linewidth=1.2, color=line.get_color()) + labels.append(label) + loss_by_step[label] = {} + colors[label] = line.get_color() + continue + + tr_steps, tr_losses, grad_norms, val_steps, val_losses = clip_to_step( + parse_many(existing), max_step) + + # curves plot TRAINING loss (no per-point marker — too dense) + (line,) = ax.plot(tr_steps, tr_losses, linewidth=1.2, label=label) + color = line.get_color() + + # grad-norm curve below, in the matching color (skip steps w/o grad_norm) + gsteps = [s for s, g in zip(tr_steps, grad_norms) if g is not None] + gvals = [g for g in grad_norms if g is not None] + ax_grad.plot(gsteps, gvals, linewidth=1.2, color=color, label=label) + + # validation-loss curve below the grad-norm plot, matching color + # (markers here since validation points are sparse) + ax_val.plot(val_steps, val_losses, linewidth=1.2, marker="o", + markersize=3, color=color, label=label) + + labels.append(label) + loss_by_step[label] = dict(zip(val_steps, val_losses)) + colors[label] = color + total_val += len(val_steps) + + if not labels: + sys.exit("No datapoints found in any input file") + + ax.set_xlabel("step") + ax.set_ylabel("training loss") + ax.set_title("Training Loss", fontweight="bold") + ax.grid(True, alpha=0.3) + ax.legend() + if max_step is not None: + ax.set_xlim(right=max_step) # ax_grad/ax_val share this x-axis + + ax_grad.set_xlabel("step") + ax_grad.set_ylabel("grad norm") + ax_grad.set_title("Gradient Norm", fontweight="bold") + ax_grad.grid(True, alpha=0.3) + ax_grad.sharex(ax) + + ax_val.set_xlabel("step") + ax_val.set_ylabel("validation loss") + ax_val.set_title("Validation Loss", fontweight="bold") + ax_val.grid(True, alpha=0.3) + ax_val.sharex(ax) + + # --- optional details text box (top-right, above the table) --- + if ax_det is not None: + ax_det.axis("off") + ax_det.text( + 0.5, 0.98, textwrap.fill(details, width=34), + transform=ax_det.transAxes, ha="center", va="top", + fontsize=8, color="white", + bbox=dict(boxstyle="round,pad=0.5", facecolor="none", edgecolor="gray"), + ) + + # --- table of VALIDATION losses per step for each method --- + all_steps = sorted({s for d in loss_by_step.values() for s in d}) + # wrap long method names so they don't overflow their table column + wrapped = ["\n".join(textwrap.wrap(lab, width=14)) or lab for lab in labels] + col_labels = ["step"] + wrapped + cell_text = [] + for s in all_steps: + row = [str(s)] + for lab in labels: + v = loss_by_step[lab].get(s) + row.append(f"{v:.4f}" if v is not None else "") + cell_text.append(row) + + ax_tbl.axis("off") + ax_tbl.set_title("Validation Loss", fontweight="bold") + table = ax_tbl.table( + cellText=cell_text, + colLabels=col_labels, + cellLoc="center", + loc="center", + ) + table.auto_set_font_size(False) + table.set_fontsize(8) + table.scale(1, 1.3) + + # dark theme: transparent cells with light-gray borders + white text + for cell in table.get_celld().values(): + cell.set_edgecolor("gray") + cell.set_facecolor("none") + cell.set_text_props(color="white") + + # give the header row enough height for the tallest wrapped label + max_lines = max(lbl.count("\n") + 1 for lbl in col_labels) + base_h = table[0, 0].get_height() + # color the header cells to match each method's line + for c, (raw, lab) in enumerate(zip(["step"] + labels, col_labels)): + cell = table[0, c] + cell.set_height(base_h * max_lines) + cell.set_text_props(weight="bold") + if raw in colors: + # method line colors are light pastels on the dark theme, + # so black header text reads better than white + cell.set_facecolor(colors[raw]) + cell.set_text_props(weight="bold", color="black") + + # figure-level supertitle above both the plot subtitle and the table title + fig.suptitle(suptitle, fontsize=16, fontweight="bold") + fig.tight_layout(rect=(0, 0, 1, 0.96)) # leave room for the suptitle + fig.savefig(out, dpi=150) + # also emit vector PDF + SVG alongside the raster output + base = os.path.splitext(out)[0] + pdf_out, svg_out = base + ".pdf", base + ".svg" + fig.savefig(pdf_out) + fig.savefig(svg_out) + n_files = sum(len(paths) for paths, *_ in runs) + print(f"Wrote {out}, {pdf_out} and {svg_out} ({total_val} validation " + f"datapoints in table across {len(runs)} run(s), {n_files} file(s))") + + +if __name__ == "__main__": + main() diff --git a/rdma_tests/RDMA.md b/rdma_tests/RDMA.md new file mode 100644 index 00000000..6c1346f0 --- /dev/null +++ b/rdma_tests/RDMA.md @@ -0,0 +1,45 @@ +# Testing RDMA between two nodes + +Note: you will need `apt-get update && apt-get install -y ibverbs-utils iproute2 perftest` for tools inside docker container. + +```bash +`ibv_devices` # list RDMA devices (e.g. mlx5_0) +`ibv_devinfo` # PortState: PORT_ACTIVE and note the link layer (IB vs. Ethernet/ROCE) +`rdma link show` # link state per device +``` + +# Raw RDMA loopback between two nodes + +```bash +# On node A (server): +ib_write_bw -d mlx5_0 -F --report_gbits + +# On node B (client), point at node A's IP: +ib_write_bw -d mlx5_0 -F --report_gbits + +``` + +You can get `` for a given RDMA interface like this: + +```bash +# Map RDMA device -> netdev: +ibdev2netdev # look for a line like: "mlx5_0 port 1 ==> rdma0 (Up)" +rdma link show # alternatively, look here. e.g. mlx5_0/1 ... netdev rdma0 + +# Then get that interface's IP: +ip -4 addr show rdma0 # look for the "inet x.x.x.x" line + +``` + +# Test RDMA using torchrun + +If you have terminal access to both machines, you can use `torchrun` to double check RDMA support. Look for a line in NCCL INFO outputs that says `NET/IB : Using [0]mlx5_1:1/RoCE ... [9]mlx5_9:1/RoCE`. If it falls back to TCP you will get an output like: `NCCL INFO NET/Socket ...`. + +```bash +export NCCL_DEBUG=INFO + +torchrun \ + --nnodes=2 --nproc-per-node=8 --node-rank= \ + --master-addr= --master-port=29500 \ + test_rdma_allreduce.py +``` \ No newline at end of file diff --git a/rdma_tests/start_container.sh b/rdma_tests/start_container.sh new file mode 100755 index 00000000..8446024d --- /dev/null +++ b/rdma_tests/start_container.sh @@ -0,0 +1,50 @@ +#!/bin/bash +# Start the ALTO multinode container with host-provided RDMA support. +# +# RDMA providers (the libibverbs plugins, e.g. libmlx5 / libbnxt_re) are ABI-tied +# to the host's rdma-core AND kernel, both of which vary across our machines. So +# instead of baking a provider into the image, we mount the host's entire +# libibverbs userspace -- the library, its provider modules, and the .driver +# registration files -- read-only over the same paths. The in-container RDMA +# stack then always matches whatever kernel/NIC this particular host has. +# +# Set USE_HOST_RDMA=0 to skip the mounts and use the image's own stack. +set -euo pipefail + +IMAGE="${IMAGE:-alto:multinode}" +CONTAINER="${CONTAINER:-alto_multinode}" +ALTO_DIR="${ALTO_DIR:-$HOME/lpt_branch/ALTO}" + +docker_args=( + -it + --rm + --name "$CONTAINER" + --network host + --ipc host + --shm-size=16g + --cap-add=IPC_LOCK + --ulimit memlock=-1:-1 + -v "$ALTO_DIR:/alto" +) + +# GPU / RDMA character devices -- add only if present on this host. +for dev in /dev/infiniband /dev/dri /dev/kfd; do + [[ -e "$dev" ]] && docker_args+=(--device="$dev") +done + +# Host RDMA userspace. Mounting the host's libibverbs.so together with its +# provider modules and driver configs keeps the library ABI self-consistent +# (host lib + host providers) while matching the host kernel's uABI. +if [[ "${USE_HOST_RDMA:-1}" == "1" && -d /etc/libibverbs.d ]]; then + for lib in $(ldconfig -p | awk '/lib(ibverbs|rdmacm|ibumad|mlx5|mlx4|bnxt_re|efa|irdma|hns|cxgb4)\.so/ {print $NF}' | sort -u); do + [[ -e "$lib" ]] || continue + docker_args+=(-v "$lib:$lib:ro") + real="$(readlink -f "$lib")" + [[ "$real" != "$lib" && -e "$real" ]] && docker_args+=(-v "$real:$real:ro") + done + for d in /usr/lib/x86_64-linux-gnu/libibverbs /usr/lib64/libibverbs /etc/libibverbs.d; do + [[ -d "$d" ]] && docker_args+=(-v "$d:$d:ro") + done +fi + +docker run "${docker_args[@]}" "$IMAGE" bash diff --git a/rdma_tests/test_rdma_allreduce.py b/rdma_tests/test_rdma_allreduce.py new file mode 100644 index 00000000..98cbaffa --- /dev/null +++ b/rdma_tests/test_rdma_allreduce.py @@ -0,0 +1,50 @@ +import os +import time + +import torch +import torch.distributed as dist + + +def main(): + dist.init_process_group("nccl") + rank = dist.get_rank() + world = dist.get_world_size() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dev = torch.device("cuda", local_rank) + + if rank == 0: + print(f"[setup] world_size={world}", flush=True) + + # 1 GiB tensor (256M fp32 elements) + numel = 256 * 1024 * 1024 + x = torch.ones(numel, device=dev) + nbytes = x.element_size() * x.numel() + + # warmup + for _ in range(5): + dist.all_reduce(x) + torch.cuda.synchronize() + + iters = 20 + t0 = time.perf_counter() + for _ in range(iters): + dist.all_reduce(x) + torch.cuda.synchronize() + dt = (time.perf_counter() - t0) / iters + + # all-reduce bus-bandwidth: 2*(n-1)/n * size / time + algbw = nbytes / dt + busbw = algbw * 2 * (world - 1) / world + if rank == 0: + print( + f"[result] size={nbytes/1e9:.2f} GB time={dt*1e3:.2f} ms " + f"algbw={algbw/1e9:.1f} GB/s busbw={busbw/1e9:.1f} GB/s", + flush=True, + ) + + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/rdma_tests/test_rdma_via_pytorch.sh b/rdma_tests/test_rdma_via_pytorch.sh new file mode 100644 index 00000000..102cc0f3 --- /dev/null +++ b/rdma_tests/test_rdma_via_pytorch.sh @@ -0,0 +1,9 @@ +export NCCL_DEBUG=INFO +# export NCCL_IB_HCA=mlx5 # use your device prefix from ibdev2netdev +# export NCCL_SOCKET_IFNAME= # the bootstrap iface (e.g. enp1s0f0 or eth0) if autodetect picks the wrong one + + +torchrun \ + --nnodes=2 --nproc-per-node=8 --node-rank=0 \ + --master-addr= --master-port=29500 \ + test_rdma_allreduce.py diff --git a/scripts/check_adahop_resume_divergence.py b/scripts/check_adahop_resume_divergence.py new file mode 100644 index 00000000..2a06a45f --- /dev/null +++ b/scripts/check_adahop_resume_divergence.py @@ -0,0 +1,173 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Diagnostic: does the AdaHOP calibration blob restore identically on every rank? + +Hypothesis under test +--------------------- +On resume, adahop jobs deadlock on the first training step inside the MoE +``ALLTOALL_BASE`` collective. The suspected cause is that the checkpointed +calibration state -- a plain pickled Python blob restored by DCP +(alto/modifiers/lpt/adahop_internals/calibration_state.py) -- is NOT delivered +identically to all ranks. If rank 0 sees ``completed=True`` (and takes the +"apply modes, skip calibration" branch) while some other rank sees +``completed=False`` (and takes ``_arm_calibration()``), the two branches issue +different collectives and the expert-parallel all-to-all hangs forever. Raising +the NCCL timeout cannot fix a control-flow divergence. + +This script reproduces ONLY the calibration-state load path -- no model, no +training, no MoE -- so it finishes in seconds instead of hanging for 30 min. +Every rank loads the blob, then all ranks all-gather (completed, step_idx, +n_modes, modes_hash) and rank 0 reports whether they agree. + +Run with torchrun so every rank participates, and DO NOT filter rank output: + + cd /home/ybouquet/projects/ALTO + CKPT=./gpt_oss_20b-pretrain-subset-mxfp4-adahop-srfix-mi300x-outputs/checkpoint/step-500 \ + torchrun --nproc_per_node=8 --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ + scripts/check_adahop_resume_divergence.py + +Each rank also writes /tmp/adahop_divergence_rank.txt so results survive any +output filtering. +""" + +import hashlib +import json +import os +import sys + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp + + +def _modes_hash(modes_by_fqn) -> str: + """Deterministic hash of the per-layer modes dict (order-independent).""" + blob = json.dumps(modes_by_fqn, sort_keys=True, default=str) + return hashlib.sha256(blob.encode("utf-8")).hexdigest()[:16] + + +def main() -> int: + ckpt = os.environ.get("CKPT") + if not ckpt: + print("ERROR: set CKPT=", file=sys.stderr) + return 2 + if not os.path.isdir(ckpt): + print(f"ERROR: CKPT is not a directory: {ckpt}", file=sys.stderr) + return 2 + + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + torch.cuda.set_device(local_rank) + # RCCL/NCCL share the "nccl" backend name on ROCm. + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world = dist.get_world_size() + + # Import AFTER process group init so any module-level state is fresh. + from alto.modifiers.lpt.adahop_internals.calibration_state import ( + get_adahop_calibration_state, + is_calibration_completed, + get_calibration_modes, + reset_calibration_state, + ) + + # Start every rank from the fresh-run default so what we observe after the + # load is purely what DCP delivered -- not leftover state. + reset_calibration_state() + + before_completed = is_calibration_completed() + + # Reproduce EXACTLY the training resume path: register the Stateful manager + # under the same key the trainer uses ("adahop_calibration") and let DCP + # populate the module-level _STATE via load_state_dict. + manager = get_adahop_calibration_state() + state = {"adahop_calibration": manager} + + load_error = "" + try: + dcp.load(state, checkpoint_id=ckpt) + except Exception as e: # noqa: BLE001 -- we want to see per-rank failures + load_error = f"{type(e).__name__}: {e}" + + from alto.modifiers.lpt.adahop_internals.calibration_state import get_calibration_state_dict + completed = is_calibration_completed() + modes = get_calibration_modes() + n_modes = len(modes) + step_idx = get_calibration_state_dict().get("step_idx", 0) + h = _modes_hash(modes) if modes else "" + + # Which training branch WOULD this rank take on step 501? + branch = "APPLY_MODES(skip calib)" if completed else "ARM_CALIBRATION" + + line = (f"[rank {rank}/{world}] before_completed={before_completed} " + f"after_completed={completed} step_idx={step_idx} " + f"n_modes={n_modes} modes_hash={h} branch={branch} " + f"load_error={load_error or ''}") + + # Persist per-rank so nothing is lost to output filtering. Default to /tmp; + # set DIVERGENCE_OUT_DIR to a bind-mounted dir to keep files after the + # container is torn down. + out_dir = os.environ.get("DIVERGENCE_OUT_DIR", "/tmp") + os.makedirs(out_dir, exist_ok=True) + with open(os.path.join(out_dir, f"adahop_divergence_rank{rank}.txt"), "w") as f: + f.write(line + "\n") + print(line, flush=True) + + # Gather a compact tuple from every rank to rank 0 for a verdict. + record = { + "rank": rank, + "completed": bool(completed), + "step_idx": int(step_idx), + "n_modes": int(n_modes), + "modes_hash": h, + "branch": branch, + "load_error": load_error, + } + gathered = [None] * world + dist.all_gather_object(gathered, record) + + rc = 0 + if rank == 0: + print("\n==================== VERDICT ====================", flush=True) + for r in sorted(gathered, key=lambda x: x["rank"]): + print(f" rank {r['rank']}: completed={r['completed']} " + f"step_idx={r['step_idx']} n_modes={r['n_modes']} hash={r['modes_hash']} " + f"branch={r['branch']} err={r['load_error'] or ''}", + flush=True) + + completed_set = {r["completed"] for r in gathered} + hash_set = {r["modes_hash"] for r in gathered} + branch_set = {r["branch"] for r in gathered} + err_ranks = [r["rank"] for r in gathered if r["load_error"]] + + print("\n ---- analysis ----", flush=True) + if err_ranks: + print(f" ✗ LOAD FAILED on ranks {err_ranks} -- blob not restorable there.", flush=True) + rc = 1 + if len(branch_set) > 1: + print(f" ✗ DIVERGENCE CONFIRMED: ranks disagree on the step-501 branch " + f"{sorted(branch_set)}. This is the deadlock: different ranks issue " + f"different collectives.", flush=True) + rc = 1 + elif len(completed_set) > 1: + print(f" ✗ DIVERGENCE: 'completed' flag differs across ranks {completed_set}.", flush=True) + rc = 1 + elif len(hash_set) > 1: + print(f" ⚠ modes agree on branch but DIFFER in content across ranks " + f"{sorted(hash_set)} -- would not deadlock on control flow but numerics " + f"diverge per rank.", flush=True) + rc = 1 + else: + print(f" ✓ ALL RANKS AGREE: completed={completed_set.pop()}, " + f"identical modes_hash, same branch. Calibration-state restore is NOT " + f"the divergence source -- look elsewhere (e.g. Hadamard seed, " + f"or a genuinely different collective ordering).", flush=True) + print("=================================================", flush=True) + + dist.barrier() + dist.destroy_process_group() + return rc + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/debug_observer_viz.py b/scripts/debug_observer_viz.py new file mode 100644 index 00000000..7e4c9469 --- /dev/null +++ b/scripts/debug_observer_viz.py @@ -0,0 +1,327 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Standalone visualizer for DebugObserverModifier dumps. + +No alto import — only torch + matplotlib required. + +Usage +----- +# Summary only (no plots written): +python scripts/debug_observer_viz.py --dump ./outputs/debug_obs_lpt.pt --summary-only + +# Full plots: +python scripts/debug_observer_viz.py \\ + --dump ./outputs/debug_obs_lpt.pt \\ + --out-dir ./viz + +# With BF16 baseline overlay: +python scripts/debug_observer_viz.py \\ + --dump ./outputs/debug_obs_lpt.pt \\ + --baseline ./outputs/debug_obs_bf16.pt \\ + --out-dir ./viz + +# Filter to specific layers: +python scripts/debug_observer_viz.py \\ + --dump ./outputs/debug_obs_lpt.pt \\ + --layer-regex 'experts' +""" + +from __future__ import annotations + +import argparse +import re +import sys +from pathlib import Path +from typing import Optional + +import torch + + +# --------------------------------------------------------------------------- +# Loading +# --------------------------------------------------------------------------- + +def load_dump(path: str) -> tuple[dict, dict]: + """Returns (layers_dict, meta_dict). + + layers_dict: {fqn: {step_idx: {"input": T, ...}}} + meta_dict: the "_meta" entry + """ + blob: dict = torch.load(path, map_location="cpu", weights_only=False) + meta = blob.pop("_meta", {}) + return blob, meta + + +# --------------------------------------------------------------------------- +# Stats helpers +# --------------------------------------------------------------------------- + +def _absmax(t: torch.Tensor) -> float: + return t.float().abs().max().item() + + +def _std(t: torch.Tensor) -> float: + return t.float().std().item() + + +def _collect_series(layer_data: dict, key: str) -> tuple[list[int], list[float], list[float]]: + """Returns (steps, absmax_values, std_values) for a given tensor key.""" + steps, abs_maxes, stds = [], [], [] + for step in sorted(s for s in layer_data if isinstance(s, int)): + tensors = layer_data[step] + if key not in tensors: + continue + t = tensors[key] + steps.append(step) + abs_maxes.append(_absmax(t)) + stds.append(_std(t)) + return steps, abs_maxes, stds + + +def _align_baseline(run_data: dict, baseline_data: dict) -> dict: + """Return a copy of baseline_data re-keyed to match run_data's step numbers. + + BF16 and MXFP4 runs may capture at different absolute step indices (e.g. + the calibration pre-steps shift the MXFP4 counter). We match by ordinal + position (1st capture ↔ 1st capture, 2nd ↔ 2nd, …) so overlays always + show the same training epoch regardless of step numbering. + """ + run_steps = sorted(s for s in run_data if isinstance(s, int)) + base_steps = sorted(s for s in baseline_data if isinstance(s, int)) + remapped = {} + for run_step, base_step in zip(run_steps, base_steps): + remapped[run_step] = baseline_data[base_step] + # preserve the "active" gate key if present + if "active" in baseline_data: + remapped["active"] = baseline_data["active"] + return remapped + + +# --------------------------------------------------------------------------- +# Summary +# --------------------------------------------------------------------------- + +def print_summary(layers: dict, meta: dict) -> None: + print(f"\n=== DebugObserver summary ===") + print(f" rank: {meta.get('rank', 'unknown')}") + print(f" iterations captured: {meta.get('iterations_captured', [])}") + print(f" layers: {len(layers)}") + print(f" capture_every: {meta.get('capture_every', '?')} max_captures: {meta.get('max_captures', '?')}") + print(f" mlp2 input captured: {meta.get('mlp2_input_captured', False)}") + print() + print(f"{'Layer':<60} {'keys':>30} {'input absmax':>12} {'gw absmax':>12}") + print("-" * 120) + for fqn, layer_data in sorted(layers.items()): + steps = [s for s in layer_data if isinstance(s, int)] + if not steps: + continue + sample = layer_data[min(steps)] + keys = ",".join(sorted(sample.keys())) + inp_absmax = _absmax(sample["input"]) if "input" in sample else float("nan") + gw_key = "grad_weight" if "grad_weight" in sample else ( + "grad_mlp1_weight" if "grad_mlp1_weight" in sample else None + ) + gw_absmax = _absmax(sample[gw_key]) if gw_key else float("nan") + print(f"{fqn:<60} {keys:>30} {inp_absmax:>12.4f} {gw_absmax:>12.4f}") + print() + + +# --------------------------------------------------------------------------- +# Plotting +# --------------------------------------------------------------------------- + +def _plot_time_series( + ax_absmax, + ax_std, + steps: list[int], + absmax_vals: list[float], + std_vals: list[float], + label: str, + color: str, + is_baseline: bool = False, +) -> None: + lw = 1.5 if not is_baseline else 2.5 + ls_abs = "-" if not is_baseline else "--" + ls_std = "--" if not is_baseline else ":" + marker = "o" if not is_baseline else "s" + ax_absmax.plot(steps, absmax_vals, marker=marker, label=label, color=color, + linewidth=lw, linestyle=ls_abs, zorder=3 if is_baseline else 2) + ax_std.plot(steps, std_vals, marker=marker, label=label, color=color, + linewidth=lw, linestyle=ls_std, zorder=3 if is_baseline else 2) + + +def _plot_histogram(ax, t: torch.Tensor, label: str, color: str, + is_baseline: bool = False) -> None: + vals = t.float().abs().flatten().numpy() + alpha = 0.35 if is_baseline else 0.6 + ax.hist(vals, bins=50, histtype="stepfilled", color=color, density=True, + alpha=alpha, label=label, edgecolor="none") + + +def plot_layer( + fqn: str, + layer_data: dict, + out_dir: Path, + baseline_data: Optional[dict] = None, +) -> None: + try: + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + except ImportError: + print("matplotlib not installed; skipping plots", file=sys.stderr) + return + + # Remap baseline step numbers to match the run's step numbers by ordinal + # position, so overlays work even when step counters differ between runs. + aligned_baseline = _align_baseline(layer_data, baseline_data) if baseline_data is not None else None + + tensor_keys = sorted({ + k for step, tensors in layer_data.items() + if isinstance(step, int) + for k in tensors + }) + + safe_fqn = re.sub(r"[^a-zA-Z0-9_.\-]", "_", fqn) + layer_out = out_dir / safe_fqn + layer_out.mkdir(parents=True, exist_ok=True) + + for key in tensor_keys: + steps, absmax_vals, std_vals = _collect_series(layer_data, key) + if not steps: + continue + + # --- time-series plot --- + fig, (ax_absmax, ax_std) = plt.subplots(2, 1, figsize=(10, 6), sharex=True) + fig.suptitle(f"{fqn}\n{key}", fontsize=9) + _plot_time_series(ax_absmax, ax_std, steps, absmax_vals, std_vals, label="run", color="blue") + if aligned_baseline is not None: + bsteps, babs, bstd = _collect_series(aligned_baseline, key) + if bsteps: + _plot_time_series(ax_absmax, ax_std, bsteps, babs, bstd, + label="baseline", color="red", is_baseline=True) + ax_absmax.set_ylabel("absmax") + ax_absmax.legend(fontsize=7) + ax_absmax.grid(True, alpha=0.3) + ax_std.set_ylabel("std") + ax_std.set_xlabel("step") + ax_std.legend(fontsize=7) + ax_std.grid(True, alpha=0.3) + ts_path = layer_out / f"{key}_timeseries.png" + fig.tight_layout() + fig.savefig(ts_path, dpi=120) + plt.close(fig) + + # --- per-iteration histogram --- + steps_sorted = sorted(steps) + n_cols = min(5, len(steps_sorted)) + n_rows = (len(steps_sorted) + n_cols - 1) // n_cols + fig, axes = plt.subplots(n_rows, n_cols, figsize=(4 * n_cols, 3 * n_rows), squeeze=False) + fig.suptitle(f"{fqn} | {key} | abs-value histograms", fontsize=9) + for idx, step in enumerate(steps_sorted): + ax = axes[idx // n_cols][idx % n_cols] + t = layer_data[step].get(key) + if t is None: + ax.set_visible(False) + continue + _plot_histogram(ax, t, label=f"step {step}", color="steelblue") + if aligned_baseline is not None and step in aligned_baseline: + bt = aligned_baseline[step].get(key) + if bt is not None: + _plot_histogram(ax, bt, label="baseline", color="red", is_baseline=True) + ax.set_title(f"step {step}", fontsize=7) + ax.legend(fontsize=6) + ax.set_xlabel("|value|", fontsize=6) + # hide unused axes + for idx in range(len(steps_sorted), n_rows * n_cols): + axes[idx // n_cols][idx % n_cols].set_visible(False) + hist_path = layer_out / f"{key}_histograms.png" + fig.tight_layout() + fig.savefig(hist_path, dpi=120) + plt.close(fig) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def _match_baseline_fqn(fqn: str, baseline_layers: dict) -> Optional[dict]: + """Look up baseline data for a given FQN. + + Tries exact match first, then falls back to matching by the longest common + suffix so that FQN differences caused by wrapper modules (e.g. FSDP shards + or LPT submodule renaming) don't break the overlay. + """ + if fqn in baseline_layers: + return baseline_layers[fqn] + # suffix match: find the baseline FQN whose suffix best matches + best_fqn, best_len = None, 0 + for bfqn in baseline_layers: + # find longest common suffix component-by-component + fqn_parts = fqn.split(".") + bfqn_parts = bfqn.split(".") + common = 0 + for a, b in zip(reversed(fqn_parts), reversed(bfqn_parts)): + if a == b: + common += 1 + else: + break + if common > best_len: + best_len, best_fqn = common, bfqn + if best_len >= 2: # require at least 2 matching suffix components + return baseline_layers[best_fqn] + return None + + +def main() -> None: + parser = argparse.ArgumentParser(description="Visualize DebugObserverModifier dumps") + parser.add_argument("--dump", required=True, help="Path to the .pt dump file") + parser.add_argument("--baseline", default=None, help="Optional BF16 baseline .pt dump") + parser.add_argument("--out-dir", default="./viz", help="Directory to write PNG files") + parser.add_argument("--layer-regex", default=None, help="Regex to filter layer FQNs") + parser.add_argument("--summary-only", action="store_true", help="Print summary, skip plotting") + parser.add_argument("--debug-fqns", action="store_true", + help="Print FQNs from both dumps and their matches, then exit") + args = parser.parse_args() + + layers, meta = load_dump(args.dump) + baseline_layers, baseline_meta = load_dump(args.baseline) if args.baseline else ({}, {}) + + if args.debug_fqns: + print(f"\n=== FQNs in dump ({len(layers)}) ===") + for fqn in sorted(layers): + matched = _match_baseline_fqn(fqn, baseline_layers) + tag = "[matched]" if matched is not None else "[NO MATCH]" + print(f" {tag} {fqn}") + print(f"\n=== FQNs in baseline ({len(baseline_layers)}) ===") + for fqn in sorted(baseline_layers): + print(f" {fqn}") + print(f"\ndump steps: {meta.get('iterations_captured', [])}") + print(f"baseline steps: {baseline_meta.get('iterations_captured', [])}") + return + + if args.layer_regex: + pat = re.compile(args.layer_regex) + layers = {fqn: v for fqn, v in layers.items() if pat.search(fqn)} + baseline_layers = {fqn: v for fqn, v in baseline_layers.items() if pat.search(fqn)} + + print_summary(layers, meta) + + if args.summary_only: + return + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + for fqn, layer_data in sorted(layers.items()): + print(f"Plotting {fqn} ...", end=" ", flush=True) + baseline_data = _match_baseline_fqn(fqn, baseline_layers) + plot_layer(fqn, layer_data, out_dir, baseline_data=baseline_data) + print("done") + + print(f"\nPlots written to {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/scripts/entrypoint.sh b/scripts/entrypoint.sh new file mode 100644 index 00000000..b7a87bd1 --- /dev/null +++ b/scripts/entrypoint.sh @@ -0,0 +1,61 @@ +#!/bin/bash + +set -ex + +mkdir /workspace +ln -s /wekafs/hanwang2/huggingface /huggingface +ln -s /wekafs/hanwang2/workspace /workspace/workspace + +cd /workspace/workspace/ALTO + +pip install --no-build-isolation -e 3rdparty/torchtitan +pip install -e . + +RECIPE_FILE="/tmp/recipe.yaml" +TRAIN_FILE=${TRAIN_FILE:-"alto.train"} +MODULE=${MODULE:-"gpt_oss"} +CONFIG=${CONFIG:-"gpt_oss_20b_lpt"} +GBS=${GBS:-16} +OUTPUT_DIR=${OUTPUT_DIR:-"/workspace/workspace/ALTO/debug-outputs"} + +cat > $RECIPE_FILE < tuple[str, str]: + """'row-col' -> ('row', 'col'); tolerant of the 'none-none' default.""" + a, _, b = (pair or "none-none").partition("-") + a = a if a in SINGLE_PATTERN_COLORS else "none" + b = b if b in SINGLE_PATTERN_COLORS else "none" + return a, b + + +# --------------------------------------------------------------------------- +# Loading + merge +# --------------------------------------------------------------------------- + +def load_and_merge(dump_paths: List[str]) -> tuple[dict, dict, dict]: + """Merge per-rank dumps into one structure. + + Returns (merged, majority, meta) where + merged[fqn][step][gemm][global_expert_id] = per-step record + majority[fqn][gemm][global_expert_id] = majority-vote record + meta = combined _meta (num_global_experts, etc.) + """ + merged: Dict[str, Any] = {} + majority: Dict[str, Any] = {} + meta: Dict[str, Any] = {} + ep_sizes = set() + global_counts = set() + ranks = [] + steps_all = set() + + for p in dump_paths: + blob = torch.load(p, map_location="cpu", weights_only=False) + m = blob.pop("_meta", {}) + maj = blob.pop("_majority", {}) + ranks.append(m.get("rank", "?")) + ep_sizes.add(m.get("ep_size", 1)) + global_counts.add(m.get("num_global_experts", 0)) + steps_all.update(m.get("iterations_captured", [])) + for fqn, per_step in blob.items(): + f = merged.setdefault(fqn, {}) + for step, per_gemm in per_step.items(): + s = f.setdefault(step, {}) + for gemm, records in per_gemm.items(): + g = s.setdefault(gemm, {}) + g.update(records) # global ids are disjoint across ranks + for fqn, per_gemm in maj.items(): + mf = majority.setdefault(fqn, {}) + for gemm, records in per_gemm.items(): + mf.setdefault(gemm, {}).update(records) # disjoint global ids + + meta = { + "ranks": sorted(ranks, key=str), + "ep_size": max(ep_sizes) if ep_sizes else 1, + "num_global_experts": max(global_counts) if global_counts else 0, + "iterations_captured": sorted(steps_all), + } + return merged, majority, meta + + +# --------------------------------------------------------------------------- +# Summary +# --------------------------------------------------------------------------- + +def print_summary(majority: dict, meta: dict) -> None: + print("\n=== MoE matmul pattern summary (majority over captured steps) ===") + print(f" ranks merged: {meta.get('ranks')}") + print(f" ep_size: {meta.get('ep_size')} global experts: {meta.get('num_global_experts')}") + print(f" steps captured: {meta.get('iterations_captured')}") + print(f" layers: {len(majority)}") + print() + for fqn in sorted(majority): + per_gemm = majority[fqn] + for gemm in GEMMS: + records = per_gemm.get(gemm) + if not records: + continue + # Tally the winning majority T-pair across experts per matmul path. + print(f" {fqn} | {gemm} | {len(records)} experts") + for path in MATMUL_PATHS: + counts: Dict[str, int] = {} + for rec in records.values(): + pair = rec.get(path, {}).get("pair", "none-none") + counts[pair] = counts.get(pair, 0) + 1 + tally = ", ".join(f"{k}:{v}" for k, v in sorted(counts.items())) + print(f" {path:<12} {tally}") + print() + + +# --------------------------------------------------------------------------- +# Plotting +# --------------------------------------------------------------------------- + +def _pair_grid(records: dict, expert_ids: List[int]): + import numpy as np + grid = np.full((len(expert_ids), len(MATMUL_PATHS)), _PAIR_INDEX["none-none"], dtype=int) + for r, eid in enumerate(expert_ids): + rec = records.get(eid, {}) + for c, path in enumerate(MATMUL_PATHS): + pair = rec.get(path, {}).get("pair", "none-none") + grid[r, c] = _PAIR_INDEX.get(pair, _PAIR_INDEX["none-none"]) + return grid + + +def plot_gemm(fqn: str, step, gemm: str, records: dict, out_dir: Path) -> None: + try: + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + from matplotlib.colors import ListedColormap, BoundaryNorm + from matplotlib.patches import Patch + import numpy as np + except ImportError: + print("matplotlib/numpy not installed; skipping plots", file=sys.stderr) + return + + expert_ids = sorted(records.keys()) + if not expert_ids: + return + grid = _pair_grid(records, expert_ids) + + cmap = ListedColormap([PATTERN_COLORS[p] for p in _PAIR_ORDER]) + norm = BoundaryNorm(np.arange(-0.5, len(_PAIR_ORDER) + 0.5, 1), cmap.N) + + height = max(3.0, 0.28 * len(expert_ids) + 1.5) + fig, ax = plt.subplots(figsize=(6, height)) + ax.imshow(grid, aspect="auto", cmap=cmap, norm=norm, interpolation="nearest") + ax.set_xticks(range(len(MATMUL_PATHS))) + ax.set_xticklabels(MATMUL_PATHS, rotation=20, ha="right", fontsize=8) + ax.set_yticks(range(len(expert_ids))) + ax.set_yticklabels([f"e{e}" for e in expert_ids], fontsize=6) + ax.set_ylabel("expert (global id)") + ax.set_title(f"{fqn}\nstep {step} | {gemm} — matmul-input outlier pattern", fontsize=9) + + legend = [Patch(facecolor=PATTERN_COLORS[p], label=p) for p in _PAIR_ORDER] + ax.legend(handles=legend, bbox_to_anchor=(1.02, 1), loc="upper left", + fontsize=6, title="T-pair", title_fontsize=7) + + safe = re.sub(r"[^a-zA-Z0-9_.\-]", "_", f"{fqn}_step{step}_{gemm}") + out_path = out_dir / f"{safe}.png" + fig.tight_layout() + fig.savefig(out_path, dpi=140, bbox_inches="tight") + plt.close(fig) + + +def plot_gemm_operands(fqn: str, step, gemm: str, records: dict, out_dir: Path) -> None: + """Per-operand view: y = the 6 matmul operands (3 matmuls, each A @ B split + into its two rows), x = experts. Each cell is colored by that operand's + single pattern (row/col/none) with the pattern name written inside.""" + try: + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + from matplotlib.colors import ListedColormap, BoundaryNorm + import numpy as np + except ImportError: + print("matplotlib/numpy not installed; skipping plots", file=sys.stderr) + return + + expert_ids = sorted(records.keys()) + if not expert_ids: + return + + # Build the row axis: one line per operand, grouped by matmul path. + row_labels: List[str] = [] + row_specs: List[tuple[str, int]] = [] # (matmul_path, operand_idx 0=A/1=B) + for path in MATMUL_PATHS: + (a_lbl, _), (b_lbl, _) = OPERAND_LABELS[path] + row_labels.append(f"{path}\nA: {a_lbl}") + row_specs.append((path, 0)) + row_labels.append(f"{path}\nB: {b_lbl}") + row_specs.append((path, 1)) + + n_rows, n_cols = len(row_specs), len(expert_ids) + grid = np.full((n_rows, n_cols), _SINGLE_INDEX["none"], dtype=int) + text = [["" for _ in range(n_cols)] for _ in range(n_rows)] + for r, (path, operand_idx) in enumerate(row_specs): + for c, eid in enumerate(expert_ids): + pair = records.get(eid, {}).get(path, {}).get("pair", "none-none") + pat = _split_pair(pair)[operand_idx] + grid[r, c] = _SINGLE_INDEX[pat] + text[r][c] = pat + + cmap = ListedColormap([SINGLE_PATTERN_COLORS[p] for p in _SINGLE_ORDER]) + norm = BoundaryNorm(np.arange(-0.5, len(_SINGLE_ORDER) + 0.5, 1), cmap.N) + + width = max(6.0, 0.6 * n_cols + 2.5) + fig, ax = plt.subplots(figsize=(width, 0.7 * n_rows + 1.5)) + ax.imshow(grid, aspect="auto", cmap=cmap, norm=norm, interpolation="nearest") + + ax.set_xticks(range(n_cols)) + ax.set_xticklabels([f"e{e}" for e in expert_ids], fontsize=7) + ax.set_xlabel("expert (global id)") + ax.set_yticks(range(n_rows)) + ax.set_yticklabels(row_labels, fontsize=7) + ax.set_title(f"{fqn}\nstep {step} | {gemm} — per-operand outlier pattern", fontsize=9) + ax.text(0.0, -0.16, + "patterns shown in as-fed-to-matmul orientation; (ᵀ) = operand enters " + "transposed, so its row/col is flipped vs the stored tensor", + transform=ax.transAxes, fontsize=6.5, style="italic", + color="#555555", ha="left", va="top", wrap=True) + + # Separator lines between the three matmul groups (every 2 rows). + for r in range(2, n_rows, 2): + ax.axhline(r - 0.5, color="white", linewidth=2.0) + + # Write the pattern name inside each cell. + for r in range(n_rows): + for c in range(n_cols): + ax.text(c, r, text[r][c], ha="center", va="center", + fontsize=7, color="black") + + safe = re.sub(r"[^a-zA-Z0-9_.\-]", "_", f"{fqn}_step{step}_{gemm}_operands") + out_path = out_dir / f"{safe}.png" + fig.tight_layout() + fig.savefig(out_path, dpi=140, bbox_inches="tight") + plt.close(fig) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def main() -> None: + parser = argparse.ArgumentParser(description="Visualize MoE matmul pattern dumps") + parser.add_argument("--dump-glob", required=True, + help="Glob for per-rank .pt dumps, e.g. './outputs/moe_patterns_rank*.pt'") + parser.add_argument("--out-dir", default="./viz", help="Directory for PNGs") + parser.add_argument("--layer-regex", default=None, help="Regex to filter layer FQNs") + parser.add_argument("--summary-only", action="store_true", help="Print summary, skip plotting") + parser.add_argument("--per-step", action="store_true", + help="Also render the raw per-step plots (default: majority only)") + args = parser.parse_args() + + dump_paths = sorted(glob.glob(args.dump_glob)) + if not dump_paths: + print(f"No dumps matched {args.dump_glob!r}", file=sys.stderr) + sys.exit(1) + print(f"Merging {len(dump_paths)} dump(s): {dump_paths}") + + merged, majority, meta = load_and_merge(dump_paths) + + if args.layer_regex: + pat = re.compile(args.layer_regex) + merged = {fqn: v for fqn, v in merged.items() if pat.search(fqn)} + majority = {fqn: v for fqn, v in majority.items() if pat.search(fqn)} + + # Older dumps predate the accumulated majority; fall back to the last step. + if not majority and merged: + print("No _majority block found; falling back to the last captured step.", + file=sys.stderr) + for fqn, per_step in merged.items(): + last = max(per_step) if per_step else None + if last is not None: + majority[fqn] = per_step[last] + + print_summary(majority, meta) + if args.summary_only: + return + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + # Majority plots (the default view) — labeled "majority" in title/filename. + for fqn in sorted(majority): + for gemm in GEMMS: + records = majority[fqn].get(gemm) + if records: + plot_gemm(fqn, "majority", gemm, records, out_dir) + plot_gemm_operands(fqn, "majority", gemm, records, out_dir) + + if args.per_step: + for fqn in sorted(merged): + for step in sorted(merged[fqn]): + for gemm in GEMMS: + records = merged[fqn][step].get(gemm) + if records: + plot_gemm(fqn, step, gemm, records, out_dir) + plot_gemm_operands(fqn, step, gemm, records, out_dir) + print(f"Plots written to {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/scripts/multinode.sh b/scripts/multinode.sh new file mode 100644 index 00000000..0b0e3d1f --- /dev/null +++ b/scripts/multinode.sh @@ -0,0 +1,63 @@ +#!/usr/bin/env bash +#SBATCH -A amd-arad +#SBATCH -p amd-arad-burst +#SBATCH --qos=low +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --time=04:00:00 +#SBATCH --job-name=alto-multinode +#SBATCH --output=alto-multinode-%j.out +#SBATCH --requeue + +set -euo pipefail + +IMAGE="alto:multinode" +ALTO_DIR="$HOME/lpt_branch/ALTO" +GPUS_PER_NODE=8 + +cd "$ALTO_DIR" +docker build -f Dockerfile.multinode -t $IMAGE . + +# First allocated node becomes the torchrun rendezvous host. +MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)" +MASTER_PORT="$((20000 + SLURM_JOB_ID % 20000))" + +export IMAGE ALTO_DIR GPUS_PER_NODE MASTER_ADDR MASTER_PORT + +echo "Nodes: $(scontrol show hostnames "$SLURM_JOB_NODELIST")" +echo "Master: ${MASTER_ADDR}:${MASTER_PORT}" + +# One srun task, and therefore one Docker container, per node. +srun --kill-on-bad-exit=1 bash -c ' + set -euo pipefail + + CONTAINER="alto_${SLURM_JOB_ID}_${SLURM_NODEID}" + + cleanup() { + docker stop --time 10 "$CONTAINER" >/dev/null 2>&1 || true + } + trap cleanup EXIT INT TERM + + docker run --rm \ + --name "$CONTAINER" \ + --network=host \ + --ipc=host \ + --shm-size=128g \ + --ulimit memlock=-1 \ + --device=/dev/kfd \ + --device=/dev/dri \ + --group-add video \ + -v /dev/infiniband:/dev/infiniband \ + -v "$ALTO_DIR:/alto" \ + -v /shared:/shared \ + -w /alto \ + "$IMAGE" \ + torchrun \ + --nnodes="$SLURM_JOB_NUM_NODES" \ + --nproc-per-node="$GPUS_PER_NODE" \ + --node-rank="$SLURM_NODEID" \ + --master-addr="$MASTER_ADDR" \ + --master-port="$MASTER_PORT" \ + rdma_tests/test_rdma_allreduce.py +' \ No newline at end of file diff --git a/scripts/train_gptoss20b.sh b/scripts/train_gptoss20b.sh new file mode 100755 index 00000000..5d227935 --- /dev/null +++ b/scripts/train_gptoss20b.sh @@ -0,0 +1,229 @@ +#!/usr/bin/env bash +# Run ALTO GPT-OSS 20B training on the current node. +# +# Example: +# NGPU=4 CONFIG=gpt_oss_debugmodel TRAINING_STEPS=20 \ +# bash ~/ALTO/train_gptoss20b.sh + +######## STEP 0: Install ALTO repository and update ALTO_DIR below +# git clone --recurse-submodules https://github.com/AMD-AGI/ALTO.git + + +######## STEP 1: Download the C4 dataset +######## make sure to update config_registry.py with appropriate data location +### OPTION 1: +# # Create desired download directory with the right permission +# cd /data/gpt_oss_20b +# # Download training and validation data +# bash <(curl -s https://raw.githubusercontent.com/mlcommons/r2-downloader/refs/heads/main/mlc-r2-downloader.sh) \ +# -d data https://training.mlcommons-storage.org/metadata/llama-3-1-8b-preprocessed-c4-dataset.uri +### OPTION 2: +# C4_CACHE="$HF_HOME_SHARED/datasets/allenai___c4" +# if [ -d "$C4_CACHE" ] && [ -n "$(ls -A "$C4_CACHE" 2>/dev/null)" ]; then +# echo "[train] C4 dataset already cached, skipping download." +# else +# echo "[train] Downloading C4 dataset (this may take a while) ..." +# docker exec "$CONTAINER" bash -c " +# python3 -c \" +# from datasets import load_dataset +# load_dataset('allenai/c4', 'en', split='train') +# load_dataset('allenai/c4', 'en', split='validation') +# \" +# " +# fi + +####### Viewing Loss Curves +# tensorboard events are saved in the checkpointing directory, one can +# view these by using the following command: +# tensorboard --logdir $CHECKPOINT_DIR --host 127.0.0.1 --port 6006 +# +# If running on remote machine, you will want to forward the port to the local machine: +# ssh -L 6006:localhost:6006 nfrumkin@useocpslog-002 + +set -euo pipefail + +# ----------------------------------------------------------------------------- +# Configuration +# ----------------------------------------------------------------------------- + +# If launched via sbatch, SLURM sets SLURM_JOB_ID; fold it into RUN_ID so the +# checkpoint dir and log filename are traceable back to the SLURM job. +SLURM_JOB_ID="${SLURM_JOB_ID:-${SLURM_JOBID:-}}" +if [[ -n "$SLURM_JOB_ID" ]]; then + RUN_ID="${RUN_ID:-${SLURM_JOB_ID}}" +else + RUN_ID="${RUN_ID:-$(date +%Y%m%d-%H%M%S)}" +fi + +### Machine-specific args +NGPU="${NGPU:-8}" +HF_HOME_DIR="${HF_HOME_DIR:-/shared_inference/alirezak/hf_home}" # HF model location +DATA_DIR="${DATA_DIR:-/shared_inference/alirezak/hf_home/data}" # expose data dir to container +HF_ENV_FILE="${HF_ENV_FILE:-$HOME/.hf.env}" # .env file has raw HF access token + +### Run-specific args +ALTO_DIR="${ALTO_DIR:-$PWD}" # assume repo dir is pwd +CONFIG="${CONFIG:-gpt_oss_20b_mxfp4_base}" + +CHECKPOINT_DIR="${CHECKPOINT_DIR:-/shared_inference/alirezak/gptoss_chkpt/${USER}/${CONFIG}_$RUN_ID}" +LOG_FILE_DEFAULT="${ALTO_DIR}/logs/${CONFIG}_$(date +%Y%m%d_%H%M%S).log" +LOG_FILE="${LOG_FILE:-$LOG_FILE_DEFAULT}" # log fname based on time + +### Other modifiable args +MODULE="${MODULE:-gpt_oss}" +TRAINING_STEPS="${TRAINING_STEPS:-15000}" +CONTAINER="${CONTAINER:-${CONFIG}_${RUN_ID}}" # container name (for user readability) + +# default Docker image from Han Wang +IMAGE="${IMAGE:-wanghanthu/torchtitan:ubuntu22.04-pytorch2.12.0dev20260217-rocm7.2-patch}" + +# ----------------------------------------------------------------------------- +# Docker Setup +# ----------------------------------------------------------------------------- + +mkdir -p \ + "$HF_HOME_DIR" \ + "$CHECKPOINT_DIR" \ + "$(dirname "$LOG_FILE")" + +echo "=== ALTO GPT-OSS 20B ===" +echo "Node: $(hostname)" +[[ -n "$SLURM_JOB_ID" ]] && echo "SLURM job: $SLURM_JOB_ID" +echo "Image: $IMAGE" +echo "Config: $CONFIG" +echo "GPUs: $NGPU" +echo "Training steps: $TRAINING_STEPS" +echo "Model directory: $HF_HOME_DIR" +echo "Checkpoints: $CHECKPOINT_DIR" +echo "Log: $LOG_FILE" +echo + +docker pull "$IMAGE" + +docker_args=( + -d + --rm + --name "$CONTAINER" + --user "$(id -u):$(id -g)" + --network host + --ipc host + --cap-add SYS_PTRACE + --shm-size 512G + --security-opt seccomp=unconfined + --env-file "$HF_ENV_FILE" + -v "$HOME:$HOME" + -v "$ALTO_DIR:/alto" + -v "$DATA_DIR:/data" + -v "$HF_HOME_DIR:/hf_home" + -v "$CHECKPOINT_DIR:/checkpoints" + -v /etc/passwd:/etc/passwd:ro + -v /etc/group:/etc/group:ro + -e HOME="$HOME" + -e USER="$(id -un)" + -e HF_HOME=/hf_home + -e HF_DATASETS_CACHE=/hf_home/datasets + -e TRITON_CACHE_DIR=/tmp/triton_cache + -e TORCHINDUCTOR_CACHE_DIR=/tmp/torchinductor_cache + -e PYTHONNOUSERSITE=1 +) + +# Hardware resources are added only when they exist. +DEVICE_PATHS="${DEVICE_PATHS:-/dev/kfd /dev/dri /dev/infiniband}" +DEVICE_GROUPS="${DEVICE_GROUPS:-render video}" + +for device in $DEVICE_PATHS; do + if [[ -e "$device" ]]; then + docker_args+=(--device "$device") + fi +done + +for group in $DEVICE_GROUPS; do + gid="$(getent group "$group" | cut -d: -f3 || true)" + + if [[ -n "$gid" ]]; then + docker_args+=(--group-add "$gid") + fi +done + +# start docker container +docker run "${docker_args[@]}" "$IMAGE" sleep infinity + +# make sure docker is gracefully stopped quickly on exit +cleanup() { + status=$? + + # Prevent cleanup from being triggered recursively. + trap - EXIT INT TERM + + echo + echo "[train] Stopping container $CONTAINER ..." + + docker stop --time 3 "$CONTAINER" >/dev/null 2>&1 || + docker kill "$CONTAINER" >/dev/null 2>&1 || + true + + exit "$status" +} + +trap cleanup EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +# ----------------------------------------------------------------------------- +# Load model and install additional docker dependencies +# ----------------------------------------------------------------------------- +MODEL_DIR="/hf_home/hub/models--openai--gpt-oss-20b/snapshots/6cee5e81ee83917806bbde320786a8fb61efebee" +MODEL_DIR_HOST="${HF_HOME_DIR}/hub/models--openai--gpt-oss-20b/snapshots/6cee5e81ee83917806bbde320786a8fb61efebee" +echo "[model] Ensuring tokenizer is available at $MODEL_DIR_HOST ..." + +if [[ -f "$MODEL_DIR_HOST/tokenizer.json" ]]; then + echo "[model] Tokenizer already present, skipping download." +else + echo "[model] Downloading tokenizer ..." + docker exec "$CONTAINER" \ + hf download openai/gpt-oss-20b \ + --include "tokenizer*" "special_tokens_map.json" "config.json" \ + --local-dir "$MODEL_DIR" +fi + + +echo "[train] Installing dependencies ..." + +docker exec "$CONTAINER" bash -c " + python3 -m pip install -q 'torchao==0.16.0' && + python3 -m pip install -q \ + --no-build-isolation \ + --no-deps \ + -e /alto/3rdparty/torchtitan +" + +# ----------------------------------------------------------------------------- +# Training +# ----------------------------------------------------------------------------- + +echo "[train] Launching $CONFIG for $TRAINING_STEPS steps on $NGPU GPUs ..." + +docker exec \ + -w /alto \ + -e PYTORCH_ALLOC_CONF=expandable_segments:True \ + -e TRANSFORMERS_OFFLINE=1 \ + "$CONTAINER" \ + torchrun \ + --standalone \ + --nproc_per_node "$NGPU" \ + --local-ranks-filter 0 \ + --tee 3 \ + -m alto.train \ + --module "$MODULE" \ + --config "$CONFIG" \ + --training.steps "$TRAINING_STEPS" \ + --comm.init_timeout_seconds 1800 \ + --hf_assets_path "$MODEL_DIR" \ + --dump_folder /checkpoints \ + --profiling.enable_profiling \ + --profiling.profile_freq 1000 \ + --profiling.profiler_warmup 3 \ + --profiling.profiler_active 1 \ + 2>&1 | tee "$LOG_FILE" + +echo "[train] Run complete." \ No newline at end of file diff --git a/scripts/train_gptoss_multinode.sh b/scripts/train_gptoss_multinode.sh new file mode 100755 index 00000000..9608d434 --- /dev/null +++ b/scripts/train_gptoss_multinode.sh @@ -0,0 +1,202 @@ +#!/usr/bin/env bash +#SBATCH -A amd-arad +#SBATCH -p amd-arad-burst +#SBATCH --qos=low +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --time=04:00:00 +#SBATCH --job-name=alto-gptoss20b-multinode +#SBATCH --output=alto-gptoss20b-multinode-%j.out +#SBATCH --requeue +# +# Multi-node ALTO GPT-OSS 20B training. +# +# Runs one Docker container per node (from the locally-built alto:multinode +# image) and launches a single torchrun job spanning all nodes, rendezvousing +# on the first allocated node. +# +# sbatch scripts/train_gptoss_20b.sh +# +# Everything below is env-overridable, e.g.: +# TRAINING_STEPS=50 CONFIG=gpt_oss_debugmodel sbatch scripts/train_gptoss_20b.sh +# +# NOTE: CHECKPOINT_DIR, LOG_DIR, HF_HOME_DIR and ALTO_DIR must live on a +# filesystem shared across all nodes (e.g. $HOME or /shared) so every rank sees +# the same repo, model assets and checkpoints. + +set -euo pipefail + +# ----------------------------------------------------------------------------- +# Configuration +# ----------------------------------------------------------------------------- + +### Machine-specific args +GPUS_PER_NODE="${GPUS_PER_NODE:-8}" +HF_HOME_DIR="${HF_HOME_DIR:-$HOME/.cache/huggingface}" # HF model / dataset cache +DATA_DIR="${DATA_DIR:-/shared_inference}" # data dir exposed to container +HF_ENV_FILE="${HF_ENV_FILE:-$HOME/.hf.env}" # .env file with raw HF token + +### Run-specific args +# *NOTE*: if you cloned multiple copies of this repo, make sure the path below is correct +ALTO_DIR="${ALTO_DIR:-$HOME/lpt_branch/ALTO}" # repo dir exposed to container +CONFIG="${CONFIG:-gpt_oss_20b_pretrain_c4}" +RUN_ID="${RUN_ID:-$(date +%Y%m%d-%H%M%S)}" +CHECKPOINT_DIR="${CHECKPOINT_DIR:-$ALTO_DIR/gptoss_chkpt/${CONFIG}_$RUN_ID}" +LOG_DIR="${LOG_DIR:-$ALTO_DIR/logs}" + +### Other modifiable args +MODULE="${MODULE:-gpt_oss}" +TRAINING_STEPS="${TRAINING_STEPS:-15000}" +MODEL_DIR="${MODEL_DIR:-$HF_HOME_DIR/models/gpt-oss-20b}" + +### Docker image built from Dockerfile.multinode +IMAGE="${IMAGE:-alto:multinode}" +DOCKERFILE="${DOCKERFILE:-Dockerfile.multinode}" + +# ----------------------------------------------------------------------------- +# Build the multinode image on every node +# ----------------------------------------------------------------------------- +# Docker images are local to each node's daemon, so the image must be built on +# every allocated node -- building only on the batch node would leave the other +# nodes unable to find alto:multinode at `docker run` time. + +cd "$ALTO_DIR" + +echo "[build] Building $IMAGE from $DOCKERFILE on all nodes ..." +srun --ntasks-per-node=1 \ + bash -c "cd '$ALTO_DIR' && docker build -f '$DOCKERFILE' -t '$IMAGE' ." + +# ----------------------------------------------------------------------------- +# Rendezvous / bookkeeping +# ----------------------------------------------------------------------------- + +# First allocated node becomes the torchrun rendezvous host. +MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)" +MASTER_PORT="$((20000 + SLURM_JOB_ID % 20000))" + +mkdir -p "$CHECKPOINT_DIR" "$LOG_DIR" "$HF_HOME_DIR" + +echo "=== ALTO GPT-OSS 20B (multinode) ===" +echo "Nodes: $(scontrol show hostnames "$SLURM_JOB_NODELIST" | paste -sd, -)" +echo "Master: ${MASTER_ADDR}:${MASTER_PORT}" +echo "Image: $IMAGE" +echo "Config: $CONFIG" +echo "GPUs per node: $GPUS_PER_NODE" +echo "Total GPUs: $((GPUS_PER_NODE * SLURM_JOB_NUM_NODES))" +echo "Training steps: $TRAINING_STEPS" +echo "Model dir: $MODEL_DIR" +echo "Checkpoints: $CHECKPOINT_DIR" +echo "Logs: $LOG_DIR" +echo + +# ----------------------------------------------------------------------------- +# Fetch tokenizer / model config once (shared FS, so all ranks reuse it) +# ----------------------------------------------------------------------------- +if [[ -f "$MODEL_DIR/tokenizer.json" ]]; then + echo "[model] Tokenizer already present at $MODEL_DIR, skipping download." +else + echo "[model] Downloading tokenizer / config to $MODEL_DIR ..." + docker run --rm \ + --user "$(id -u):$(id -g)" \ + --network host \ + --env-file "$HF_ENV_FILE" \ + -v "$HOME:$HOME" \ + -v "$HF_HOME_DIR:/hf_home" \ + -v /etc/passwd:/etc/passwd:ro \ + -v /etc/group:/etc/group:ro \ + -e HOME="$HOME" \ + -e USER="$(id -un)" \ + -e HF_HOME=/hf_home \ + "$IMAGE" \ + hf download openai/gpt-oss-20b \ + --include "tokenizer*" "special_tokens_map.json" "config.json" \ + --local-dir "$MODEL_DIR" +fi + +# Export everything the per-node srun step needs. +export IMAGE ALTO_DIR HF_HOME_DIR DATA_DIR HF_ENV_FILE +export CONFIG MODULE TRAINING_STEPS MODEL_DIR CHECKPOINT_DIR LOG_DIR RUN_ID +export GPUS_PER_NODE MASTER_ADDR MASTER_PORT + +# ----------------------------------------------------------------------------- +# Launch: one srun task -> one container -> torchrun per node +# ----------------------------------------------------------------------------- +srun --kill-on-bad-exit=1 bash -c ' + set -euo pipefail + + CONTAINER="alto_${SLURM_JOB_ID}_${SLURM_NODEID}" + NODE_LOG="$LOG_DIR/gpt_oss_20b-${RUN_ID}-node${SLURM_NODEID}.log" + + cleanup() { + docker stop --time 10 "$CONTAINER" >/dev/null 2>&1 || true + } + trap cleanup EXIT INT TERM + + # Assemble docker args, adding hardware resources only when present on this node. + docker_args=( + --rm + --name "$CONTAINER" + --user "$(id -u):$(id -g)" + --network host + --ipc host + --shm-size 128g + --ulimit memlock=-1 + --cap-add SYS_PTRACE + --security-opt seccomp=unconfined + --env-file "$HF_ENV_FILE" + -v "$HOME:$HOME" + -v "$ALTO_DIR:/alto" + -v "$DATA_DIR:$DATA_DIR" + -v "$HF_HOME_DIR:/hf_home" + -v "$CHECKPOINT_DIR:$CHECKPOINT_DIR" + -v /shared:/shared + -v /etc/passwd:/etc/passwd:ro + -v /etc/group:/etc/group:ro + -e HOME="$HOME" + -e USER="$(id -un)" + -e HF_HOME=/hf_home + -e HF_DATASETS_CACHE=/hf_home/datasets + -e TRITON_CACHE_DIR=/tmp/triton_cache + -e TORCHINDUCTOR_CACHE_DIR=/tmp/torchinductor_cache + -e PYTHONNOUSERSITE=1 + -e TRANSFORMERS_OFFLINE=1 + -e PYTORCH_ALLOC_CONF=expandable_segments:True + -w /alto + ) + + for device in /dev/kfd /dev/dri /dev/infiniband; do + [[ -e "$device" ]] && docker_args+=(--device "$device") + done + + for group in render video; do + gid="$(getent group "$group" | cut -d: -f3 || true)" + [[ -n "$gid" ]] && docker_args+=(--group-add "$gid") + done + + echo "[node $SLURM_NODEID] launching torchrun on $(hostname) -> $NODE_LOG" + + docker run "${docker_args[@]}" "$IMAGE" \ + torchrun \ + --nnodes "$SLURM_JOB_NUM_NODES" \ + --nproc-per-node "$GPUS_PER_NODE" \ + --node-rank "$SLURM_NODEID" \ + --master-addr "$MASTER_ADDR" \ + --master-port "$MASTER_PORT" \ + --local-ranks-filter 0 \ + --tee 3 \ + -m alto.train \ + --module "$MODULE" \ + --config "$CONFIG" \ + --training.steps "$TRAINING_STEPS" \ + --comm.init_timeout_seconds 1800 \ + --hf_assets_path "$MODEL_DIR" \ + --dump_folder "$CHECKPOINT_DIR" \ + --profiling.enable_profiling \ + --profiling.profile_freq 1000 \ + --profiling.profiler_warmup 3 \ + --profiling.profiler_active 1 \ + 2>&1 | tee "$NODE_LOG" +' + +echo "[train] Multinode run complete." diff --git a/tests/integration/llama3_1b_adahop.sh b/tests/integration/llama3_1b_adahop.sh new file mode 100644 index 00000000..e972d8b4 --- /dev/null +++ b/tests/integration/llama3_1b_adahop.sh @@ -0,0 +1,17 @@ +#!/bin/bash +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +# +# Phase-3 real-model run: Llama-3.2-1B under MXFP4 + AdaHOP two-phase swap. +# 1000 training steps (calibration window controlled by the recipe). +# Mirrors the llama3_1b_lpt smoke but with the AdaHOP recipe — same model, +# same data, same target subset (Linear modules, output excluded). + +SCRIPT_DIR=$(dirname "$0") +cd $SCRIPT_DIR/../.. + +NGPU=${NGPU:-8} \ +MODULE=llama3 \ +CONFIG=llama3_1b_adahop \ +./examples/run.sh "$@" diff --git a/tests/integration/llama3_debugmodel_adahop.sh b/tests/integration/llama3_debugmodel_adahop.sh new file mode 100755 index 00000000..7f7a92c3 --- /dev/null +++ b/tests/integration/llama3_debugmodel_adahop.sh @@ -0,0 +1,17 @@ +#!/bin/bash +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +# +# Phase-3 smoke: llama3 debug model under MXFP4 + AdaHOP two-phase swap. +# Runs 30 calibration steps (Phase A) followed by 5 post-calibration training +# steps (Phase B). Confirms the calibration callback + backward hook + re-swap +# pipeline runs end-to-end without crashing. + +SCRIPT_DIR=$(dirname "$0") +cd $SCRIPT_DIR/../.. + +NGPU=2 \ +MODULE=llama3 \ +CONFIG=llama3_debugmodel_adahop \ +./examples/run.sh diff --git a/tests/integration/llama3_debugmodel_adahop_short.sh b/tests/integration/llama3_debugmodel_adahop_short.sh new file mode 100755 index 00000000..b397c719 --- /dev/null +++ b/tests/integration/llama3_debugmodel_adahop_short.sh @@ -0,0 +1,16 @@ +#!/bin/bash +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +# +# Fast-iteration smoke for AdaHOP two-phase swap: 3 calibration steps +# followed by 5 post-Phase-B training steps. Use to debug the calibration +# → re-swap → frozen-mode pipeline without waiting 30 steps of calibration. + +SCRIPT_DIR=$(dirname "$0") +cd $SCRIPT_DIR/../.. + +NGPU=2 \ +MODULE=llama3 \ +CONFIG=llama3_debugmodel_adahop_short \ +./examples/run.sh diff --git a/tests/integration/llama3_debugmodel_baseline.sh b/tests/integration/llama3_debugmodel_baseline.sh new file mode 100755 index 00000000..dcb8e38b --- /dev/null +++ b/tests/integration/llama3_debugmodel_baseline.sh @@ -0,0 +1,17 @@ +#!/bin/bash +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +# +# Step-1 baseline companion to llama3_debugmodel_lpt.sh: runs the same llama3 +# debug config without the LowPrecisionTrainingModifier, so loss/throughput +# can be A/B'd against the mxfp4 path. + +SCRIPT_DIR=$(dirname "$0") +cd $SCRIPT_DIR/../.. + +NGPU=2 \ +MODULE=llama3 \ +CONFIG=llama3_debugmodel \ +./examples/run.sh \ + --training.steps 10 diff --git a/tests/unittest/adahop/test_adahop_modifier_helpers.py b/tests/unittest/adahop/test_adahop_modifier_helpers.py new file mode 100644 index 00000000..dc7dc048 --- /dev/null +++ b/tests/unittest/adahop/test_adahop_modifier_helpers.py @@ -0,0 +1,113 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""CPU-side tests for AdaHOPModifier's pure-Python helpers. + +The full modifier needs ALTO's triton-backed kernels at import time, but the +helper functions live in a sibling module (``calibration_hooks``) with +stdlib-only deps so they can be tested standalone.""" + +import importlib.util +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_HELPERS_PATH = (Path(__file__).resolve().parents[3] / "alto" / "modifiers" / "lpt" / "adahop_internals" / + "calibration_hooks.py") + + +def _load_helpers(): + name = "alto_adahop_calibration_hooks_under_test" + if name in sys.modules: + return sys.modules[name] + spec = importlib.util.spec_from_file_location(name, _HELPERS_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def helpers(): + return _load_helpers() + + +class _Det: + """Stand-in tensor that echoes a string from .detach().""" + + def __init__(self, v): + self.v = v + + def detach(self): + return self.v + + +def test_forward_pre_hook_writes_x_and_w_patterns(helpers): + # module.weight.data._data is the underlying tensor the hook inspects. + weight = SimpleNamespace(data=SimpleNamespace(_data=_Det("col"))) + module = SimpleNamespace(weight=weight) + modifier = SimpleNamespace(_per_step_patterns=[{}]) + + hook = helpers.make_forward_pre_hook(modifier, "layers.0.wq", lambda t: t) + hook(module, (_Det("row"),)) + assert modifier._per_step_patterns[-1]["layers.0.wq"] == {"x": "row", "w": "col"} + + +def test_forward_pre_hook_noop_when_no_step_dict(helpers): + modifier = SimpleNamespace(_per_step_patterns=[]) + hook = helpers.make_forward_pre_hook(modifier, "layers.0.wq", lambda t: t) + hook(object(), (_Det("row"),)) # early-returns before touching module + assert modifier._per_step_patterns == [] + + +def test_backward_hook_writes_grad_output_pattern(helpers): + + class FakeGrad: + + def detach(self): + return "row" + + modifier = SimpleNamespace(_per_step_patterns=[{}]) + hook = helpers.make_backward_hook(modifier, "layers.0.wq", lambda t: t) + hook(None, None, (FakeGrad(),)) + assert modifier._per_step_patterns[-1]["layers.0.wq"]["grad_output"] == "row" + + +def test_backward_hook_ignores_none_grad(helpers): + modifier = SimpleNamespace(_per_step_patterns=[{}]) + hook = helpers.make_backward_hook(modifier, "l0", lambda t: t) + hook(None, None, (None,)) + assert modifier._per_step_patterns[-1] == {} + + +def test_backward_hook_handles_bare_tensor(helpers): + + class FakeGrad: + + def detach(self): + return "col" + + modifier = SimpleNamespace(_per_step_patterns=[{}]) + hook = helpers.make_backward_hook(modifier, "l0", lambda t: t) + hook(None, None, FakeGrad()) # not a tuple + assert modifier._per_step_patterns[-1]["l0"]["grad_output"] == "col" + + +def test_write_and_load_modes_json(tmp_path, helpers): + path = tmp_path / "modes.json" + aggregated = {"l0": {"x": "row", "w": "col", "grad_output": "none"}} + modes = {"l0": {"forward_y": "hadamard", "backward_gx": "none", "backward_gw": "full_precision"}} + helpers.write_modes_json(str(path), aggregated, modes) + + raw = json.loads(path.read_text()) + assert raw["aggregated_patterns"] == aggregated + assert raw["per_layer_modes"] == modes + + # load_modes_from_json expects the bare {fqn: {slot: mode}} format used + # by the transform_config_path recipe option. + modes_only = tmp_path / "modes_only.json" + modes_only.write_text(json.dumps(modes)) + assert helpers.load_modes_from_json(str(modes_only)) == modes diff --git a/tests/unittest/adahop/test_adahop_numerical_parity.py b/tests/unittest/adahop/test_adahop_numerical_parity.py new file mode 100644 index 00000000..8b307f6b --- /dev/null +++ b/tests/unittest/adahop/test_adahop_numerical_parity.py @@ -0,0 +1,276 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Phase-1 numerical parity test: AdaHOP per-slot modes vs plain MXFP4. + +Motivation +---------- +On identical MI350X hardware, ``gpt_oss_20b_adahop`` trains WORSE than plain +``gpt_oss_20b_lpt`` even though AdaHOP is supposed to be a strict superset +(Hadamard incoherence processing + outlier extraction). The validation-loss +gap *widens* over training, which points at biased gradients rather than a +one-off noisier forward. + +This test isolates the numerical behaviour of the two autograd paths for a +single Linear, so we can see *which output* (``y`` / ``grad_x`` / ``grad_w``) +and *which AdaHOP mode* diverges from the bf16 reference more than the plain +MXFP4 baseline does. + +Suspects being probed (see plan keen-questing-ember): + * S1 — AdaHOP ``hadamard`` mode drops 2D-block weight scaling (forward/y). + * S2 — ``inner_outlier_extract_right`` (dominant backward_gw mode) quantizes + grad_output with NO stochastic rounding → biased grad_w. + * S4 — ``inner_outlier_extract_left`` (backward_gx) SR propagation. + +Interpretation +-------------- +For an *unbiased* quantizer, averaging over many random gradients drives the +mean relative error of grad_w toward ~0 (the ``_bias`` metrics below). A +quantizer that silently drops stochastic rounding will show a grad_w mean-bias +that does NOT shrink with averaging and is materially larger than the plain +MXFP4 baseline's. That is the fingerprint of S2. + +Requires a real GPU (MI350X/CDNA4 for the fused kernels; CDNA3 exercises the +QDQ fallbacks). Skips on CPU-only boxes. +""" + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("triton") + +if not torch.cuda.is_available(): + pytest.skip("AdaHOP numerical parity test requires a GPU", allow_module_level=True) + +try: + from alto.kernels.fp4.mxfp4.mxfp_linear import MXFP4LinearFunction + from alto.modifiers.lpt.adahop_internals.mxfp4_linear_function import ( + MXFP4AdaHOPLinearFunction, + ) + from alto._adahop_bridge import HadamardFactory +except RuntimeError as exc: # triton driver init can fail even with a GPU present + pytest.skip(f"kernel import failed: {exc}", allow_module_level=True) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +# Shapes divisible by 32 (MXFP4 block) and >= OUTLIER_K=64 on the extracted +# axis so outlier extraction has room. gpt_oss_20b wq is 2880 -> 4096. +M, K, N = 2048, 2880, 4096 +DTYPE = torch.bfloat16 +DEVICE = "cuda" + + +def _rel_err(approx: torch.Tensor, ref: torch.Tensor) -> float: + """Frobenius relative error ||approx - ref|| / ||ref||.""" + approx = approx.float() + ref = ref.float() + return (torch.linalg.vector_norm(approx - ref) / + torch.linalg.vector_norm(ref).clamp_min(1e-12)).item() + + +def _mean_bias(approx: torch.Tensor, ref: torch.Tensor) -> float: + """Normalized mean signed error — detects a *systematic* offset (bias) + that stochastic rounding is supposed to cancel. Near 0 = unbiased.""" + approx = approx.float() + ref = ref.float() + return ((approx - ref).mean() / ref.abs().mean().clamp_min(1e-12)).item() + + +def _make_inputs(seed: int): + g = torch.Generator(device=DEVICE).manual_seed(seed) + x = torch.randn(M, K, generator=g, device=DEVICE, dtype=DTYPE) + w = torch.randn(N, K, generator=g, device=DEVICE, dtype=DTYPE) * 0.02 + # Inject genuine column outliers into x so the "col" pattern / outlier + # extraction paths have something real to act on (mirrors what the + # calibration classifier detected on the real run: x=col on most layers). + x[:, ::128] *= 25.0 + grad_y = torch.randn(M, N, generator=g, device=DEVICE, dtype=DTYPE) + return x, w, grad_y + + +def _reference(x, w, grad_y): + """bf16 ground truth for y, grad_x, grad_w.""" + xr = x.detach().float().requires_grad_(True) + wr = w.detach().float().requires_grad_(True) + y = xr @ wr.T + y.backward(grad_y.float()) + return y.detach(), xr.grad.detach(), wr.grad.detach() + + +def _run_baseline(x, w, grad_y): + """Plain MXFP4 path with the recipe's flags (2dblock_w, hadamard, sr_grad).""" + xin = x.detach().clone().requires_grad_(True) + win = w.detach().clone().requires_grad_(True) + with torch.no_grad(): + ht = HadamardFactory.create_transform(device=win.device) + y = MXFP4LinearFunction.apply( + xin, win, + False, # use_2dblock_x + True, # use_2dblock_w + True, # use_sr_grad + False, # use_dge + "none", # clip_mode + False, # use_macro_block_scaling + ht, + None, # module_id + ) + y.backward(grad_y) + return y.detach(), xin.grad.detach(), win.grad.detach() + + +def _run_adahop(x, w, grad_y, fy, gx, gw, use_sr_grad=True): + """AdaHOP path with explicit per-slot modes and SR toggle.""" + xin = x.detach().clone().requires_grad_(True) + win = w.detach().clone().requires_grad_(True) + with torch.no_grad(): + ht = HadamardFactory.create_transform(device=win.device) + y = MXFP4AdaHOPLinearFunction.apply( + xin, win, + use_sr_grad, + ht, + fy, gx, gw, + ) + y.backward(grad_y) + return y.detach(), xin.grad.detach(), win.grad.detach() + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + +def test_adahop_none_matches_baseline(): + """Sanity: AdaHOP with all slots 'none' must track plain MXFP4 closely. + A large gap here means the AdaHOP autograd Function has a plumbing bug + independent of any transform (rules S1/S2 in vs out).""" + x, w, grad_y = _make_inputs(seed=0) + y_ref, gx_ref, gw_ref = _reference(x, w, grad_y) + + y_b, gx_b, gw_b = _run_baseline(x, w, grad_y) + y_a, gx_a, gw_a = _run_adahop(x, w, grad_y, "none", "none", "none") + + for name, b, a in (("y", y_b, y_a), ("grad_x", gx_b, gx_a), ("grad_w", gw_b, gw_a)): + eb, ea = _rel_err(b, y_ref if name == "y" else (gx_ref if name == "grad_x" else gw_ref)), \ + _rel_err(a, y_ref if name == "y" else (gx_ref if name == "grad_x" else gw_ref)) + print(f"[none] {name:7s} baseline_relerr={eb:.4f} adahop_relerr={ea:.4f}") + # AdaHOP 'none' should not be materially worse than baseline. + assert ea <= eb * 1.5 + 0.02, ( + f"AdaHOP none-mode {name} rel-err {ea:.4f} >> baseline {eb:.4f}: " + "plumbing bug independent of transforms") + + +def test_grad_w_bias_across_modes(): + """Core S2 probe. Averages grad_w over many random gradients and reports + the *mean signed bias* per backward_gw mode. Unbiased quantization (SR + on) → bias shrinks toward 0. If 'inner_outlier_extract_right' shows a + persistent bias much larger than the baseline / 'hadamard' modes, SR was + silently dropped on that path (S2).""" + n_avg = 16 + modes = ["hadamard", "inner_outlier_extract_right", "full_precision"] + + baseline_bias = 0.0 + acc_bias = {m: 0.0 for m in modes} + acc_relerr = {m: 0.0 for m in modes} + baseline_relerr = 0.0 + + for i in range(n_avg): + x, w, grad_y = _make_inputs(seed=100 + i) + _, _, gw_ref = _reference(x, w, grad_y) + + _, _, gw_b = _run_baseline(x, w, grad_y) + baseline_bias += _mean_bias(gw_b, gw_ref) / n_avg + baseline_relerr += _rel_err(gw_b, gw_ref) / n_avg + + for m in modes: + # Keep the forward + gx legs on 'hadamard' so we isolate backward_gw. + _, _, gw_a = _run_adahop(x, w, grad_y, "hadamard", "hadamard", m) + acc_bias[m] += _mean_bias(gw_a, gw_ref) / n_avg + acc_relerr[m] += _rel_err(gw_a, gw_ref) / n_avg + + print(f"\n[grad_w bias over {n_avg} draws]") + print(f" baseline(mxfp4,sr) bias={baseline_bias:+.5f} relerr={baseline_relerr:.4f}") + for m in modes: + print(f" gw={m:28s} bias={acc_bias[m]:+.5f} relerr={acc_relerr[m]:.4f}") + + # The dominant real-run mode. If its |bias| is an order of magnitude worse + # than baseline, that's the widening-gap culprit (S2). + ioe = abs(acc_bias["inner_outlier_extract_right"]) + base = abs(baseline_bias) + print(f" => inner_outlier_extract_right |bias|={ioe:.5f} vs baseline |bias|={base:.5f}") + # Not a hard assert on the ratio (we want the number reported even when it + # passes); flag only an egregious systematic bias. + assert ioe < 0.05, ( + f"inner_outlier_extract_right grad_w has a large systematic bias " + f"({ioe:.5f}); stochastic rounding likely dropped on this path (S2)") + + +def test_forward_y_hadamard_vs_baseline(): + """S1 probe: compare forward-y rel-err of AdaHOP 'hadamard' mode against + the baseline (which keeps 2D-block weight scaling). AdaHOP should be + <= baseline; if it's worse, the 1D-only iht_quantization is a real + forward precision downgrade.""" + x, w, grad_y = _make_inputs(seed=7) + y_ref, _, _ = _reference(x, w, grad_y) + + y_b, _, _ = _run_baseline(x, w, grad_y) + y_a, _, _ = _run_adahop(x, w, grad_y, "hadamard", "hadamard", "hadamard") + + eb = _rel_err(y_b, y_ref) + ea = _rel_err(y_a, y_ref) + print(f"\n[forward y] baseline_relerr={eb:.4f} adahop_hadamard_relerr={ea:.4f}") + assert ea <= eb * 1.5 + 0.02, ( + f"AdaHOP hadamard forward-y rel-err {ea:.4f} >> baseline {eb:.4f} (S1)") + + +def _which_extract_branch(): + """Report which inner_outlier_extract branch this GPU runs, so the test + output states what was actually verified (CDNA3 QDQ vs CDNA4 fused).""" + try: + from alto.kernels.fp4.mxfp4.mxfp_quantization import is_cdna4 + return "CDNA4 (fused/dot_scaled)" if is_cdna4() else "CDNA3 (pytorch QDQ)" + except Exception: + return "unknown" + + +def test_sr_reduces_grad_w_bias_outlier_right(): + """The SR fix: threading use_sr into inner_outlier_extract_right's grad + quantization should make grad_w UNBIASED (mean signed error -> 0 with + averaging). With SR off (the pre-fix behavior), grad_w carries a persistent + systematic bias. This directly verifies the fix on whatever branch this GPU + runs (CDNA3 pytorch-QDQ or CDNA4 fused) — both were patched. + + backward_gw = inner_outlier_extract_right is the dominant real-run mode + (86/96 layers in the calibrated run), so this is the path that matters. + """ + branch = _which_extract_branch() + n_avg = 32 + bias_sr_on = 0.0 + bias_sr_off = 0.0 + relerr_sr_on = 0.0 + for i in range(n_avg): + x, w, grad_y = _make_inputs(seed=500 + i) + _, _, gw_ref = _reference(x, w, grad_y) + # Isolate backward_gw = inner_outlier_extract_right; keep fwd + gx on hadamard. + _, _, gw_on = _run_adahop(x, w, grad_y, "hadamard", "hadamard", + "inner_outlier_extract_right", use_sr_grad=True) + _, _, gw_off = _run_adahop(x, w, grad_y, "hadamard", "hadamard", + "inner_outlier_extract_right", use_sr_grad=False) + bias_sr_on += _mean_bias(gw_on, gw_ref) / n_avg + bias_sr_off += _mean_bias(gw_off, gw_ref) / n_avg + relerr_sr_on += _rel_err(gw_on, gw_ref) / n_avg + + print(f"\n[SR fix / grad_w, backward_gw=inner_outlier_extract_right, {branch}]") + print(f" SR on : mean_bias={bias_sr_on:+.6f} relerr={relerr_sr_on:.4f}") + print(f" SR off: mean_bias={bias_sr_off:+.6f}") + print(f" => |bias| SR on={abs(bias_sr_on):.6f} SR off={abs(bias_sr_off):.6f}") + + # The SR fix must actually change the result (proves use_sr is threaded through + # to the kernel's grad quantization on this branch) AND make it less biased. + assert abs(bias_sr_on) != abs(bias_sr_off), ( + "SR on/off produced identical grad_w bias -> use_sr is NOT reaching the " + f"outlier kernel's grad quant on {branch}") + assert abs(bias_sr_on) <= abs(bias_sr_off) + 1e-6, ( + f"SR did not reduce grad_w bias on {branch}: " + f"|bias| on={abs(bias_sr_on):.6f} off={abs(bias_sr_off):.6f}") diff --git a/tests/unittest/adahop/test_adahop_wrapper.py b/tests/unittest/adahop/test_adahop_wrapper.py new file mode 100644 index 00000000..cb8b4602 --- /dev/null +++ b/tests/unittest/adahop/test_adahop_wrapper.py @@ -0,0 +1,118 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""CPU-side smoke tests for the AdaHOP wrapper subclasses. + +These tests construct the wrappers and exercise the flatten / unflatten +round-trip without firing any MXFP4 kernel. Running on a CPU-only box +without triton drivers — full dispatch tests live in cluster integration. +""" + +import importlib.util +import sys +from pathlib import Path + +import pytest +import torch + +_ALTO_ROOT = Path(__file__).resolve().parents[3] + + +def _import_alto_kernel_free(name: str, path_parts: list[str]): + """Import a single ALTO module by file path so ``alto/__init__.py`` (which + transitively loads triton kernels) does not run.""" + path = _ALTO_ROOT.joinpath(*path_parts) + qualified = f"alto_under_test_{name}" + if qualified in sys.modules: + return sys.modules[qualified] + spec = importlib.util.spec_from_file_location(qualified, path) + module = importlib.util.module_from_spec(spec) + sys.modules[qualified] = module + spec.loader.exec_module(module) + return module + + +pytest.importorskip("triton") # adahop_tensor imports through alto.kernels.fp4 -> triton + +try: + from alto.kernels.dispatch.adahop_tensor import MXFP4AdaHOPWrapper + from alto.kernels.dispatch.config import TrainingOpConfig +except RuntimeError as exc: + # ALTO's kernels init triton at import; on a no-GPU box that raises + # "0 active drivers". Skip the whole module — these tests need a GPU box. + pytest.skip(f"alto import requires triton driver: {exc}", allow_module_level=True) + + +@pytest.fixture +def mxfp4_config(): + return TrainingOpConfig( + precision="mxfp4", + use_2dblock_x=False, + use_2dblock_w=True, + use_hadamard=False, + use_sr_grad=False, + use_dge=False, + ) + + +def test_wrapper_defaults_to_calibrating_none_modes(mxfp4_config): + t = torch.randn(8, 8) + w = MXFP4AdaHOPWrapper(t, mxfp4_config) + # Default (Phase-A) state: all modes "none" == plain MXFP4. + assert w._forward_y_mode == "none" + assert w._backward_gx_mode == "none" + assert w._backward_gw_mode == "none" + assert w.is_calibrating is True + + +def test_wrapper_set_modes_in_place(mxfp4_config): + t = torch.randn(8, 8) + w = MXFP4AdaHOPWrapper(t, mxfp4_config) + w.set_modes( + forward_y_mode="hadamard", + backward_gx_mode="none", + backward_gw_mode="full_precision", + ) + assert w._forward_y_mode == "hadamard" + assert w._backward_gw_mode == "full_precision" + assert w.is_calibrating is False + + +def test_adahop_wrapper_rejects_unsupported_mode(mxfp4_config): + t = torch.randn(8, 8) + with pytest.raises(NotImplementedError): + MXFP4AdaHOPWrapper(t, mxfp4_config, forward_y_mode="inner_outlier_extract_left") + + +def test_adahop_wrapper_rejects_garbage_mode(mxfp4_config): + t = torch.randn(8, 8) + with pytest.raises(ValueError): + MXFP4AdaHOPWrapper(t, mxfp4_config, forward_y_mode="not_a_mode") + + +def test_adahop_wrapper_flatten_unflatten_round_trips_modes(mxfp4_config): + t = torch.randn(8, 8) + w = MXFP4AdaHOPWrapper( + t, + mxfp4_config, + hadamard_transform=None, + forward_y_mode="hadamard", + backward_gx_mode="outer_hadamard", + backward_gw_mode="full_precision", + ) + inner_names, meta = w.__tensor_flatten__() + assert inner_names == ["_data"] + assert meta["forward_y_mode"] == "hadamard" + assert meta["backward_gx_mode"] == "outer_hadamard" + assert meta["backward_gw_mode"] == "full_precision" + + rebuilt = MXFP4AdaHOPWrapper.__tensor_unflatten__( + {"_data": t}, + meta, + outer_size=t.size(), + outer_stride=t.stride(), + ) + assert rebuilt._forward_y_mode == "hadamard" + assert rebuilt._backward_gx_mode == "outer_hadamard" + assert rebuilt._backward_gw_mode == "full_precision" + assert rebuilt._hadamard_transform is None # not serialized; modifier rebinds on load diff --git a/tests/unittest/adahop/test_bridge_isolation.py b/tests/unittest/adahop/test_bridge_isolation.py new file mode 100644 index 00000000..3b2dd59a --- /dev/null +++ b/tests/unittest/adahop/test_bridge_isolation.py @@ -0,0 +1,68 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Verify ``alto._adahop_bridge`` loads AdaHOP modules without shadowing +ALTO's own ``torchtitan`` submodule.""" + +import importlib.util +import sys +from pathlib import Path + +import pytest + +BRIDGE_PATH = Path(__file__).resolve().parents[3] / "alto" / "_adahop_bridge.py" + + +def _load_bridge_isolated(): + """Load the bridge by file path so this test can run on CPU-only boxes + where ``import alto`` fails on triton driver initialization.""" + spec = importlib.util.spec_from_file_location("alto_adahop_bridge_under_test", BRIDGE_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def bridge(): + pre = sys.modules.get("torchtitan") + module = _load_bridge_isolated() + yield module, pre + + +def test_bridge_does_not_pollute_torchtitan_namespace(bridge): + _, pre = bridge + post = sys.modules.get("torchtitan") + assert post is pre, ("Bridge inserted or replaced sys.modules['torchtitan']; " + "AdaHOP's torchtitan would shadow ALTO's submodule.") + + +def test_bridge_exports_callables(bridge): + module, _ = bridge + for name in [ + "HadamardFactory", + "HadamardTransform", + "detect_outlier_pattern", + "configure_global_transforms", + "configure_layer_transforms", + "get_layer_transform_config", + "should_apply_transform", + "clear_all_configs", + ]: + assert hasattr(module, name), f"bridge missing export: {name}" + + +def test_transform_config_registry_roundtrip(bridge): + module, _ = bridge + module.clear_all_configs() + module.configure_layer_transforms({ + "layers.0.attention.wq": { + "forward_y": "hadamard", + "backward_gw": "none", + "backward_gx": "none", + }, + }) + assert module.get_layer_transform_config("layers.0.attention.wq", "forward_y") == "hadamard" + assert module.should_apply_transform("layers.0.attention.wq", "forward_y") is True + assert module.should_apply_transform("layers.0.attention.wq", "backward_gw") is False + assert module.get_layer_transform_config("unknown.layer", "forward_y") == "none" + module.clear_all_configs() diff --git a/tests/unittest/adahop/test_calibration_hooks.py b/tests/unittest/adahop/test_calibration_hooks.py new file mode 100644 index 00000000..bb83f776 --- /dev/null +++ b/tests/unittest/adahop/test_calibration_hooks.py @@ -0,0 +1,78 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Tests for the JSON round-trip schemas in calibration_hooks.""" + +import importlib.util +import json +import sys +from pathlib import Path + +import pytest + +_HELPERS_PATH = (Path(__file__).resolve().parents[3] / "alto" / "modifiers" / "lpt" / "adahop_internals" / + "calibration_hooks.py") + + +def _load_helpers(): + name = "alto_adahop_calibration_hooks_under_test" + if name in sys.modules: + return sys.modules[name] + spec = importlib.util.spec_from_file_location(name, _HELPERS_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def helpers(): + return _load_helpers() + + +def test_load_accepts_bare_form(tmp_path, helpers): + path = tmp_path / "bare.json" + bare = {"l0": {"forward_y": "hadamard", "backward_gx": "none", "backward_gw": "full_precision"}} + path.write_text(json.dumps(bare)) + assert helpers.load_modes_from_json(str(path)) == bare + + +def test_load_accepts_wrapped_form(tmp_path, helpers): + """``write_modes_json`` output (with ``aggregated_patterns`` + + ``per_layer_modes``) must be loadable as a modes-only dict.""" + path = tmp_path / "wrapped.json" + aggregated = {"l0": {"x": "row", "w": "col", "grad_output": "none"}} + modes = {"l0": {"forward_y": "hadamard", "backward_gx": "none", "backward_gw": "full_precision"}} + helpers.write_modes_json(str(path), aggregated, modes) + assert helpers.load_modes_from_json(str(path)) == modes + + +def test_write_then_load_round_trips(tmp_path, helpers): + """End-to-end: dump → load yields the same per-layer modes.""" + path = tmp_path / "rt.json" + aggregated = { + "layers.0.attention.wq": { + "x": "row", + "w": "row", + "grad_output": "col" + }, + "layers.1.feed_forward.w1": { + "x": "none", + "w": "none", + "grad_output": "none" + }, + } + modes = { + "layers.0.attention.wq": { + "forward_y": "hadamard", + "backward_gx": "hadamard", + "backward_gw": "hadamard" + }, + "layers.1.feed_forward.w1": { + "forward_y": "none", + "backward_gx": "none", + "backward_gw": "none" + }, + } + helpers.write_modes_json(str(path), aggregated, modes) + assert helpers.load_modes_from_json(str(path)) == modes diff --git a/tests/unittest/adahop/test_pattern_aggregation.py b/tests/unittest/adahop/test_pattern_aggregation.py new file mode 100644 index 00000000..82405d99 --- /dev/null +++ b/tests/unittest/adahop/test_pattern_aggregation.py @@ -0,0 +1,128 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""CPU-side tests for adahop_internals.pattern_aggregation.""" + +import importlib.util +import sys +from pathlib import Path + +import pytest + +_ADAHOP_INTERNALS = Path(__file__).resolve().parents[3] / "alto" / "modifiers" / "lpt" / "adahop_internals" + + +def _load(name: str): + """Load a module from adahop_internals by file path so this test can run + on CPU-only boxes where ``import alto`` triggers triton driver init.""" + path = _ADAHOP_INTERNALS / f"{name}.py" + spec = importlib.util.spec_from_file_location(f"alto_adahop_test_{name}", path) + module = importlib.util.module_from_spec(spec) + # Some adahop_internals files use relative imports (`from .transform_mode import ...`). + # Register the package + sibling so relative resolution works. + pkg_name = "alto_adahop_internals_under_test" + if pkg_name not in sys.modules: + pkg_spec = importlib.util.spec_from_file_location( + pkg_name, + _ADAHOP_INTERNALS / "__init__.py", + submodule_search_locations=[str(_ADAHOP_INTERNALS)], + ) + pkg = importlib.util.module_from_spec(pkg_spec) + sys.modules[pkg_name] = pkg + pkg_spec.loader.exec_module(pkg) + # Re-create spec under the package + qualified = f"{pkg_name}.{name}" + if qualified in sys.modules: + return sys.modules[qualified] + spec = importlib.util.spec_from_file_location(qualified, path) + module = importlib.util.module_from_spec(spec) + sys.modules[qualified] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def aggregation(): + return _load("pattern_aggregation") + + +@pytest.fixture(scope="module") +def transform_mode(): + return _load("transform_mode") + + +def test_aggregate_picks_majority_per_tensor(aggregation): + per_step = [ + { + "layers.0.wq": { + "x": "row", + "w": "col", + "grad_output": "none" + } + }, + { + "layers.0.wq": { + "x": "row", + "w": "col", + "grad_output": "row" + } + }, + { + "layers.0.wq": { + "x": "col", + "w": "col", + "grad_output": "row" + } + }, + ] + out = aggregation.aggregate_patterns(per_step) + assert out == {"layers.0.wq": {"x": "row", "w": "col", "grad_output": "row"}} + + +def test_aggregate_missing_layers_and_tensors_default_to_none(aggregation): + per_step = [ + { + "layers.0.wq": { + "x": "row" + } + }, + { + "layers.1.wk": { + "w": "col" + } + }, + ] + out = aggregation.aggregate_patterns(per_step) + assert out["layers.0.wq"] == {"x": "row", "w": "none", "grad_output": "none"} + assert out["layers.1.wk"] == {"x": "none", "w": "col", "grad_output": "none"} + + +def test_patterns_to_modes_t1_t2_mapping(aggregation): + # x=row, w=col → T1=row, T2=opposite(col)=row → "row-row" + aggregated = {"l0": {"x": "row", "w": "col", "grad_output": "none"}} + cfg = {"row-row": "hadamard", "none-none": "none"} + out = aggregation.patterns_to_modes(aggregated, cfg) + assert out["l0"]["forward_y"] == "hadamard" + + +def test_patterns_to_modes_unknown_pair_falls_back_to_none(aggregation): + aggregated = {"l0": {"x": "row", "w": "row", "grad_output": "row"}} + cfg = {"row-row": "hadamard"} + # T4 = opposite(row) = col, T5 = row → "col-row" not in cfg → "none" + out = aggregation.patterns_to_modes(aggregated, cfg) + assert out["l0"]["backward_gw"] == "none" + + +def test_assert_mode_supported_accepts_basic_four(transform_mode): + for mode in ["none", "hadamard", "outer_hadamard", "full_precision"]: + transform_mode.assert_mode_supported(mode, "forward_y") + + +def test_assert_mode_supported_rejects_unported_mode(transform_mode): + with pytest.raises(NotImplementedError): + transform_mode.assert_mode_supported("inner_outlier_extract_left", "forward_y") + + +def test_assert_mode_supported_rejects_garbage(transform_mode): + with pytest.raises(ValueError): + transform_mode.assert_mode_supported("not_a_mode", "backward_gx") diff --git a/tests/unittest/debug/__init__.py b/tests/unittest/debug/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unittest/debug/test_debug_observer_modifier.py b/tests/unittest/debug/test_debug_observer_modifier.py new file mode 100644 index 00000000..09ea4f26 --- /dev/null +++ b/tests/unittest/debug/test_debug_observer_modifier.py @@ -0,0 +1,317 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""End-to-end tests for DebugObserverModifier. + +These tests mock the alto package imports that trigger the triton driver so +they can run on a CPU-only box, following the pattern established in +tests/unittest/adahop/test_adahop_modifier_helpers.py. + +The modifier's external dependencies (alto.modifiers.Modifier, HooksMixin, +match_named_modules, logger) are stubbed at the sys.modules level before the +module is exec'd. +""" + +import importlib.util +import os +import sys +import tempfile +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch +import torch.nn as nn + + +# --------------------------------------------------------------------------- +# Stub the heavy alto / torchtitan imports before loading the modifier module +# --------------------------------------------------------------------------- + +def _install_stubs(): + """Install minimal stubs so debug_observer.py can be loaded without GPU.""" + if "_alto_stubs_installed" in sys.modules: + return + + # --- torchtitan.tools.logging --- + tt_tools = ModuleType("torchtitan.tools") + tt_logging = ModuleType("torchtitan.tools.logging") + tt_logging.logger = MagicMock() + sys.modules.setdefault("torchtitan", ModuleType("torchtitan")) + sys.modules.setdefault("torchtitan.tools", tt_tools) + sys.modules["torchtitan.tools.logging"] = tt_logging + + # --- compressed_tensors.utils.match_named_modules --- + def _match_named_modules(model, targets, ignore): + """Simple FQN walker that matches by class name.""" + ignore_patterns = [] + for ig in ignore: + ignore_patterns.append(ig.lstrip("re:")) + + for fqn, module in model.named_modules(): + if not fqn: + continue + cls_name = module.__class__.__name__ + if cls_name not in targets and "Linear" not in targets: + continue + matched = any(cls_name == t or cls_name.endswith(t) for t in targets) + if not matched: + continue + if any(p in fqn for p in ignore_patterns): + continue + yield fqn, module + + ct_utils = ModuleType("compressed_tensors.utils") + ct_utils.match_named_modules = _match_named_modules + sys.modules.setdefault("compressed_tensors", ModuleType("compressed_tensors")) + sys.modules["compressed_tensors.utils"] = ct_utils + + # --- alto.modifiers (Modifier base class) --- + # We need a real Pydantic base class so PrivateAttr and Field work. + from pydantic import BaseModel, PrivateAttr, Field + from torch.utils.hooks import RemovableHandle + + class _FakeHooksMixin(BaseModel): + model_config = {"extra": "forbid"} + index: int | None = None + group: str | None = None + start: float | None = None + end: float | None = None + update: float | None = None + initialized_: bool = False + finalized_: bool = False + started_: bool = False + ended_: bool = False + _hooks: set = PrivateAttr(default_factory=set) + + def register_hook(self, target, hook, hook_type, **kwargs): + if hook_type in ("forward_pre", "forward", "full_backward"): + handle = getattr(target, f"register_{hook_type}_hook")(hook, **kwargs) + else: + handle = getattr(target, f"register_{hook_type}_hook")(hook, **kwargs) + self._hooks.add(handle) + return handle + + def remove_hooks(self, handles=None): + if handles is None: + handles = set(self._hooks) + for h in handles: + h.remove() + self._hooks -= handles + + class _FakeModifier(_FakeHooksMixin): + + @property + def initialized(self): + return self.initialized_ + + @property + def requires_training_mode(self): + return False + + def initialize(self, model_parts, **kwargs): + self.initialized_ = self.on_initialize(model_parts, **kwargs) + + def finalize(self, model_parts, **kwargs): + self.finalized_ = self.on_finalize(model_parts, **kwargs) + + def pre_step(self, model_parts, **kwargs): + self.started_ = self.on_pre_step(model_parts, **kwargs) + + def post_step(self, model_parts, **kwargs): + self.ended_ = self.on_post_step(model_parts, **kwargs) + + def convert(self, model, **kwargs): + return self.on_convert(model, **kwargs) + + def on_initialize(self, model_parts, **kwargs): + raise NotImplementedError + def on_finalize(self, model_parts, **kwargs): + raise NotImplementedError + def on_pre_step(self, model_parts, **kwargs): + raise NotImplementedError + def on_post_step(self, model_parts, **kwargs): + raise NotImplementedError + def on_convert(self, model, **kwargs): + raise NotImplementedError + + alto_mod = ModuleType("alto") + alto_modifiers = ModuleType("alto.modifiers") + alto_modifiers.Modifier = _FakeModifier + sys.modules["alto"] = alto_mod + sys.modules["alto.modifiers"] = alto_modifiers + + sys.modules["_alto_stubs_installed"] = True + + +def _load_file_as_module(name: str, path: Path) -> ModuleType: + if name in sys.modules: + return sys.modules[name] + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +# --------------------------------------------------------------------------- +# Load the debug_observer module in isolation +# --------------------------------------------------------------------------- + +def _load_debug_observer(): + _install_stubs() + + debug_root = Path(__file__).resolve().parents[3] / "alto" / "modifiers" / "debug" + + # Register observer_hooks so that `from alto.modifiers.debug.observer_hooks import ...` + # inside debug_observer.py resolves correctly. + alto_debug = _load_file_as_module( + "alto.modifiers.debug", debug_root / "__init__.py" + ) + _load_file_as_module( + "alto.modifiers.debug.observer_hooks", debug_root / "observer_hooks.py" + ) + + return _load_file_as_module( + "_debug_observer_under_test", debug_root / "debug_observer.py" + ) + + +@pytest.fixture(scope="module") +def obs_module(): + return _load_debug_observer() + + +@pytest.fixture(scope="module") +def DebugObserverModifier(obs_module): + return obs_module.DebugObserverModifier + + +# --------------------------------------------------------------------------- +# Tiny test model (plain nn.Linear, no triton) +# --------------------------------------------------------------------------- + +class _TinyMLP(nn.Module): + + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(8, 16, bias=False) + self.fc2 = nn.Linear(16, 8, bias=False) + + def forward(self, x): + return self.fc2(torch.relu(self.fc1(x))) + + +def _run_n_steps(model, modifier, n: int): + opt = torch.optim.SGD(model.parameters(), lr=0.01) + model_parts = [model] + for _ in range(n): + modifier.pre_step(model_parts) + x = torch.randn(4, 8, requires_grad=True) + loss = model(x).sum() + loss.backward() + opt.step() + opt.zero_grad() + modifier.post_step(model_parts) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + +class TestDebugObserverModifierLifecycle: + + def _make_modifier(self, DebugObserverModifier, **kwargs): + defaults = dict( + targets=["Linear"], + ignore=[], + capture_every=1, + max_captures=5, + output_path="/tmp/_test_obs.pt", + ) + defaults.update(kwargs) + return DebugObserverModifier(**defaults) + + def test_initialize_registers_layers(self, DebugObserverModifier): + model = _TinyMLP() + mod = self._make_modifier(DebugObserverModifier) + mod.initialize([model]) + assert len(mod._captures) == 2 + for fqn, data in mod._captures.items(): + assert "active" in data + + def test_captures_accumulate_over_steps(self, DebugObserverModifier): + model = _TinyMLP() + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "obs.pt") + mod = self._make_modifier(DebugObserverModifier, output_path=path, max_captures=3) + mod.initialize([model]) + _run_n_steps(model, mod, n=3) + assert mod._n_captured == 3 + + def test_max_captures_hard_cap(self, DebugObserverModifier): + model = _TinyMLP() + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "obs.pt") + mod = self._make_modifier(DebugObserverModifier, output_path=path, max_captures=2) + mod.initialize([model]) + _run_n_steps(model, mod, n=5) + assert mod._n_captured <= 2 + assert mod._detached is True + + def test_finalize_writes_file(self, DebugObserverModifier): + model = _TinyMLP() + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "obs.pt") + mod = self._make_modifier(DebugObserverModifier, output_path=path, max_captures=2) + mod.initialize([model]) + _run_n_steps(model, mod, n=2) + mod.finalize([model]) + assert os.path.exists(path) + + def test_dump_has_expected_schema(self, DebugObserverModifier): + model = _TinyMLP() + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "obs.pt") + mod = self._make_modifier(DebugObserverModifier, output_path=path, max_captures=2) + mod.initialize([model]) + _run_n_steps(model, mod, n=2) + mod.finalize([model]) + blob = torch.load(path, map_location="cpu", weights_only=False) + assert "_meta" in blob + meta = blob["_meta"] + assert "rank" in meta + assert len(meta["iterations_captured"]) == 2 + for fqn, layer_data in blob.items(): + if fqn == "_meta": + continue + for step, tensors in layer_data.items(): + assert "input" in tensors + # grad_weight is only present when MXFP4LinearFunction is in + # the graph (grad_input[1] populated). Plain nn.Linear does + # not populate it, so we only assert it exists if captured. + assert "grad_output" in tensors + + def test_hooks_removed_after_finalize(self, DebugObserverModifier): + model = _TinyMLP() + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "obs.pt") + mod = self._make_modifier(DebugObserverModifier, output_path=path, max_captures=1) + mod.initialize([model]) + _run_n_steps(model, mod, n=1) + mod.finalize([model]) + assert len(mod._hooks) == 0 + assert len(mod._param_hook_handles) == 0 + + def test_capture_every_skips_steps(self, DebugObserverModifier): + model = _TinyMLP() + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "obs.pt") + mod = self._make_modifier( + DebugObserverModifier, output_path=path, capture_every=2, max_captures=10 + ) + mod.initialize([model]) + _run_n_steps(model, mod, n=6) + # capture_every=2, 6 steps → captures at steps 2, 4, 6 + assert mod._n_captured == 3 diff --git a/tests/unittest/debug/test_grad_clip.py b/tests/unittest/debug/test_grad_clip.py new file mode 100644 index 00000000..aba1e5e5 --- /dev/null +++ b/tests/unittest/debug/test_grad_clip.py @@ -0,0 +1,336 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Tests for grad-clip registry, apply_clip helper, and GradientClippingModifier. + +CPU-safe: tests that require GPU (actual MXFP4LinearFunction backward injection) +are skipped when no GPU is available. + +Uses the same importlib stub pattern as test_debug_observer_modifier.py to avoid +triggering the triton driver on import. +""" + +import importlib.util +import sys +from pathlib import Path +from types import ModuleType +from unittest.mock import MagicMock + +import pytest +import torch +import torch.nn as nn + + +# --------------------------------------------------------------------------- +# Stub installation (mirrors test_debug_observer_modifier.py) +# --------------------------------------------------------------------------- + +def _install_stubs(): + if "_alto_stubs_installed" in sys.modules: + return + + tt_tools = ModuleType("torchtitan.tools") + tt_logging = ModuleType("torchtitan.tools.logging") + tt_logging.logger = MagicMock() + sys.modules.setdefault("torchtitan", ModuleType("torchtitan")) + sys.modules.setdefault("torchtitan.tools", tt_tools) + sys.modules["torchtitan.tools.logging"] = tt_logging + + def _match_named_modules(model, targets, ignore): + ignore_patterns = [ig.lstrip("re:") for ig in ignore] + for fqn, module in model.named_modules(): + if not fqn: + continue + cls_name = module.__class__.__name__ + matched = any(cls_name == t or cls_name.endswith(t) for t in targets) + if not matched: + continue + if any(p in fqn for p in ignore_patterns): + continue + yield fqn, module + + ct_utils = ModuleType("compressed_tensors.utils") + ct_utils.match_named_modules = _match_named_modules + sys.modules.setdefault("compressed_tensors", ModuleType("compressed_tensors")) + sys.modules["compressed_tensors.utils"] = ct_utils + + from pydantic import BaseModel, PrivateAttr, Field + + class _FakeModifier(BaseModel): + model_config = {"extra": "forbid"} + index: int | None = None + group: str | None = None + start: float | None = None + end: float | None = None + update: float | None = None + initialized_: bool = False + finalized_: bool = False + + def initialize(self, model_parts, **kwargs): + self.initialized_ = self.on_initialize(model_parts, **kwargs) + + def finalize(self, model_parts, **kwargs): + self.finalized_ = self.on_finalize(model_parts, **kwargs) + + def pre_step(self, model_parts, **kwargs): + return self.on_pre_step(model_parts, **kwargs) + + def post_step(self, model_parts, **kwargs): + return self.on_post_step(model_parts, **kwargs) + + def convert(self, model, **kwargs): + return self.on_convert(model, **kwargs) + + alto_mod = ModuleType("alto") + alto_modifiers = ModuleType("alto.modifiers") + alto_modifiers.Modifier = _FakeModifier + sys.modules["alto"] = alto_mod + sys.modules["alto.modifiers"] = alto_modifiers + + sys.modules["_alto_stubs_installed"] = True + + +def _load_file(name: str, path: Path) -> ModuleType: + if name in sys.modules: + return sys.modules[name] + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _load_grad_clip_modules(): + _install_stubs() + fp4_common = Path(__file__).resolve().parents[3] / "alto" / "kernels" / "fp4" / "fp4_common" + cfg_mod = _load_file("alto.kernels.fp4.fp4_common.grad_clip_config", fp4_common / "grad_clip_config.py") + reg_mod = _load_file("alto.kernels.fp4.fp4_common.grad_clip_registry", fp4_common / "grad_clip_registry.py") + return cfg_mod, reg_mod + + +def _load_modifier_module(): + _install_stubs() + cfg_mod, reg_mod = _load_grad_clip_modules() + + # Ensure the fp4_common package stub is in sys.modules so grad_clip.py can import from it. + fp4_pkg = sys.modules.get("alto.kernels.fp4.fp4_common") + if fp4_pkg is None: + fp4_pkg = ModuleType("alto.kernels.fp4.fp4_common") + sys.modules["alto.kernels.fp4.fp4_common"] = fp4_pkg + fp4_pkg.grad_clip_registry = reg_mod + + alto_kernels = sys.modules.get("alto.kernels", ModuleType("alto.kernels")) + alto_kernels_fp4 = sys.modules.get("alto.kernels.fp4", ModuleType("alto.kernels.fp4")) + sys.modules.setdefault("alto.kernels", alto_kernels) + sys.modules.setdefault("alto.kernels.fp4", alto_kernels_fp4) + + lpt_root = Path(__file__).resolve().parents[3] / "alto" / "modifiers" / "lpt" + return _load_file("_grad_clip_modifier_under_test", lpt_root / "grad_clip.py") + + +@pytest.fixture(scope="module") +def modules(): + cfg_mod, reg_mod = _load_grad_clip_modules() + return cfg_mod, reg_mod + + +@pytest.fixture(scope="module") +def modifier_module(): + return _load_modifier_module() + + +@pytest.fixture(scope="module") +def GradClipConfig(modules): + cfg_mod, _ = modules + return cfg_mod.GradClipConfig + + +@pytest.fixture(scope="module") +def registry(modules): + _, reg_mod = modules + return reg_mod + + +@pytest.fixture(scope="module") +def GradientClippingModifier(modifier_module): + return modifier_module.GradientClippingModifier + + +# --------------------------------------------------------------------------- +# Tiny model for modifier tests +# --------------------------------------------------------------------------- + +class _TinyMLP(nn.Module): + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(8, 16, bias=False) + self.fc2 = nn.Linear(16, 8, bias=False) + + def forward(self, x): + return self.fc2(torch.relu(self.fc1(x))) + + +# --------------------------------------------------------------------------- +# apply_clip tests +# --------------------------------------------------------------------------- + +class TestApplyClip: + + def test_apply_clip_norm(self, registry): + t = torch.randn(32, 32) + t = t * (10.0 / t.norm()) # set norm to 10 + clipped = registry.apply_clip(t, max_norm=0.5, clip_value=None) + assert clipped.norm().item() <= 0.5 + 1e-5 + + def test_apply_clip_value(self, registry): + t = torch.full((4, 4), 5.0) + clipped = registry.apply_clip(t, max_norm=None, clip_value=2.0) + assert clipped.abs().max().item() <= 2.0 + 1e-6 + + def test_apply_clip_both(self, registry): + # Norm applied first, then value clamp. + t = torch.full((4, 4), 5.0) # norm = 5*4 = 20 + clipped = registry.apply_clip(t, max_norm=1.0, clip_value=0.1) + # After norm clip: all elements ≈ 1/16; after value clamp still ≤ 0.1 + assert clipped.abs().max().item() <= 0.1 + 1e-6 + + def test_apply_clip_noop(self, registry): + t = torch.randn(4, 4) + result = registry.apply_clip(t, max_norm=None, clip_value=None) + assert result is t + + +# --------------------------------------------------------------------------- +# Registry tests +# --------------------------------------------------------------------------- + +class TestRegistry: + + def test_register_and_get(self, registry, GradClipConfig): + cfg = GradClipConfig(clip_grad_output=True, grad_output_max_norm=1.0) + fake_id = 999999 + registry.register(fake_id, cfg) + retrieved = registry.get(fake_id) + assert retrieved is cfg + registry.deregister(fake_id) + assert registry.get(fake_id) is None + + def test_get_none_id(self, registry): + assert registry.get(None) is None + + def test_deregister_missing_is_noop(self, registry): + registry.deregister(0) # should not raise + + +# --------------------------------------------------------------------------- +# GradientClippingModifier lifecycle tests +# --------------------------------------------------------------------------- + +class TestGradientClippingModifierLifecycle: + + def _make_modifier(self, GradientClippingModifier, **kwargs): + defaults = dict( + targets=["Linear"], + ignore=[], + clip_grad_output=True, + grad_output_max_norm=1.0, + grad_output_clip_value=None, + clip_grad_weight=True, + grad_weight_max_norm=1.0, + grad_weight_clip_value=None, + ) + defaults.update(kwargs) + return GradientClippingModifier(**defaults) + + def test_initialize_registers_modules(self, GradientClippingModifier, registry): + model = _TinyMLP() + mod = self._make_modifier(GradientClippingModifier) + mod.initialize([model]) + for _fqn, module in model.named_modules(): + if isinstance(module, nn.Linear): + assert registry.get(id(module)) is not None + + def test_forward_pre_hook_stamps_module_id(self, GradientClippingModifier): + """The forward_pre_hook must stamp module_id on the weight each forward. + + At initialize time module_id is NOT set (FSDP drops instance attrs on + each all-gather). Instead, a forward_pre_hook re-stamps it before + __torch_function__ is called. Simulate that by running a forward pass + with a wrapped weight that has module_id=None before the hook fires. + """ + + class _WrappedParam(nn.Parameter): + module_id: int | None = None + + model = _TinyMLP() + stub = _WrappedParam(model.fc1.weight.data) + stub.module_id = None + model.fc1.weight = stub + + mod = self._make_modifier(GradientClippingModifier) + mod.initialize([model]) + + # module_id is still None at this point — that's expected. + assert stub.module_id is None + + # Trigger one forward pass; the pre-hook should stamp module_id. + model(torch.randn(2, 8)) + assert stub.module_id == id(model.fc1) + + def test_finalize_deregisters_all(self, GradientClippingModifier, registry): + model = _TinyMLP() + mod = self._make_modifier(GradientClippingModifier) + mod.initialize([model]) + ids_before = [id(m) for _, m in model.named_modules() if isinstance(m, nn.Linear)] + mod.finalize([model]) + for mid in ids_before: + assert registry.get(mid) is None + assert len(mod._hook_handles) == 0 + + +# --------------------------------------------------------------------------- +# GPU-only: actual backward injection test +# --------------------------------------------------------------------------- + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +class TestClippingBoundsGrad: + + def test_clipping_bounds_grad_weight(self): + """Verify that the clip actually fires in MXFP4LinearFunction.backward. + + This test imports the real kernel stack (triton required) and checks that + grad_weight's abs max is bounded after a forward+backward with a tight clip. + """ + import importlib + mxfp_linear = importlib.import_module("alto.kernels.fp4.mxfp4.mxfp_linear") + grad_clip_registry = importlib.import_module("alto.kernels.fp4.fp4_common.grad_clip_registry") + GradClipConfig = importlib.import_module("alto.kernels.fp4.fp4_common.grad_clip_config").GradClipConfig + + M, N, K = 16, 16, 16 + x = torch.randn(M, K, device="cuda", requires_grad=True) + w = torch.randn(N, K, device="cuda", requires_grad=True) + + clip_value = 1e-3 + fake_module_id = id(w) + cfg = GradClipConfig( + clip_grad_weight=True, + grad_weight_clip_value=clip_value, + ) + grad_clip_registry.register(fake_module_id, cfg) + + try: + y = mxfp_linear._to_mxfp4_then_scaled_mm( + x, w, + use_2dblock_x=False, + use_2dblock_w=True, + use_sr_grad=False, + use_dge=False, + clip_mode="none", + use_hadamard=False, + module_id=fake_module_id, + ) + y.sum().backward() + assert w.grad is not None + assert w.grad.abs().max().item() <= clip_value + 1e-6 + finally: + grad_clip_registry.deregister(fake_module_id) diff --git a/tests/unittest/debug/test_moe_pattern_hooks.py b/tests/unittest/debug/test_moe_pattern_hooks.py new file mode 100644 index 00000000..1813aca3 --- /dev/null +++ b/tests/unittest/debug/test_moe_pattern_hooks.py @@ -0,0 +1,248 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Unit tests for moe_pattern_hooks.py — loaded directly to avoid triton imports. + +moe_pattern_hooks.py is torch-only (no alto/triton imports) so it can be exec'd +in isolation on a CPU-only box, same pattern as test_observer_hooks.py. The +``detect`` callable is injected, so we use the real AdaHOP algorithm reimplemented +inline here to keep the test self-contained (no adahop submodule needed) AND a +deterministic stub to assert the T-pair combination logic exactly. +""" + +import importlib.util +import math +import sys +from pathlib import Path + +import torch +import pytest + +_HOOKS_PATH = (Path(__file__).resolve().parents[3] + / "alto" / "modifiers" / "debug" / "moe_pattern_hooks.py") + + +def _load_hooks(): + name = "_moe_pattern_hooks_under_test" + if name in sys.modules: + return sys.modules[name] + spec = importlib.util.spec_from_file_location(name, _HOOKS_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def hooks(): + return _load_hooks() + + +# --- reference detector (verbatim algorithm from AdaHOP outlier_detection.py) --- + +def _kurtosis(x): + mean, std = x.mean(), x.std() + if std == 0: + return torch.tensor(0.0) + z = (x - mean) / std + return (z ** 4).mean() + + +def detect_ref(X, threshold_ratio=2.0, kurtosis_threshold=0.0): + X = X.float() + if X.dim() > 2: + X = X.reshape(-1, X.shape[-1]) + row_var = X.var(dim=1) + col_var = X.var(dim=0) + cv_row = row_var.std() / (row_var.mean() + 1e-8) + cv_col = col_var.std() / (col_var.mean() + 1e-8) + nr, nc = X.shape + cv_row = cv_row / math.sqrt(2.0 / max(nc - 1, 1)) + cv_col = cv_col / math.sqrt(2.0 / max(nr - 1, 1)) + if cv_row / (cv_col + 1e-8) > threshold_ratio: + return "row" if _kurtosis(row_var) >= kurtosis_threshold else "none" + elif cv_col / (cv_row + 1e-8) > threshold_ratio: + return "col" if _kurtosis(col_var) >= kurtosis_threshold else "none" + return "none" + + +# --------------------------------------------------------------------------- +# offs -> bounds +# --------------------------------------------------------------------------- + +class TestOffsBounds: + def test_basic(self, hooks): + # offs is a cumulative sum; expert 0 = [0,3), expert 1 = [3,7) + assert hooks._offs_to_bounds([3, 7], 2) == [(0, 3), (3, 7)] + + def test_tail_padding_dropped(self, hooks): + # trailing padding beyond offs[-1] is never assigned to an expert + assert hooks._offs_to_bounds([2, 5], 2) == [(0, 2), (2, 5)] + + def test_empty_expert(self, hooks): + # expert 1 gets zero tokens (offs repeats) + assert hooks._offs_to_bounds([4, 4], 2) == [(0, 4), (4, 4)] + + +# --------------------------------------------------------------------------- +# T-pair combination logic (deterministic stub detector) +# --------------------------------------------------------------------------- + +class TestTPairLogic: + def _run(self, hooks, x_pat, w_pat, g_pat): + # stub detect: return a fixed pattern keyed by tensor identity via shape + # We tag tensors by a sentinel scalar in [0,0] and map to the desired pat. + def detect(t): + tag = int(round(t.reshape(-1)[0].item())) + return {0: x_pat, 1: w_pat, 2: g_pat}[tag] + + T = 4 + K, N, E = 3, 5, 1 + x = torch.zeros(T, K); x[0, 0] = 0.0 # tag 0 + # w as fed to _grouped_mm is [E, K, N] + w = torch.zeros(E, K, N); w[0, 0, 0] = 1.0 # tag 1 (any expert slice) + go = torch.zeros(T, N); go[0, 0] = 2.0 # tag 2 + recs = hooks.build_expert_records(x, w, go, offs=[T], detect=detect) + return recs[0] + + def test_all_none(self, hooks): + r = self._run(hooks, "none", "none", "none") + assert r["forward_y"]["pair"] == "none-none" + assert r["backward_gx"]["pair"] == "none-none" + assert r["backward_gw"]["pair"] == "none-none" + + def test_row_row_row(self, hooks): + r = self._run(hooks, "row", "row", "row") + # forward_y = x_pat - opposite(w_pat) = row - col + assert r["forward_y"]["pair"] == "row-col" + # backward_gx = g_pat - w_pat = row - row + assert r["backward_gx"]["pair"] == "row-row" + # backward_gw = opposite(g_pat) - x_pat = col - row + assert r["backward_gw"]["pair"] == "col-row" + + def test_col_none_row(self, hooks): + r = self._run(hooks, "col", "none", "row") + assert r["forward_y"]["pair"] == "col-none" # col - opposite(none) + assert r["backward_gx"]["pair"] == "row-none" # row - none + assert r["backward_gw"]["pair"] == "col-col" # opposite(row) - col + + +# --------------------------------------------------------------------------- +# Per-expert slicing with an injected real outlier +# --------------------------------------------------------------------------- + +class TestPerExpertDetection: + def test_slicing_matches_direct_detect(self, hooks): + """build_expert_records must classify each expert on exactly its own + token slice — verified by re-running the reference detector on the + hand-sliced operands and reconstructing the T-pairs independently.""" + torch.manual_seed(0) + E, K, N = 2, 8, 6 + t0, t1 = 40, 50 + # expert 0: inject strong per-column outliers into the activation + x0 = torch.randn(t0, K); x0[:, 2:4] *= 40.0 + x1 = torch.randn(t1, K) + x = torch.cat([x0, x1], dim=0) + w = torch.randn(E, K, N) + go = torch.randn(t0 + t1, N) + offs = [t0, t0 + t1] + + recs = hooks.build_expert_records(x, w, go, offs, detect_ref) + assert recs[0]["n_tokens"] == t0 + assert recs[1]["n_tokens"] == t1 + + # Independently reconstruct expert 0's forward_y T-pair from its slice. + opp = {"row": "col", "col": "row", "none": "none"} + x0_pat = detect_ref(x[0:t0]) + w0_pat = detect_ref(w[0].transpose(-2, -1)) + expected = f"{x0_pat}-{opp[w0_pat]}" + assert recs[0]["forward_y"]["pair"] == expected + # the injected column outlier should be picked up on the slice + assert x0_pat == "col" + # stats present and finite for a routed expert + assert recs[0]["stats"]["x"]["absmax"] > 0 + + def test_empty_expert_marks_none_activation(self, hooks): + E, K, N = 2, 4, 3 + x = torch.randn(5, K) + w = torch.randn(E, K, N) + go = torch.randn(5, N) + # expert 1 gets zero tokens + recs = hooks.build_expert_records(x, w, go, offs=[5, 5], detect=detect_ref) + assert recs[1]["n_tokens"] == 0 + # activation-derived T1 must be 'none' when no tokens routed + assert recs[1]["forward_y"]["pair"].startswith("none-") + assert recs[1]["backward_gw"]["pair"] == "none-none" + + +# --------------------------------------------------------------------------- +# accumulate_majority +# --------------------------------------------------------------------------- + +class TestAccumulateMajority: + def _rec(self, fy, gx, gw, n_tokens=4): + return { + "n_tokens": n_tokens, + "forward_y": {"pair": fy}, + "backward_gx": {"pair": gx}, + "backward_gw": {"pair": gw}, + } + + def test_clear_majority(self, hooks): + steps = [ + {0: self._rec("row-col", "col-row", "none-none")}, + {0: self._rec("row-col", "col-row", "row-row")}, + {0: self._rec("none-none", "col-row", "col-col")}, + ] + out = hooks.accumulate_majority(steps) + assert out[0]["forward_y"]["pair"] == "row-col" # 2 of 3 + assert out[0]["backward_gx"]["pair"] == "col-row" # 3 of 3 + assert out[0]["n_steps"] == 3 + assert out[0]["n_tokens_total"] == 12 + assert out[0]["forward_y"]["votes"] == {"row-col": 2, "none-none": 1} + assert out[0]["forward_y"]["n"] == 3 + + def test_tie_breaks_deterministically(self, hooks): + # 1 vs 1 tie: winner is the alphabetically-first pair string. + steps = [ + {0: self._rec("row-col", "none-none", "none-none")}, + {0: self._rec("col-row", "none-none", "none-none")}, + ] + out = hooks.accumulate_majority(steps) + assert out[0]["forward_y"]["pair"] == "col-row" + + def test_expert_missing_from_some_steps(self, hooks): + # Expert 1 is only routed in one of the two steps. + steps = [ + {0: self._rec("row-col", "col-row", "none-none"), + 1: self._rec("col-col", "row-row", "none-none")}, + {0: self._rec("row-col", "col-row", "none-none")}, + ] + out = hooks.accumulate_majority(steps) + assert set(out.keys()) == {0, 1} + assert out[0]["n_steps"] == 2 + assert out[1]["n_steps"] == 1 + assert out[1]["forward_y"]["pair"] == "col-col" + + def test_empty_input(self, hooks): + assert hooks.accumulate_majority([]) == {} + + +# --------------------------------------------------------------------------- +# extract_offs +# --------------------------------------------------------------------------- + +class TestExtractOffs: + def test_kwarg(self, hooks): + offs = torch.tensor([2, 4], dtype=torch.int32) + got = hooks.extract_offs((torch.randn(4, 3), torch.randn(1, 3, 5)), {"offs": offs}) + assert torch.equal(got, offs) + + def test_positional_int_tensor(self, hooks): + offs = torch.tensor([2, 4], dtype=torch.int64) + got = hooks.extract_offs((torch.randn(4, 3), torch.randn(1, 3, 5), offs), {}) + assert torch.equal(got, offs) + + def test_missing(self, hooks): + got = hooks.extract_offs((torch.randn(4, 3),), {}) + assert got is None diff --git a/tests/unittest/debug/test_observer_hooks.py b/tests/unittest/debug/test_observer_hooks.py new file mode 100644 index 00000000..f6c423f8 --- /dev/null +++ b/tests/unittest/debug/test_observer_hooks.py @@ -0,0 +1,208 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Unit tests for observer_hooks.py — loaded directly to avoid triton imports. + +observer_hooks.py has no alto or triton imports so it can be exec'd in isolation +on a CPU-only box, following the same pattern as test_adahop_modifier_helpers.py. +""" + +import importlib.util +import sys +from pathlib import Path + +import torch +import torch.nn as nn +import pytest + +_HOOKS_PATH = (Path(__file__).resolve().parents[3] / "alto" / "modifiers" / "debug" / "observer_hooks.py") + + +def _load_hooks(): + name = "_debug_observer_hooks_under_test" + if name in sys.modules: + return sys.modules[name] + spec = importlib.util.spec_from_file_location(name, _HOOKS_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def hooks(): + return _load_hooks() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_captures(fqn: str, active: bool = True) -> dict: + return {fqn: {"active": active}} + + +def _step_ref(step: int = 0) -> list: + return [step] + + +# --------------------------------------------------------------------------- +# nn.Linear forward pre-hook +# --------------------------------------------------------------------------- + +class TestLinearFwdPreHook: + + def test_captures_input(self, hooks): + fqn = "layers.0.linear" + captures = _make_captures(fqn) + step_ref = _step_ref(1) + module = nn.Linear(4, 4, bias=False) + hook = hooks.make_linear_fwd_pre_hook( + captures, fqn, step_ref, capture_input=True, capture_weight=False + ) + x = torch.randn(2, 4, requires_grad=True) + hook(module, (x,)) + assert 1 in captures[fqn] + assert "input" in captures[fqn][1] + assert torch.allclose(captures[fqn][1]["input"], x.detach().cpu()) + + def test_captures_weight(self, hooks): + fqn = "layers.0.linear" + captures = _make_captures(fqn) + step_ref = _step_ref(2) + module = nn.Linear(4, 4, bias=False) + hook = hooks.make_linear_fwd_pre_hook( + captures, fqn, step_ref, capture_input=False, capture_weight=True + ) + hook(module, (torch.randn(2, 4),)) + assert "weight" in captures[fqn][2] + assert captures[fqn][2]["weight"].shape == module.weight.shape + + def test_gate_prevents_capture(self, hooks): + fqn = "layers.0.linear" + captures = _make_captures(fqn, active=False) + step_ref = _step_ref(0) + module = nn.Linear(4, 4, bias=False) + hook = hooks.make_linear_fwd_pre_hook( + captures, fqn, step_ref, capture_input=True, capture_weight=True + ) + hook(module, (torch.randn(2, 4),)) + assert 0 not in captures[fqn] + + def test_captured_tensor_not_requires_grad(self, hooks): + fqn = "l" + captures = _make_captures(fqn) + step_ref = _step_ref(0) + module = nn.Linear(4, 4, bias=False) + hook = hooks.make_linear_fwd_pre_hook( + captures, fqn, step_ref, capture_input=True, capture_weight=False + ) + x = torch.randn(2, 4, requires_grad=True) + hook(module, (x,)) + assert not captures[fqn][0]["input"].requires_grad + + +# --------------------------------------------------------------------------- +# nn.Linear backward hook +# --------------------------------------------------------------------------- + +class TestLinearBwdHook: + + def test_captures_grad_output(self, hooks): + fqn = "layers.0" + captures = _make_captures(fqn) + step_ref = _step_ref(3) + module = nn.Linear(4, 4, bias=False) + hook = hooks.make_linear_bwd_hook( + captures, fqn, step_ref, capture_grad_output=True, capture_grad_weight=False + ) + go = torch.randn(2, 4) + hook(module, (None,), (go,)) + assert "grad_output" in captures[fqn][3] + assert torch.allclose(captures[fqn][3]["grad_output"], go.detach().cpu()) + + def test_captures_grad_weight_from_grad_input(self, hooks): + fqn = "layers.0" + captures = _make_captures(fqn) + step_ref = _step_ref(3) + module = nn.Linear(4, 4, bias=False) + hook = hooks.make_linear_bwd_hook( + captures, fqn, step_ref, capture_grad_output=False, capture_grad_weight=True + ) + gw = torch.randn(4, 4) + # grad_input = (grad_x, grad_weight) — mirrors what autograd.Function.backward returns + hook(module, (torch.randn(2, 4), gw), (torch.randn(2, 4),)) + assert "grad_weight" in captures[fqn][3] + assert torch.allclose(captures[fqn][3]["grad_weight"], gw.detach().cpu()) + + def test_inactive_gate_bwd(self, hooks): + fqn = "l" + captures = _make_captures(fqn, active=False) + step_ref = _step_ref(0) + module = nn.Linear(4, 4, bias=False) + hook = hooks.make_linear_bwd_hook( + captures, fqn, step_ref, capture_grad_output=True, capture_grad_weight=True + ) + hook(module, (None, torch.randn(4, 4)), (torch.randn(2, 4),)) + assert 0 not in captures[fqn] + + +# --------------------------------------------------------------------------- +# GptOssGroupedExperts hooks +# --------------------------------------------------------------------------- + +class _FakeGroupedExperts(nn.Module): + """Minimal stand-in for GptOssGroupedExperts (no triton ops).""" + + def __init__(self): + super().__init__() + self.mlp1_weight = nn.Parameter(torch.randn(2, 8, 4)) + self.mlp2_weight = nn.Parameter(torch.randn(2, 4, 8)) + + def forward(self, x): + return x + + +class TestGroupedExpertsHooks: + + def test_fwd_pre_captures_input(self, hooks): + fqn = "moe.experts" + captures = _make_captures(fqn) + step_ref = _step_ref(1) + module = _FakeGroupedExperts() + hook = hooks.make_grouped_experts_fwd_pre_hook( + captures, fqn, step_ref, capture_input=True, capture_weight=False + ) + x = torch.randn(4, 4) + hook(module, (x,)) + assert "input" in captures[fqn][1] + assert torch.allclose(captures[fqn][1]["input"], x.detach().cpu()) + + def test_fwd_pre_captures_both_mlp_weights(self, hooks): + fqn = "moe.experts" + captures = _make_captures(fqn) + step_ref = _step_ref(1) + module = _FakeGroupedExperts() + hook = hooks.make_grouped_experts_fwd_pre_hook( + captures, fqn, step_ref, capture_input=False, capture_weight=True + ) + hook(module, (torch.randn(4, 4),)) + assert "mlp1_weight" in captures[fqn][1] + assert "mlp2_weight" in captures[fqn][1] + + def test_grad_weight_hooks_stored_under_separate_keys(self, hooks): + fqn = "moe.experts" + captures = _make_captures(fqn) + step_ref = _step_ref(7) + hook1 = hooks.make_grouped_experts_grad_weight_hook( + captures, fqn, step_ref, key="grad_mlp1_weight", capture_grad_weight=True + ) + hook2 = hooks.make_grouped_experts_grad_weight_hook( + captures, fqn, step_ref, key="grad_mlp2_weight", capture_grad_weight=True + ) + g1 = torch.randn(2, 8, 4) + g2 = torch.randn(2, 4, 8) + hook1(g1) + hook2(g2) + assert captures[fqn][7]["grad_mlp1_weight"].shape == g1.shape + assert captures[fqn][7]["grad_mlp2_weight"].shape == g2.shape diff --git a/tests/unittest/mxfp4/test_midmax_quantization.py b/tests/unittest/mxfp4/test_midmax_quantization.py new file mode 100644 index 00000000..32c300e6 --- /dev/null +++ b/tests/unittest/mxfp4/test_midmax_quantization.py @@ -0,0 +1,300 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT + +""" +Unit tests for the midmax scale-selection logic in the mxfp4 quantization kernel. + +Baseline reference: ROCm/tensorcast tcast/number.py _decode(), which defines + emax = 2^ebits - 1 - bias (for E2M1: emax = 2) + maxfloat = 2^emax * (2 - 2^-mbits) (for E2M1: 6.0) + midmax = (2^(emax+1) - maxfloat)/2 + maxfloat (for E2M1: 7.0) + +The kernel uses midmax to decide whether to bump the uint8 scale exponent by 1: + bump if amax_normalized > 7.0 (i.e. amax_normalized > midmax) +""" + +import pytest +import torch + +from alto.kernels.fp4.mxfp4.mxfp_quantization import ( + convert_to_mxfp4, + convert_from_mxfp4, + is_cdna4, +) +from .utils import ( + prepare_data, + convert_to_mxfp4_pytorch, + convert_from_mxfp4_pytorch, +) + +# E2M1 constants (baseline from tensorcast number.py) +_EBITS, _MBITS, _BIAS = 2, 1, 1 +_EMAX = 2**_EBITS - 1 - _BIAS # 2 +_MAXFLOAT = 2**_EMAX * (2.0 - 2**-_MBITS) # 6.0 +_MIDMAX = (2**(_EMAX + 1) - _MAXFLOAT) / 2.0 + _MAXFLOAT # 7.0 +_TARGET_MAX_POW2 = 2 # matches kernel constant + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + +def _make_block(value: float, block_size: int = 32, dtype=torch.float32) -> torch.Tensor: + """Return a 1-D tensor of length block_size filled with `value`, on CUDA.""" + return torch.full((block_size,), value, dtype=dtype, device="cuda") + + +def _midmax_scale_ref(amax: float, block_size: int = 32) -> int: + """ + Pure-Python reference for the midmax uint8 scale given a block's amax. + Mirrors the triton kernel logic in _calculate_scales with SCALE_SELECTION="midmax". + """ + import struct + + amax_f32 = float(amax) + bits = struct.unpack("I", struct.pack("f", abs(amax_f32)))[0] + f32_exp = (bits >> 23) & 0xFF + if f32_exp >= 0xFF: + f32_exp = 0xFE # cap NaN/Inf + + scale = f32_exp - _TARGET_MAX_POW2 + + # normalize amax into [2^target_max_pow2, 2^(target_max_pow2+1)) + mantissa = bits & 0x7FFFFF + norm_bits = mantissa | ((127 + _TARGET_MAX_POW2) << 23) + amax_scaled = struct.unpack("f", struct.pack("I", norm_bits))[0] + + if amax_scaled > _MIDMAX: + scale += 1 + + scale = max(scale, 1) # clamp to minimum normal + return scale + + +# --------------------------------------------------------------------------- +# 1. Baseline constant sanity +# --------------------------------------------------------------------------- + +def test_e2m1_midmax_value(): + """E2M1 midmax must equal 7.0 per the tensorcast baseline formula.""" + assert _MAXFLOAT == 6.0 + assert _MIDMAX == 7.0 + + +# --------------------------------------------------------------------------- +# 2. Scale correctness: midmax vs round-even on crafted inputs +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("amax,expect_bump", [ + (6.0, False), # exactly maxfloat — amax_normalized == 6.0, no bump + (6.9, False), # below midmax — no bump + (7.0, False), # at midmax (not strictly greater) — no bump + (7.1, True), # just above midmax — bump + (8.0, True), # well above midmax — bump + (12.0, True), # large value — bump +]) +def test_midmax_scale_bump(amax, expect_bump): + """ + Verify the pure-Python reference bumps the scale exactly when + amax_normalized > 7.0 (strict greater-than, matching the kernel). + """ + ref_scale = _midmax_scale_ref(amax) + ref_no_bump = _midmax_scale_ref(6.0) + ref_bump = _midmax_scale_ref(7.1) + if expect_bump: + assert ref_scale == ref_bump, f"amax={amax} expected bump, got scale={ref_scale}" + else: + assert ref_scale == ref_no_bump, f"amax={amax} expected no bump, got scale={ref_scale}" + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_midmax_scale_matches_ref_random(dtype): + """ + Triton kernel (blockscale_selection="midmax") scales must match the pure-Python reference + on random data. + """ + torch.manual_seed(42) + x = torch.randn(128, 64, dtype=dtype, device="cuda") + _, scales_triton = convert_to_mxfp4(x, blockscale_selection="midmax") + + # build reference scales block-by-block + x_f32 = x.float() + block_size = 32 + M, N = 128, 64 + scales_ref = torch.zeros(M, N // block_size, dtype=torch.uint8, device="cpu") + for m in range(M): + for nb in range(N // block_size): + block = x_f32[m, nb * block_size:(nb + 1) * block_size] + amax = block.abs().max().item() + scales_ref[m, nb] = _midmax_scale_ref(amax) + + assert torch.all(scales_triton.cpu() == scales_ref).item(), \ + "Triton midmax scales differ from pure-Python reference" + + +# --------------------------------------------------------------------------- +# 3. Kernel output: blockscale_selection="midmax" vs blockscale_selection="default" differ appropriately +# --------------------------------------------------------------------------- + +def test_midmax_differs_from_round_even_on_outliers(): + """ + For blocks whose amax == 7.0 (exactly E2M1 midmax), round-even rounds up + but midmax does not (strict > comparison), so scales must differ. + """ + block_size = 32 + x = torch.full((4, 64), 7.0, dtype=torch.float32, device="cuda") + + _, scales_midmax = convert_to_mxfp4(x, blockscale_selection="midmax") + _, scales_round_even = convert_to_mxfp4(x, blockscale_selection="default") + + assert not torch.equal(scales_midmax, scales_round_even), \ + "Expected scales to differ at amax==7.0 (midmax boundary)" + # midmax should be strictly lower (no bump) than round-even (bumps at 7.0) + assert torch.all(scales_midmax < scales_round_even).item(), \ + "midmax scales should be lower than round-even at amax==7.0" + + +def test_midmax_false_matches_pytorch_ref(): + """ + With blockscale_selection="default" the triton kernel must match the existing pytorch + reference implementation (which implements round-even only). + """ + x = prepare_data((128, 64), torch.float32) + _, scales_triton = convert_to_mxfp4(x, blockscale_selection="default") + _, scales_ref = convert_to_mxfp4_pytorch(x) + assert torch.all(scales_triton == scales_ref).item(), \ + "Triton round-even scales differ from pytorch reference" + + +# --------------------------------------------------------------------------- +# 4. Quantize → dequantize roundtrip with blockscale_selection="midmax" +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("tensor_shape", [(128, 64), (4, 128, 64)]) +@pytest.mark.parametrize("axis", [-1, -2]) +@pytest.mark.parametrize("is_2d_block", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_midmax_roundtrip(tensor_shape, axis, is_2d_block, dtype): + """ + Roundtrip (quant → dequant) with blockscale_selection="midmax" should reconstruct + the input with reasonable accuracy (MAE within one E2M1 quantum of the + block scale). + """ + x = prepare_data(tensor_shape, dtype) + data_lp, scales = convert_to_mxfp4(x, axis=axis, is_2d_block=is_2d_block, blockscale_selection="midmax") + x_dq = convert_from_mxfp4(data_lp, scales, output_dtype=dtype, axis=axis, is_2d_block=is_2d_block) + + mae = (x.float() - x_dq.float()).abs().mean().item() + # E2M1 has 4 representable values per octave; 25% relative error is generous + amax = x.float().abs().max().item() + assert mae < 0.25 * amax + 1e-3, f"Roundtrip MAE {mae:.4f} too large (amax={amax:.4f})" + + +# --------------------------------------------------------------------------- +# 5. Edge cases +# --------------------------------------------------------------------------- + +def test_midmax_all_zeros(): + """A block of all zeros must not produce NaN/Inf scales or outputs.""" + x = torch.zeros(32, 32, dtype=torch.float32, device="cuda") + data_lp, scales = convert_to_mxfp4(x, blockscale_selection="midmax") + assert not torch.any(torch.isnan(scales.float())), "NaN in scales for zero input" + assert torch.all(scales >= 1).item(), "scale below minimum-normal clamp" + + x_dq = convert_from_mxfp4(data_lp, scales, output_dtype=torch.float32) + assert not torch.any(torch.isnan(x_dq)), "NaN in dequantized output for zero input" + assert torch.all(x_dq == 0).item(), "Zero input should dequantize to zeros" + + +def test_midmax_large_values(): + """Very large values (near FP32 max) must not produce inf/nan scales.""" + x = torch.full((32, 32), 1e30, dtype=torch.float32, device="cuda") + data_lp, scales = convert_to_mxfp4(x, blockscale_selection="midmax") + assert torch.all(scales < 255).item(), "scale hit 0xFF (inf/nan exponent)" + assert not torch.any(torch.isnan(scales.float())) + + +def test_midmax_negative_values(): + """Negative-valued blocks should produce the same scales as their positive counterpart.""" + torch.manual_seed(7) + x_pos = torch.abs(torch.randn(64, 64, dtype=torch.float32, device="cuda")) + x_neg = -x_pos + + _, scales_pos = convert_to_mxfp4(x_pos, blockscale_selection="midmax") + _, scales_neg = convert_to_mxfp4(x_neg, blockscale_selection="midmax") + assert torch.all(scales_pos == scales_neg).item(), \ + "Negating all values should not change midmax scales" + + +def test_midmax_single_outlier_block(): + """ + A single block containing one value just above midmax (7.1) surrounded by + small values should bump only that block's scale. + """ + block_size = 32 + x = torch.ones(1, 2 * block_size, dtype=torch.float32, device="cuda") * 0.5 + # First block: push amax above midmax + x[0, :block_size] = 7.1 + + _, scales = convert_to_mxfp4(x, blockscale_selection="midmax") + scale_bumped_block = scales[0, 0].item() + scale_small_block = scales[0, 1].item() + assert scale_bumped_block > scale_small_block, \ + "Block with amax > midmax should have a higher scale than a block with small values" + + +# --------------------------------------------------------------------------- +# 6. Scale minimum clamp (all-zero block) +# --------------------------------------------------------------------------- + +def test_midmax_scale_minimum_clamp(): + """ + The kernel clamps scales to minimum 1 (= 2^1 in FP32 representation). + An all-zero block exercises this path. + """ + x = torch.zeros(32, 64, dtype=torch.float32, device="cuda") + _, scales = convert_to_mxfp4(x, blockscale_selection="midmax") + assert torch.all(scales >= 1).item(), "All scales must be >= 1 (minimum-normal clamp)" + + +# --------------------------------------------------------------------------- +# 7. Consistency: blockscale_selection="midmax" gives same result across dtypes (f32 vs bf16) +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("use_asm", [False]) +def test_midmax_dtype_consistency(use_asm): + """ + Quantizing the same tensor in float32 and bfloat16 with blockscale_selection="midmax" + should yield scales that are close (within ±1) due to bf16 precision loss. + """ + if use_asm and not is_cdna4(): + pytest.skip("ASM mode only on CDNA4") + + torch.manual_seed(0) + x_f32 = torch.randn(64, 64, dtype=torch.float32, device="cuda") + x_bf16 = x_f32.to(torch.bfloat16) + + _, scales_f32 = convert_to_mxfp4(x_f32, blockscale_selection="midmax", use_asm=use_asm) + _, scales_bf16 = convert_to_mxfp4(x_bf16, blockscale_selection="midmax", use_asm=use_asm) + + diff = (scales_f32.int() - scales_bf16.int()).abs() + assert diff.max().item() <= 1, \ + f"Scale difference between f32 and bf16 exceeded ±1: max={diff.max().item()}" + + +# --------------------------------------------------------------------------- +# 8. ASM path (CDNA4 only) +# --------------------------------------------------------------------------- + +def test_midmax_asm_matches_non_asm(): + """On CDNA4, use_asm=True with blockscale_selection="midmax" must match use_asm=False.""" + if not is_cdna4(): + pytest.skip("ASM path requires CDNA4 hardware") + + x = prepare_data((128, 64), torch.float32) + data_lp_asm, scales_asm = convert_to_mxfp4(x, blockscale_selection="midmax", use_asm=True) + data_lp_ref, scales_ref = convert_to_mxfp4(x, blockscale_selection="midmax", use_asm=False) + + assert torch.all(scales_asm == scales_ref).item(), "ASM/non-ASM scales differ under midmax" + assert torch.all(data_lp_asm == data_lp_ref).item(), "ASM/non-ASM fp4 values differ under midmax"