Integrate RAD's LPT recipes

.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

.gitignore

@@ -212,3 +212,19 @@ __marimo__/
 datasets/
 /models/
 comm_traces/
+
+
+# Slurm
+*.out
+
+# Checkpoints
+*.distcp
+gptoss_chkpt/
+plotting/plots
+
+*.safetensors
+
+docker/
+gpt_*/
+*/plotting/
+docs/

.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

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

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 <NODE_A_IP>
+
+```
+
+You can get `<NODE_A_IP>` 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
+
+```

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",
+]

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",
 ]