diff --git a/.github/actionlint.yaml b/.github/actionlint.yaml deleted file mode 100644 index b0d95d88e34537bfddd380dd054393190e7cebca..0000000000000000000000000000000000000000 --- a/.github/actionlint.yaml +++ /dev/null @@ -1,3 +0,0 @@ -self-hosted-runner: - labels: - - docker-builder-01 diff --git a/.github/workflows/build-and-commit.yml b/.github/workflows/build-and-commit.yml deleted file mode 100644 index 2b8d63d05640696ccc4fc70d3e2fcf0b35d74044..0000000000000000000000000000000000000000 --- a/.github/workflows/build-and-commit.yml +++ /dev/null @@ -1,120 +0,0 @@ -name: Nix build and commit - -on: - pull_request: - types: [opened, synchronize, reopened] - workflow_dispatch: - -permissions: - contents: write - -jobs: - check-commit: - runs-on: ubuntu-latest - outputs: - skip: ${{ steps.check.outputs.skip }} - steps: - - uses: actions/checkout@v4 - with: - fetch-depth: 0 - - id: check - run: | - if [ "${{ github.event_name }}" = "pull_request" ]; then - msg=$(git log -1 --pretty=%B "${{ github.event.pull_request.head.sha }}") - else - msg="manual dispatch" - fi - echo "Commit message: $msg" - if echo "$msg" | grep -q '\[skip-build\]'; then - echo "skip=true" >> "$GITHUB_OUTPUT" - else - echo "skip=false" >> "$GITHUB_OUTPUT" - fi - - build_and_commit: - needs: check-commit - if: needs.check-commit.outputs.skip == 'false' - runs-on: docker-builder-01 - steps: - - name: Show disk usage - run: df -h - - - name: Notify build start on Slack - id: slack_start - run: | - msg="*Build started* for \`${{ github.repository }}\`\nBranch: \`${{ github.ref_name }}\`\n<${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}|View Workflow>" - response=$(curl -s -X POST \ - -H "Authorization: Bearer ${{ secrets.SLACK_TOKEN }}" \ - -H "Content-type: application/json; charset=utf-8" \ - --data "{\"channel\":\"${{ secrets.SLACK_CHANNEL_ID }}\",\"text\":\"$msg\"}" \ - https://slack.com/api/chat.postMessage) - ts=$(echo "$response" | jq -r '.ts') - echo "thread_ts=$ts" >> "$GITHUB_OUTPUT" - echo "$response" - - - name: Checkout repository - uses: actions/checkout@v4 - with: - fetch-depth: 0 - lfs: true - ref: ${{ github.head_ref || github.ref }} - - - name: Install Nix - uses: cachix/install-nix-action@v31 - - - name: Setup huggingface cachix - uses: cachix/cachix-action@v15 - with: - name: huggingface - - - name: Clean build directory - run: | - rm -rf build - - - name: Build with Nix - run: | - nix run .#build-and-copy \ - --override-input kernel-builder github:huggingface/kernel-builder \ - --max-jobs 8 \ - -j 8 \ - -L - - - name: List built binaries - run: | - ls build - - - name: Commit build artifact - run: | - git config user.name "github-actions[bot]" - git config user.email "41898282+github-actions[bot]@users.noreply.github.com" - git add build/* - git commit -m "Add built binary [skip-build]" - - - name: Push changes - run: | - git push origin HEAD:"$HEAD_REF" - env: - HEAD_REF: ${{ github.head_ref || github.ref }} - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - - - name: Notify success on Slack (thread) - if: success() - run: | - ts="${{ steps.slack_start.outputs.thread_ts }}" - msg="*Build succeeded* for \`${{ github.repository }}\`\nBranch: \`${{ github.ref_name }}\`\n<${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}|View Workflow>" - curl -s -X POST \ - -H "Authorization: Bearer ${{ secrets.SLACK_TOKEN }}" \ - -H "Content-type: application/json; charset=utf-8" \ - --data "{\"channel\":\"${{ secrets.SLACK_CHANNEL_ID }}\",\"text\":\"$msg\",\"thread_ts\":\"$ts\"}" \ - https://slack.com/api/chat.postMessage - - - name: Notify failure on Slack (thread) - if: failure() - run: | - ts="${{ steps.slack_start.outputs.thread_ts }}" - msg="*Build failed* for \`${{ github.repository }}\`\nBranch: \`${{ github.ref_name }}\`\n<${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}|View Workflow>" - curl -s -X POST \ - -H "Authorization: Bearer ${{ secrets.SLACK_TOKEN }}" \ - -H "Content-type: application/json; charset=utf-8" \ - --data "{\"channel\":\"${{ secrets.SLACK_CHANNEL_ID }}\",\"text\":\"$msg\",\"thread_ts\":\"$ts\"}" \ - https://slack.com/api/chat.postMessage diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml deleted file mode 100644 index 6d2ff056b79fbd65f2ed61018c5730b234fee9b7..0000000000000000000000000000000000000000 --- a/.github/workflows/pre-commit.yml +++ /dev/null @@ -1,30 +0,0 @@ -name: pre-commit - -on: - pull_request: - push: - branches: [ main, master ] - -jobs: - run-pre-commit: - runs-on: ubuntu-latest - permissions: - contents: read - pull-requests: read - steps: - - uses: actions/checkout@v4 - - - uses: actions/setup-python@v5 - with: - python-version: "3.11" - - - name: Cache pre-commit - uses: actions/cache@v4 - with: - path: ~/.cache/pre-commit - key: pre-commit-${{ runner.os }}-${{ hashFiles('.pre-commit-config.yaml') }} - restore-keys: | - pre-commit-${{ runner.os }}- - - - name: Run pre-commit - uses: pre-commit/action@v3.0.1 diff --git a/.github/workflows/push-to-hf.yml b/.github/workflows/push-to-hf.yml deleted file mode 100644 index ae31d9649bab120a948956ea6cf422dfa33b74a8..0000000000000000000000000000000000000000 --- a/.github/workflows/push-to-hf.yml +++ /dev/null @@ -1,40 +0,0 @@ -name: Push to HF Repo - -on: - push: - branches: - - main - workflow_dispatch: - -jobs: - push_to_hf: - runs-on: ubuntu-latest - steps: - # 1. Checkout the repo - - name: Checkout repository - uses: actions/checkout@v4 - with: - fetch-depth: 0 - - name: Install Git LFS - run: | - git lfs install - git lfs fetch --all - git lfs pull - # 2. Set up Git - - name: Configure Git - run: | - git config user.name "MotifTech" - git config user.email "huggingface@motiftech.io" - - # 3. Add HF remote - - name: Add Hugging Face remote - run: | - git remote add hf https://huggingface.co/Motif-Technologies/optimizer - git fetch hf || true - - # 4. Push to HF repo - - name: Push to Hugging Face - env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} - run: | - git push "https://hf_token:${HF_TOKEN}@huggingface.co/Motif-Technologies/optimizer" HEAD:main diff --git a/.gitignore b/.gitignore deleted file mode 100644 index 3c818f24455eb15948604ef3c9a32c5e351cb27c..0000000000000000000000000000000000000000 --- a/.gitignore +++ /dev/null @@ -1,21 +0,0 @@ -__pycache__ -.idea -.DS_Store -*.egg-info -outputs -dist/* -.vscode - -# data -data -out -wandb - -torchtitan/datasets/**/*.model -torchtitan/experiments/flux/assets/* - -# temp files -*.log -error.json -_remote_module_non_scriptable.py -.git_disabled/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml deleted file mode 100644 index 55f8e34c04aac06db5a3137a475e13e3e5ecf8d5..0000000000000000000000000000000000000000 --- a/.pre-commit-config.yaml +++ /dev/null @@ -1,33 +0,0 @@ -default_install_hook_types: - - pre-commit - - commit-msg -default_stages: - - pre-commit # Run locally - - manual # Run in CI -exclude: '(build|result)/.*|__pycache__/.*|.*\.(png|html)$' -repos: -- repo: https://github.com/google/yapf - rev: v0.43.0 - hooks: - - id: yapf - args: [--in-place, --verbose] -- repo: https://github.com/crate-ci/typos - rev: v1.34.0 - hooks: - - id: typos - exclude: '.gitattributes' -- repo: https://github.com/PyCQA/isort - rev: 6.0.1 - hooks: - - id: isort -- repo: https://github.com/pre-commit/mirrors-clang-format - rev: v20.1.3 - hooks: - - id: clang-format - types_or: [c++, cuda] - args: [--style=file, --verbose] -- repo: https://github.com/jackdewinter/pymarkdown - rev: v0.9.29 - hooks: - - id: pymarkdown - args: [fix] diff --git a/CLAUDE.md b/CLAUDE.md deleted file mode 100644 index cf12b9a281680b32e8247ba103bfe86e53d8b7ff..0000000000000000000000000000000000000000 --- a/CLAUDE.md +++ /dev/null @@ -1,108 +0,0 @@ -# CLAUDE.md - -This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. - -## Project Overview - -Optimizer is a PyTorch package implementing the **Muon optimizer** with support for N-D sharding parallelism for large-scale distributed training. Based on the paper at https://arxiv.org/abs/2511.07464. It supports general N-D sharding configurations (FSDP2 through hybrid setups like 2 TP + 2 DP-Replicate + 2 DP-Shard). - -## Commands - -### Lint & Format - -```bash -pre-commit run --all-files # Run all pre-commit hooks -pre-commit run isort --all-files # Run a specific hook (e.g., isort) -``` - -Hooks: yapf (Python formatter), isort (import sorter), typos (spell checker), clang-format (C++/CUDA), pymarkdown (Markdown linter), actionlint (GitHub Actions). - -### Tests - -Tests require **8 GPUs**, access to `Motif-Technologies/Motif-2.6B-4layer-random` on HuggingFace (`HF_TOKEN` env var), and PyTorch >= 2.8.0. - -```bash -cd test && ./run_test.sh -# Equivalent to: -cd test && torchrun --nproc-per-node=8 --local-ranks-filter=0 -m pytest test_muon.py -``` - -Useful pytest flags: `--measure-perf` (timing/memory), `--do-profile` (profiling, requires `--measure-perf`), `--skip-verify` (skip correctness check against sequential implementation). - -### Build - -Uses kernel-builder infrastructure (`build.toml`, `flake.nix`). Pre-built binaries for various PyTorch/CUDA/ROCm combinations are stored in `build/`. - -### Commit Convention - -**Always append `[skip-build]` to every commit message.** This prevents CI from triggering unnecessary build jobs on development branches. - -## Architecture - -### Source Layout - -``` -torch-ext/optimizer/ -├── __init__.py # Public API: exports Muon -├── muon.py # Muon optimizer class (~430 lines) -├── newton_schulz.py # Newton-Schulz iteration (~50 lines) -├── qk_clip.py # QK clipping for attention heads (~130 lines) -├── core.py # Shared state, helpers, param grouping (~110 lines) -├── pipeline.py # Async generator pipeline for parallel mode (~290 lines) -├── async_utils.py # AsyncTask / AsyncRuntime scheduling (~75 lines) -├── adamw.py # Fused AdamW for non-Muon parameters (~160 lines) -├── matmul_transpose_triton.py # Triton kernel for X @ X.T (~130 lines) -└── distributed/ - └── utils.py # Shard mesh construction, DTensor slicing (~175 lines) -``` - -### Optimizer Modes - -The `Muon` optimizer has three execution paths selected per-parameter based on its tensor type and mesh structure: - -1. **Base mode** (`base()`) — Single-device / non-sharded tensors. Standard Muon with Newton-Schulz orthogonalization. -2. **Distributed mode** (`distributed_muon()`) — Gathers full tensors via all-gather, computes updates, redistributes. Used for small parameters or fallback. -3. **Parallel mode** (`parallel()`) — Pipelined all2all communication overlapped with compute. Uses an async generator pipeline scheduled by `run_pipeline()`. This is the main advanced feature. - -### Parallel Mode Pipeline - -The parallel pipeline is implemented as a single generator function `muon_chunk_pipeline()` in `pipeline.py`. Parameters are split into chunks, and each chunk flows through: - -``` -build bufs + async all2all_gather → yield → wait + Newton-Schulz compute + async all2all_scatter → yield → wait + update_param -``` - -The generator yields 2 times (after launching async gather and async scatter via `async_op=True`), allowing `run_pipeline()` to interleave multiple chunks for communication overlap. `work.wait()` completes each async operation after the yield. - -`warmup_step` maps to `max_concurrent_tasks = warmup_step + 1` in `run_pipeline()`. - -For detailed implementation documentation (pipeline internals, distributed utilities, QK clipping with strided sharding, etc.), see [`docs/implementation.md`](docs/implementation.md). - -### Key Abstractions - -- **`get_default_muon_param_groups(model, is_muon_func)`** (`core.py`) — Separates parameters into Muon-optimizable (2D+) and AdamW groups. Skips embeddings and output layers by default. -- **`_muon_state` dataclass** (`core.py`) — Per-parameter config: rank ownership (`worker_rank`), process group, precomputed shard indices (`rank_indices`, `rank_numels`), and optional QK clip state. Config-only; no transient pipeline state. -- **`muon_chunk_pipeline()` generator** (`pipeline.py`) — Processes one chunk through the full gather→compute→scatter→update pipeline. Uses `async_op=True` for non-blocking all-to-all and yields to allow chunk interleaving. All intermediate buffers are generator-local variables. -- **`run_pipeline()`** (`async_utils.py`) — Generator-based pipeline scheduling with bounded concurrency. Interleaves multiple chunk pipelines at yield points. -- **`construct_shard_mesh()` / `get_slices_of_dtensor()`** (`distributed/utils.py`) — Utilities for building shard meshes from DTensor placements and computing per-rank local slices. Handles both `Shard` and `_StridedShard` (PyTorch 2.10+). -- **Newton-Schulz iteration** (`newton_schulz.py`) — `_zeropower_via_newtonschulz5()`: 5 quintic iterations in bfloat16 with pre-optimized coefficients for gradient orthogonalization. Uses Triton kernel `matmul_transpose_assign` for efficient X @ X.T. -- **QK Clipping** (`qk_clip.py`) — Optional dynamic clipping of attention head projections when QK logits exceed a threshold. Configured via `q_indices`, `k_indices`, `head_dim`, `threshold`. -- **Fused AdamW** (`adamw.py`) — Uses PyTorch's `torch._fused_adamw_` for non-Muon parameters, grouping tensors by device/dtype and DTensor placement. - -### Dependency Graph - -``` -matmul_transpose_triton.py (leaf) - │ - newton_schulz.py (leaf + triton) - │ - core.py ──── qk_clip.py (leaf, distributed/utils) - │ │ │ - │ pipeline.py ─── async_utils.py - │ │ - │ adamw.py - │ │ - muon.py (all above) - │ - __init__.py -``` diff --git a/README.md b/README.md index 5416fc8ec3ca3469a807f04d2622d88a39a1c64c..e24a75af8cfd719a94a499a644188b3164b2d1cb 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,6 @@ --- tags: -- kernels -license: apache-2.0 +- kernel --- # Optimizer @@ -10,14 +9,8 @@ Optimizer is a python package that provides: - PyTorch implementation of recent optimizer algorithms - with support for parallelism techniques for efficient large-scale training. -## Currently implemented -- Parallel Muon with N-D sharding - - [arxiv URL](https://arxiv.org/abs/2511.07464) - - Supports **general N-D sharding configurations** - - The implementation is not tied to any specific parallel strategy. - - Verified from basic FSDP2 setups up to hybrid configurations such as - **(2 TP + 2 DP-Replicate + 2 DP-Shard)**. - - Verified configurations can be found in [test_muon.py](./test/test_muon.py) +### Currently implemented +- [Parallel Muon with FSDP2](./docs/muon/parallel_muon.pdf) ## Usage @@ -27,78 +20,14 @@ from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from kernels import get_kernel optimizer = get_kernel("motif-technologies/optimizer") -get_default_muon_param_groups = optimizer.muon.get_default_muon_param_groups model = None # your model here fsdp_model = FSDP(model) -# muon, in nature, cannot use 1-d tensor -# we provide helper function to group such tensors -# you can use your own function, if necessary -params = get_default_muon_param_groups(model) # user can write own is_muon_func, if necessary - optim = optimizer.Muon( - params, + fsdp_model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4, ) ``` - -## Documentation - -- [Implementation Guide](./docs/implementation.md) — Detailed walkthrough of the internal architecture, parallel pipeline, distributed utilities, and QK clipping. Recommended for code reviewers and new contributors. -- [PyTorch 2.10 TP Fix](./docs/pytorch-2.10-tp-fix.md) — Root cause analysis and fixes for `_StridedShard` compatibility with PyTorch 2.10+. - -## Test - -- Check [test/README.md](./test/README.md) for how to run the tests. - -## Pre-commit Hooks - -This project uses [pre-commit](https://pre-commit.com/) to automatically check and format code before commits. - -### Setup - -1. Install pre-commit: - - ```bash - pip install pre-commit - ``` - -2. Install the git hooks: - -```bash - pre-commit install - ``` - -Once installed, the configured hooks will run automatically on each commit. - -### Included Hooks - -The following tools are run via pre-commit: - -- **[yapf](https://github.com/google/yapf)** – Python code formatter -- **[typos](https://github.com/crate-ci/typos)** – Spell checker for common typos -- **[isort](https://github.com/PyCQA/isort)** – Organizes and sorts Python imports -- **[clang-format](https://clang.llvm.org/docs/ClangFormat.html)** – Formats C++/CUDA code (`--style=file`) -- **[pymarkdown](https://github.com/jackdewinter/pymarkdown)** – Lints and auto-fixes Markdown files -- **[actionlint](https://github.com/rhysd/actionlint)** – Validates GitHub Actions workflows - -### Usage - -- Run all checks on the entire codebase: - - ```bash - pre-commit run --all-files - ``` - -- Run a specific hook (example: isort): - - ```bash - pre-commit run isort --all-files - ``` - -### Test - -- There is a [simple unittest for Parallel Muon](./test/test_muon/README.md) diff --git a/_typos.toml b/_typos.toml deleted file mode 100644 index 13dde69f3cd50f77d4607c7912e7a6427ec34de9..0000000000000000000000000000000000000000 --- a/_typos.toml +++ /dev/null @@ -1,3 +0,0 @@ -[default.extend-words] -# Math notation used in docs/muon-clip.md (O subscript t, update step output) -Ot = "Ot" diff --git a/build.toml b/build.toml index ebabc676bfe40eb07e2bb447ff0c17605ac42844..b80854db0a67cdde4e5c3dcb8d95f18704812383 100644 --- a/build.toml +++ b/build.toml @@ -1,33 +1,23 @@ [general] name = "optimizer" -backends = [ - "cuda", - "rocm", -] +universal = false [torch] src = [ - "torch-ext/torch_binding.cpp", - "torch-ext/torch_binding.h", + "torch-ext/torch_binding.cpp", + "torch-ext/torch_binding.h", ] -[kernel.optimizer] -backend = "cuda" -depends = ["torch"] -src = ["optimizer/dummy.cu"] - -[kernel.optimizer_rocm] +[kernel.activation] backend = "rocm" -rocm-archs = [ - "gfx906", - "gfx908", - "gfx90a", - "gfx940", - "gfx941", - "gfx942", - "gfx1030", - "gfx1100", - "gfx1101", +src = [ + "optimizer/dummy.cu", +] +depends = [ "torch" ] + +[kernel.activation_cuda] +backend = "cuda" +src = [ + "optimizer/dummy.cu", ] -depends = ["torch"] -src = ["optimizer/dummy.cu"] +depends = [ "torch" ] diff --git a/build/torch210-cxx11-cu126-x86_64-linux/_ops.py b/build/torch210-cxx11-cu126-x86_64-linux/_ops.py deleted file mode 100644 index 034bff088659b5df6f6d401feb18c89dc5f33b29..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _optimizer_8d53b78_dirty -ops = torch.ops._optimizer_8d53b78_dirty - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_optimizer_8d53b78_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch210-cxx11-cu126-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so b/build/torch210-cxx11-cu126-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so deleted file mode 100755 index 8f9c672be7ef1f016613e205dfc115f9af85d8e2..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:075fc73dbb2750aed7598cc3e13b593b6b1e7a78a78491e1b852fbd2a9af8f8d -size 1940944 diff --git a/build/torch210-cxx11-cu126-x86_64-linux/adamw.py b/build/torch210-cxx11-cu126-x86_64-linux/adamw.py deleted file mode 100644 index b5a95816a9f5b9e1889eaadae65373bfbced809a..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-x86_64-linux/adamw.py +++ /dev/null @@ -1,271 +0,0 @@ -import logging -from collections import defaultdict -from typing import cast - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -def fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: float | torch.Tensor, - weight_decay: float, - eps: float, - maximize: bool, -) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: dict | None = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) - - -def _to_local(t): - """Unwrap DTensor to local tensor for fused ops.""" - return t._local_tensor if isinstance(t, DTensor) else t - - -# --------------------------------------------------------------------------- -# Caches for eliminating per-step Python overhead. -# -# Placement grouping and tensor list assembly are identical every step -# (params don't change placement, moment/step tensors are the same objects -# after initialisation). We cache them keyed by id() of the param list -# stored in param_groups (stable across steps). -# -# Only gradients change each step and must be collected fresh. -# --------------------------------------------------------------------------- - -# id(group["params"]) → dict[placement_key, list[param]] -_placement_cache: dict[int, dict[tuple, list]] = {} - -# id(placement_group_list) → (params_local, moment1, moment2, state_steps) -_tensor_cache: dict[int, tuple[list, list, list, list]] = {} - - -def _step_adamw_params_slow(optimizer_state, params, group): - """Uncached fallback for the rare case where some params lack grads.""" - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - g = p.grad - if g is None: - continue - state = optimizer_state[p] - params_with_grads.append(_to_local(p)) - grads.append(_to_local(g)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - if not params_with_grads: - return - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw_params(optimizer_state, params, group): - """Run fused AdamW on a list of parameters sharing the same placement. - - After the first call, cached tensor lists (params_local, moment1, - moment2, state_steps) are reused — only gradients are collected fresh. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - params: List of parameters to update. - group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay. - """ - # Collect grads — the only thing that changes each step. - with record_function("adamw::collect_grads"): - grads = [] - for p in params: - g = p.grad - if g is None: - # Rare: fall back to slow path that filters per-param. - _step_adamw_params_slow(optimizer_state, params, group) - return - grads.append(_to_local(g)) - - tensor_key = id(params) - if tensor_key not in _tensor_cache: - with record_function("adamw::init_tensor_cache"): - params_local = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - state = optimizer_state[p] - params_local.append(_to_local(p)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(p.grad) - state["moment2"] = torch.zeros_like(p.grad) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - _tensor_cache[tensor_key] = (params_local, moment1, moment2, - state_steps) - - params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key] - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - with record_function("adamw::fused_adamw"): - fused_adamw( - params_local, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw(optimizer_state, group): - """Dispatch AdamW step, grouping parameters by type and placement. - - Placement grouping is cached after the first call since params never - change their placement between steps. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - group: Parameter group dict. - """ - params = group["params"] - placement_key = id(params) - - if placement_key not in _placement_cache: - with record_function("adamw::group_by_placement"): - placement_to_params: dict[tuple, - list[torch.Tensor]] = defaultdict(list) - for p in params: - match p: - case DTensor(): - logger.debug( - "[AdamW] DTensor param: shape=%s, placements=%s, " - "mesh=%s, grad=%s", p.shape, p.placements, - p.device_mesh.mesh_dim_names, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple( - [p.placements, p.device_mesh])].append(p) - case torch.Tensor(): - logger.debug( - "[AdamW] plain param: shape=%s, grad=%s", p.shape, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple([torch.Tensor, - None])].append(p) - - logger.debug("[AdamW] %d placement groups, %d total params", - len(placement_to_params), len(params)) - - _placement_cache[placement_key] = dict(placement_to_params) - - for group_params in _placement_cache[placement_key].values(): - step_adamw_params(optimizer_state, group_params, group) diff --git a/build/torch210-cxx11-cu126-x86_64-linux/async_utils.py b/build/torch210-cxx11-cu126-x86_64-linux/async_utils.py deleted file mode 100644 index a45c530ac9cad88e3555ec1047a6aa59f225347e..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-x86_64-linux/async_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import logging -from typing import Generator - -logger = logging.getLogger(__name__) - - -class _Task: - """Internal: wraps a generator, advances one yield at a time.""" - - def __init__(self, generator: Generator[None, None, None], index: int): - self._generator = generator - self._index = index - self._steps_completed = 0 - self.step() # run to first yield - - def step(self) -> bool: - try: - next(self._generator) - self._steps_completed += 1 - logger.debug("pipeline[%d] completed stage %d", self._index, - self._steps_completed) - return True - except StopIteration: - logger.debug("pipeline[%d] finished after %d stages", self._index, - self._steps_completed) - return False - - def close(self): - self._generator.close() - - -def run_pipeline( - pipelines: Generator[Generator[None, None, None], None, None], - max_concurrent: int, -) -> None: - """Run generator-based pipelines with bounded concurrency. - - Each pipeline is a generator that yields at stage boundaries. - The runtime interleaves pipelines so communication and computation - overlap across chunks. - """ - if max_concurrent <= 0: - raise ValueError(f"max_concurrent must be > 0, got {max_concurrent}") - - have_new = True - task_index = 0 - previous_tasks: list[_Task] = [] - - try: - while have_new or previous_tasks: - running_tasks: list[_Task] = [] - - # Admit one new pipeline per iteration (staggered admission). - # Admitting one at a time ensures that while chunk N does NS - # compute on the default stream, chunk N+1's NCCL all-to-all - # runs concurrently on the NCCL stream — creating real - # communication/computation overlap on the GPU. - if have_new and len(previous_tasks) < max_concurrent: - try: - gen = next(pipelines) - task = _Task(gen, task_index) - task_index += 1 - running_tasks.append(task) - except StopIteration: - have_new = False - - # Advance every previously-yielded task by one step. - for task in previous_tasks: - if task.step(): - running_tasks.append(task) - - previous_tasks = running_tasks - except BaseException: - # Clean up all in-flight generators to release GPU resources. - for task in previous_tasks: - task.close() - raise diff --git a/build/torch210-cxx11-cu126-x86_64-linux/core.py b/build/torch210-cxx11-cu126-x86_64-linux/core.py deleted file mode 100644 index c69d515afef305ad0ed66374095fa2d2468d99cc..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-x86_64-linux/core.py +++ /dev/null @@ -1,219 +0,0 @@ -import logging -import math -from dataclasses import dataclass -from typing import List - -import torch -from torch.distributed import ProcessGroup -from torch.distributed.tensor import DTensor - -# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into -# parameter FQNs. Activation checkpointing similarly inserts -# "_checkpoint_wrapped_module". Strip these so name-based matching (skip_keys, -# expert_keys, QK layer parsing) works regardless of wrapper nesting. -_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) - -logger = logging.getLogger(__name__) - - -def normalize_fqn(name: str) -> str: - """Strip torch.compile / checkpoint wrapper components from a parameter FQN.""" - return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS) - - -@dataclass -class _muon_state: - worker_rank: int - process_group: ProcessGroup - rank_indices: dict[int, tuple] # local_rank -> per-dim indices - rank_numels: dict[int, int] # local_rank -> numel - name: str - qk_clip_state: torch.Tensor | None = None - - -def _batch_momentum( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update (no nesterov).""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - - -def _batch_momentum_nesterov( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update with nesterov correction.""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - nesterov_terms = torch._foreach_mul(momentum_bufs, momentum) - torch._foreach_add_(grads, nesterov_terms) - - -_compiled_momentum: dict[bool, callable] = {} -_use_momentum_compile = True - - -def set_momentum_compile(enabled: bool): - """Toggle torch.compile for batched momentum.""" - global _use_momentum_compile - _use_momentum_compile = enabled - - -def batch_pre_ortho( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, - nesterov: bool, -) -> None: - """Batched momentum update on lists of plain tensors. - - Mirrors dion's ``muon_update_pre_orthogonalize``. - Inputs must be plain CUDA tensors (not DTensor). - Modifies ``momentum_bufs`` and (for nesterov) ``grads`` in-place. - - When compile is enabled, uses separately compiled functions for - nesterov=True/False to avoid graph breaks from the branch. - """ - fn = _batch_momentum_nesterov if nesterov else _batch_momentum - if _use_momentum_compile: - if nesterov not in _compiled_momentum: - _compiled_momentum[nesterov] = torch.compile(fn) - fn = _compiled_momentum[nesterov] - fn(grads, momentum_bufs, momentum) - - -def _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay): - """Weight-decay + update on plain tensors. - - Not compiled: per-param @torch.compile caused ~0.25ms TorchDynamo cache - lookup per call × 256+ params = massive overhead. The pipeline path uses - batched _foreach_* ops instead; this function remains for base() and - distributed_muon(). - """ - p_data.mul_(1 - lr * weight_decay) - p_data.add_(u_data, alpha=-adjusted_lr) - - -def update_p(p, u, lr, adjusted_lr, weight_decay): - """Apply weight decay and orthogonalized update to parameter. - - Args: - p: Parameter (torch.nn.Parameter or DTensor). - u: Orthogonalized update tensor. - lr: Base learning rate. - adjusted_lr: Size-adjusted learning rate. - weight_decay: Weight decay coefficient. - """ - # Unwrap Parameter -> underlying data tensor. - p_data = p.data if isinstance(p, torch.nn.Parameter) else p - # Unwrap DTensor -> local CUDA tensor for compiled kernel. - if isinstance(p_data, DTensor): - p_data = p_data._local_tensor - u_data = u._local_tensor if isinstance(u, DTensor) else u - _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay) - - -def adjust_lr_for_muon(lr, param_shape): - """Scale learning rate based on parameter matrix dimensions. - - Args: - lr: Base learning rate. - param_shape: Shape of the parameter tensor. - - Returns: - Adjusted learning rate. - """ - A, B = param_shape[:2] - # We adjust the learning rate and weight decay based on the size of the parameter matrix - # as described in the paper - adjusted_ratio = 0.2 * math.sqrt(max(A, B)) - adjusted_lr = lr * adjusted_ratio - return adjusted_lr - - -def _match_key(parts, key): - """Check if key matches as contiguous components in parts. - - Single-component keys (e.g. "experts") match any single component. - Multi-component keys (e.g. "experts.w1") match as a contiguous subsequence. - """ - key_parts = key.split(".") - key_len = len(key_parts) - if key_len == 1: - return key in parts - return any(parts[i:i + key_len] == key_parts - for i in range(len(parts) - key_len + 1)) - - -def is_expert_param(name, expert_keys): - """Check if a parameter name matches any expert key (component-level).""" - if not expert_keys: - return False - parts = normalize_fqn(name).split(".") - return any(_match_key(parts, key) for key in expert_keys) - - -def default_is_muon(name, x, expert_keys=None): - normalized = normalize_fqn(name) - parts = normalized.split(".") - skip_keys = [ - "embed_tokens", - "lm_head", - "tok_embeddings", - "output", - "mhc_attn", - "mhc_ffn", - "lambda_proj", - ] - if any(key in parts for key in skip_keys): - logger.info( - "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d", - normalized, name, x.ndim) - return False - effective_ndim = x.ndim - is_expert = is_expert_param(name, expert_keys) - if is_expert: - effective_ndim -= 1 - result = effective_ndim >= 2 - logger.info( - "[is_muon] %s (orig: %s): ndim=%d, expert=%s, effective_ndim=%d → %s", - normalized, name, x.ndim, is_expert, effective_ndim, - "Muon" if result else "AdamW") - return result - - -def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None): - if is_muon_func is None: - is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys) - - muon_params, muon_names = [], [] - non_muon_params, non_muon_names = [], [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - non_muon_names.append(n) - - logger.info("[param_groups] expert_keys=%s, Muon=%d, AdamW=%d", - expert_keys, len(muon_names), len(non_muon_names)) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] diff --git a/build/torch210-cxx11-cu126-x86_64-linux/cpu_offload.py b/build/torch210-cxx11-cu126-x86_64-linux/cpu_offload.py deleted file mode 100644 index fb5e69154a1d4a6c884491413a37a9acf0f66c80..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-x86_64-linux/cpu_offload.py +++ /dev/null @@ -1,206 +0,0 @@ -"""CPU offloading for optimizer states. - -Manages a pinned CPU memory pool and async CUDA streams to offload -optimizer state tensors (momentum buffers, Adam moments) to CPU between -optimizer steps, freeing GPU memory. - -All tracked tensors are packed into a single flat pinned CPU buffer -(per dtype). D2H and H2D copies are performed per-tensor directly -between individual GPU tensors and their slice of the CPU flat buffer -— no GPU staging buffer is allocated, so there is **no temporary GPU -memory spike** during offload or reload. - -Individual tensor storages are freed after offload via -``untyped_storage().resize_(0)``, preserving tensor identity so -downstream caches remain valid. -""" - -import logging -from collections import defaultdict - -import torch -from torch.distributed.tensor import DTensor - -logger = logging.getLogger(__name__) - - -class CPUOffloadPool: - """Pinned CPU memory pool for async optimizer state offloading. - - Tracked tensors are grouped by dtype. Each group gets a single flat - pinned CPU buffer. D2H / H2D copies are per-tensor (into slices of - the flat buffer) to avoid allocating a GPU staging buffer. - """ - - def __init__(self): - self._managed: list[torch.Tensor] = [] - self._storage_nbytes: dict[int, int] = {} # id(t) → bytes - - # Per-dtype group: populated on first offload. - # dtype → dict with keys: - # "indices" : list[int] managed-list indices - # "offsets" : list[tuple[int,int]] (start, numel) in flat buf - # "total" : int total numel - # "cpu_flat" : Tensor pinned CPU buffer - self._groups: dict[torch.dtype, dict] = {} - - self._offload_stream: torch.cuda.Stream | None = None - self._device: torch.device | None = None - self._initialized: bool = False - self._logged: bool = False - - # ------------------------------------------------------------------ - @staticmethod - def _local(t: torch.Tensor) -> torch.Tensor: - """Unwrap DTensor to its local CUDA tensor.""" - return t._local_tensor if isinstance(t, DTensor) else t - - def _ensure_stream(self): - if self._offload_stream is None: - self._offload_stream = torch.cuda.Stream(device=self._device) - - # ------------------------------------------------------------------ - def track(self, tensor: torch.Tensor): - """Register a GPU tensor for CPU offloading. Idempotent.""" - tid = id(tensor) - if tid in self._storage_nbytes: - return - local = self._local(tensor) - if self._device is None: - self._device = local.device - storage = local.untyped_storage() - # Skip tensors with empty storage (e.g. empty FSDP shards) - if storage.size() == 0: - return - self._storage_nbytes[tid] = storage.size() - self._managed.append(tensor) - - # ------------------------------------------------------------------ - def _init_buffers(self): - """Build per-dtype flat buffers on first offload.""" - # Group managed tensors by dtype. - dtype_map: dict[torch.dtype, list[tuple[int, int]]] = defaultdict(list) - for idx, t in enumerate(self._managed): - local = self._local(t) - dtype_map[local.dtype].append((idx, local.numel())) - - total_cpu_bytes = 0 - for dtype, entries in dtype_map.items(): - offsets: list[tuple[int, int]] = [] - indices: list[int] = [] - off = 0 - for idx, n in entries: - indices.append(idx) - offsets.append((off, n)) - off += n - cpu_flat = torch.empty(off, dtype=dtype, device="cpu", pin_memory=True) - self._groups[dtype] = { - "indices": indices, - "offsets": offsets, - "total": off, - "cpu_flat": cpu_flat, - } - total_cpu_bytes += off * cpu_flat.element_size() - - self._initialized = True - logger.info( - "[CPUOffload] Pool initialized: %d tensors, %d dtype group(s), " - "%.2f MB pinned CPU memory", - len(self._managed), - len(self._groups), - total_cpu_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def offload(self): - """Per-tensor async D2H into CPU flat buffer, then free GPU storage.""" - if not self._managed: - return - if not self._initialized: - self._init_buffers() - self._ensure_stream() - - # Offload stream waits for compute to finish. - compute_event = torch.cuda.current_stream(self._device).record_event() - self._offload_stream.wait_event(compute_event) - - offloaded_bytes = 0 - - # Per-tensor D2H copies directly into CPU flat buffer slices. - # No GPU staging buffer → no temporary GPU memory spike. - with torch.cuda.stream(self._offload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - cpu_flat[off : off + n].copy_(local.reshape(-1), non_blocking=True) - - offloaded_bytes += grp["total"] * cpu_flat.element_size() - - # Wait for all D2H copies to land, then free GPU storage. - self._offload_stream.synchronize() - for t in self._managed: - storage = self._local(t).untyped_storage() - if storage.size() != 0: - storage.resize_(0) - else: - raise RuntimeError( - f"Tensor storage is already freed (size=0) before offload. " - f"This indicates a double-free or external interference. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - - if not self._logged: - logger.info( - "[CPUOffload] Offloaded %.2f MB (GPU → CPU)", - offloaded_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def reload(self): - """Per-tensor H2D from CPU flat buffer on the default stream. - - Runs on the current (default) CUDA stream to avoid stream - interaction issues with the parallel Muon pipeline. Since - pinned CPU memory is the source, the copies overlap with - GPU idle time between steps. - """ - if not self._managed or not self._initialized: - return - - reloaded_bytes = 0 - - # Re-allocate all GPU storages first. - for t in self._managed: - local = self._local(t) - storage = local.untyped_storage() - if storage.size() != 0: - raise RuntimeError( - f"Storage should have been freed (size=0) before reload, " - f"but got size={storage.size()}. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - storage.resize_(self._storage_nbytes[id(t)]) - - # Per-tensor H2D copies from CPU flat buffer slices. - # non_blocking=True with pinned source allows DMA overlap. - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off : off + n], non_blocking=True) - - reloaded_bytes += grp["total"] * cpu_flat.element_size() - - if not self._logged: - logger.info( - "[CPUOffload] Reloaded %.2f MB (CPU → GPU)", reloaded_bytes / (1024**2) - ) diff --git a/build/torch210-cxx11-cu126-x86_64-linux/distributed/utils.py b/build/torch210-cxx11-cu126-x86_64-linux/distributed/utils.py deleted file mode 100644 index 683096bcd50489e7a09959753fdbc29ba0f3b6a6..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-x86_64-linux/distributed/utils.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import (Placement, Shard, - _StridedShard) - - -def _is_shard(placement: Placement) -> bool: - """Check if a placement is a shard type (Shard or _StridedShard). - - In PyTorch 2.10+, _StridedShard no longer inherits from Shard, so - ``placement.is_shard()`` returns False for _StridedShard. This helper - handles both old and new hierarchies. - """ - return isinstance(placement, (Shard, _StridedShard)) - - -def get_slices_of_dtensor( - target: DTensor | torch.Tensor, - local_rank: int, - shard_mesh: DeviceMesh, - shard_placements: tuple[Placement], -) -> tuple[slice | torch.Tensor, ...]: - """ - Get per-dimension indices for a given rank's shard of the target tensor. - - Uses ``Shard.local_shard_size_and_offset`` and - ``_StridedShard.local_shard_size_and_offset`` for correct handling of - both contiguous and strided (non-contiguous) sharding. - - Args: - target (DTensor | torch.Tensor): The target tensor (for its shape). - local_rank (int): The local rank within the shard group. - shard_mesh (DeviceMesh): The shard mesh (only shard dimensions). - shard_placements (tuple[Placement]): The shard placements. - - Returns: - A tuple of indices (one per tensor dim). Each element is either: - - A ``slice`` (for contiguous or unsharded dims) - - A 1-D ``torch.LongTensor`` of indices (for strided sharding) - """ - - # find the global rank of the local rank in the shard mesh - rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank] - - rank_coords = (shard_mesh.mesh == rank).nonzero() - - assert len(rank_coords) == 1 - rank_coords = tuple(rank_coords[0].tolist()) - - assert len(rank_coords) == len(shard_placements) - - # Track per-shard-dim indices. - # None means "not yet sharded on this dim". - dim_indices: dict[int, torch.Tensor] = {} - - # Caution: Assuming replicate-to-shard of the shard mesh goes with - # left-to-right sharding. This is ensured by the sorting logic of - # construct_shard_mesh function. - for mesh_dim_idx, (rank_coord, placement) in enumerate( - zip(rank_coords, shard_placements)): - assert _is_shard(placement) - - num_chunks = shard_mesh.mesh.shape[mesh_dim_idx] - shard_dim = placement.dim - - # Current effective size on this dim (may already be sub-sharded) - if shard_dim in dim_indices: - curr_size = len(dim_indices[shard_dim]) - else: - curr_size = target.size()[shard_dim] - - # Compute indices for this level of sharding - if isinstance(placement, _StridedShard): - _shard_size, offsets = _StridedShard.local_shard_size_and_offset( - placement, - curr_size, - num_chunks, - rank_coord, - return_first_offset=False) - new_indices = torch.tensor(offsets, dtype=torch.long) - else: - shard_size, offset = Shard.local_shard_size_and_offset( - curr_size, num_chunks, rank_coord) - new_indices = torch.arange(offset, - offset + shard_size, - dtype=torch.long) - - # Compose with previous indices on this dim - if shard_dim in dim_indices: - dim_indices[shard_dim] = dim_indices[shard_dim][new_indices] - else: - dim_indices[shard_dim] = new_indices - - # Build result tuple - result: list[slice | torch.Tensor] = [] - for d in range(len(target.size())): - if d not in dim_indices: - result.append(slice(None)) - else: - indices = dim_indices[d] - # Convert contiguous indices to slice for efficiency - if len(indices) > 0: - start = indices[0].item() - expected = torch.arange(start, - start + len(indices), - dtype=torch.long) - if torch.equal(indices, expected): - result.append(slice(start, start + len(indices))) - else: - result.append(indices) - else: - result.append(slice(0, 0)) - - return tuple(result) - - -_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, - ProcessGroup]] = dict() - - -def construct_shard_mesh( - placements: tuple[Placement], - mesh: DeviceMesh, -) -> tuple[DeviceMesh, ProcessGroup, tuple[Placement, ...]]: - """Construct shard sub-mesh and ProcessGroup for all-to-all communication. - - Given a DTensor's placements and device mesh, extracts the "shard group" - — the set of ranks that together hold all shards of the same replica — - and creates a ProcessGroup for all-to-all among them. - - Steps: - 1. Sort placements: Replicate first, then Shard by (dim, granularity). - 2. Permute the mesh tensor to match the sorted order. - 3. Collapse Replicate dims → list of shard sub-meshes (one per replica). - 4. Create/retrieve a cached ProcessGroup for the current rank's sub-mesh. - - Example — 8 GPUs, mesh shape (2, 2, 2), - placements ``[Shard(0), Replicate, _StridedShard(0)]``:: - - Step 1 — Sort: [Replicate, _StridedShard(0), Shard(0)] - Permutation: [1, 2, 0] - - Step 2 — Permute mesh dims by [1, 2, 0]: - Original: Permuted: - [[[0,1],[2,3]], [[[0,2],[1,3]], - [[4,5],[6,7]]] [[4,6],[5,7]]] - - Step 3 — Unbind replicate dim (dim 0), giving 2 shard sub-meshes: - sub-mesh 0 = [[0,2],[1,3]] (replica group 0) - sub-mesh 1 = [[4,6],[5,7]] (replica group 1) - shard_placements = (_StridedShard(0), Shard(0)) - - Step 4 — Rank 0 → ProcessGroup([0,1,4,5]) - Rank 2 → ProcessGroup([2,3,6,7]) - - Returns: - ``(shard_mesh, process_group, shard_placements)`` - """ - my_rank = dist.get_rank() - assert mesh.mesh.device.type == 'cpu' - - # -- Fast path: 1D all-shard mesh → reuse existing PG. ---------------- - # Reuses the mesh's existing ProcessGroup directly, avoiding the - # overhead of dist.new_group(). The standard path below also handles - # subset calls safely via use_local_synchronization=True, but this - # fast path is still beneficial for the common 1D shard case. - if mesh.ndim == 1 and len(placements) == 1 and _is_shard(placements[0]): - key = (*mesh.mesh.shape, *mesh.mesh.flatten().tolist()) - if key not in _ranks_to_dist_cache: - _ranks_to_dist_cache[key] = (mesh, mesh.get_group()) - return (*_ranks_to_dist_cache[key], tuple(placements)) - - mesh_tensor = mesh.mesh.clone() - - # -- Step 1: Sort placements (Replicate first, then Shard by dim). ------ - # _StridedShard comes BEFORE regular Shard on the same dim so that - # get_slices_of_dtensor applies the outer sharding first, matching - # DTensor's left-to-right (outer-to-inner) composition order. - def _sort_key(item): - index, placement = item - assert not placement.is_partial(), "Partial placement not supported" - if placement.is_replicate(): - return (-1, 0, index) - assert _is_shard(placement), f"Unsupported: {type(placement)}" - split = (-1 / placement.split_factor if isinstance( - placement, _StridedShard) else 0) - return (placement.dim, split, index) - - indexed = sorted(enumerate(placements), key=_sort_key) - perm, sorted_placements = zip(*indexed) - - # -- Step 2: Permute mesh to match sorted placement order. -------------- - sorted_mesh = mesh_tensor.permute(perm) - - # -- Step 3: Collapse replicate dims → list of shard sub-meshes. -------- - # E.g. mesh (2, 3, 4, 4) with [R, R, S(0), S(1)] → 6 sub-meshes of (4, 4) - num_rep = sum(1 for p in sorted_placements if p.is_replicate()) - if num_rep > 0: - if num_rep > 1: - sorted_mesh = sorted_mesh.flatten(0, num_rep - 1) - shard_meshes = list(torch.unbind(sorted_mesh, dim=0)) - else: - shard_meshes = [sorted_mesh] - shard_placements = sorted_placements[num_rep:] - assert len(shard_placements) == len(set(shard_placements)) - - # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. -- - # Each rank only creates the group it belongs to, using - # use_local_synchronization=True so that only group members need to - # coordinate. This avoids deadlocks when different PP stages call - # construct_shard_mesh for different parameters. - def _cache_key(t: torch.Tensor) -> tuple: - return (*t.shape, *t.flatten().tolist()) - - my_key = None - for sm in shard_meshes: - if (my_rank == sm).any().item(): - key = _cache_key(sm) - assert my_key is None, "Rank appears in multiple shard groups" - my_key = key - if key not in _ranks_to_dist_cache: - pg = dist.new_group(sm.flatten().tolist(), - use_local_synchronization=True) - _ranks_to_dist_cache[key] = ( - DeviceMesh(device_type="cuda", mesh=sm), - pg, - ) - - return (*_ranks_to_dist_cache[my_key], shard_placements) diff --git a/build/torch210-cxx11-cu126-x86_64-linux/matmul_transpose_triton.py b/build/torch210-cxx11-cu126-x86_64-linux/matmul_transpose_triton.py deleted file mode 100644 index 792de23d82c3fb45fe33d397ab9b76a0787259d0..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-x86_64-linux/matmul_transpose_triton.py +++ /dev/null @@ -1,122 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], - restore_value=['y'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -@torch.library.custom_op("muon::matmul_transpose_assign", - mutates_args=("d_out", )) -def matmul_transpose_assign(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """Compute d_out = d_in @ d_in.T using an optimized Triton kernel.""" - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -@matmul_transpose_assign.register_fake -def _(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """FakeTensor impl: d_out is already allocated, mutation is declared.""" - pass diff --git a/build/torch210-cxx11-cu126-x86_64-linux/metadata.json b/build/torch210-cxx11-cu126-x86_64-linux/metadata.json deleted file mode 100644 index c55a35717622f1dd5c8ba376ea3a814cbcc10d78..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-x86_64-linux/metadata.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch210-cxx11-cu126-x86_64-linux/muon.py b/build/torch210-cxx11-cu126-x86_64-linux/muon.py deleted file mode 100644 index 14c0e22471fa6d47a51ed95e0e0c341dc18d5194..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-x86_64-linux/muon.py +++ /dev/null @@ -1,1068 +0,0 @@ -import logging -import types -from collections import defaultdict -from typing import Any - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Replicate, Shard -from torch.profiler import record_function - -from .adamw import _placement_cache, _tensor_cache, step_adamw -from .async_utils import run_pipeline -from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho, - get_default_muon_param_groups, is_expert_param, update_p) -from .cpu_offload import CPUOffloadPool -from .distributed.utils import (_is_shard, construct_shard_mesh, - get_slices_of_dtensor) -from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO, - _zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5_batched) -from .pipeline import muon_chunk_pipeline, prelaunch_first_gather -from .qk_clip import compute_scales, get_qk_clip_info, qk_clip - -logger = logging.getLogger(__name__) - - -def _expand_expert_params(names, params, expert_keys): - """Expand expert params by splitting on dim 0 (expert dimension). - - Params whose name matches any key in ``expert_keys`` are treated as - expert-parallel tensors. Their outermost dimension is the expert - dimension: an ``(E, out, in)`` tensor becomes ``E`` separate 2D - ``nn.Parameter`` views so that in-place updates propagate back to - the original storage. - - Non-expert params with ``ndim > 2`` trigger an ``AssertionError`` — - if they are expert params, their key must be added to ``expert_keys``. - - The grad must already be set on each expert param (e.g. after momentum). - - For DTensor expert params, placements that shard on dim 0 (expert dim) - are consumed by the split. Non-dim-0 shard placements (e.g. TP) are - preserved: each 2D slice is wrapped as a DTensor on the corresponding - submesh so the parallel pipeline handles the TP communication. - """ - expanded_names = [] - expanded_params = [] - - for n, p in zip(names, params): - is_expert = is_expert_param(n, expert_keys) - is_dtensor = isinstance(p.data, DTensor) - - if is_expert: - if is_dtensor: - logger.debug( - "[expand_expert] %s: expert DTensor, shape=%s, " - "placements=%s, mesh=%s, local_shape=%s", n, p.shape, - p.placements, p.device_mesh.mesh_dim_names, - p.to_local().shape) - else: - logger.debug( - "[expand_expert] %s: expert plain tensor, shape=%s", n, - p.data.shape) - - if not is_expert: - assert p.data.ndim <= 2, ( - f"Param {n} has ndim={p.data.ndim} but does not match " - f"expert_keys={expert_keys}. If this is an expert param, " - f"add its key to expert_keys.") - expanded_names.append(n) - expanded_params.append(p) - continue - - g = p.grad - assert g is not None, ( - f"Expert param {n} must have grad set before expansion") - - tp_mesh = None - tp_placements_2d = None - - if is_dtensor: - local_data = p.to_local() - local_grad = g.to_local() if isinstance(g, DTensor) else g - - # Find non-dim-0 shard placements (e.g. TP sharding). - # After splitting on dim 0, Shard(k) becomes Shard(k-1). - tp_dim_indices = [] - tp_placements_2d = [] - for i, pl in enumerate(p.placements): - if _is_shard(pl) and pl.dim != 0: - tp_dim_indices.append(i) - tp_placements_2d.append(Shard(pl.dim - 1)) - - if tp_dim_indices: - tp_dim_names = tuple(p.device_mesh.mesh_dim_names[i] - for i in tp_dim_indices) - if len(tp_dim_names) == 1: - tp_mesh = p.device_mesh[tp_dim_names[0]] - else: - tp_mesh = p.device_mesh[tp_dim_names] - else: - local_data = p.data - local_grad = g - - # Expand: split dim 0, reshape each slice to 2D. - num_local_experts = local_data.shape[0] - for i in range(num_local_experts): - slice_data = local_data[i] - slice_grad = local_grad[i] - - if tp_mesh is not None: - # Wrap as DTensor on TP submesh so the pipeline handles - # TP communication (gather/scatter across TP ranks). - dt_data = DTensor.from_local(slice_data, - device_mesh=tp_mesh, - placements=tp_placements_2d) - dt_grad = DTensor.from_local(slice_grad, - device_mesh=tp_mesh, - placements=tp_placements_2d) - expert_param = torch.nn.Parameter(dt_data, requires_grad=False) - expert_param.grad = dt_grad - else: - expert_param = torch.nn.Parameter(slice_data, - requires_grad=False) - expert_param.grad = slice_grad - - expanded_names.append(f"{n}[{i}]") - expanded_params.append(expert_param) - - p.grad = None # allow expert grad storage to be freed after pipeline - - return expanded_names, expanded_params - - -class Muon(torch.optim.Optimizer): - """ - Muon - MomentUm Orthogonalized by Newton-schulz - - Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- - processing step, in which each 2D parameter's update is replaced with the nearest orthogonal - matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has - the advantage that it can be stably run in bfloat16 on the GPU. - - Some warnings: - - We believe this optimizer is unlikely to work well for training with small batch size. - - We believe it may not work well for finetuning pretrained models, but we haven't tested this. - - Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. - lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) - momentum: The momentum used by the internal SGD. (0.95 is a good default) - nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) - ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. - Parameters that are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW instead. - adamw_lr: The learning rate for the internal AdamW. - adamw_betas: The betas for the internal AdamW. - adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - warmup_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher warmup_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. - chunk_size : Batch size of parameters to process in each - all2all gather/compute/scatter step. - Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified. - use_distributed_muon: Use distributed muon by Liu et al. (2024). - For testing purpose only. - expert_keys: List of strings to identify expert-parallel parameters. - If any key appears in a parameter's name, its outermost - dimension is treated as the expert dimension and expanded - into per-expert 2D params for Muon. For example, - ``expert_keys=["experts"]`` matches any param whose name - contains "experts". 3D+ params not matched by any key - will raise an error. - """ - - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config=None, - warmup_step=5, - chunk_size=-1, - use_distributed_muon=False, - expert_keys=None): - defaults = dict( - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - nesterov=nesterov, - ns_steps=ns_steps, - adamw_betas=adamw_betas, - adamw_eps=adamw_eps, - none_grad=none_grad, - use_muon=True, - ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) - super().__init__(params, defaults) - - self.debug = debug - self.clip_config = clip_config if clip_config is not None else { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100, - } - self.warmup_step = warmup_step - self.chunk_size = chunk_size - self.use_distributed_muon = use_distributed_muon - self.expert_keys = expert_keys - self.cpu_offload = False - self._cpu_offload_pool: CPUOffloadPool | None = None - self._offload_initialized = False - self._parallel_cache: dict[tuple[str, ...], dict] = {} - self._expert_expand_cache: dict[tuple[int, ...], dict] = {} - - def _calc_flops(self, G, steps): - assert len(G.shape) == 2 - M, N = G.shape - if M > N: - M, N = N, M - - return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) - - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - p.placements, p.device_mesh) - - return shard_mesh, shard_pg, shard_placements - - def init_state_and_assign_params(self, names, params, group, qk_logits): - param_to_state = {} - param_to_flops = {} - - total_flops = 0 - for p in params: - g = p.grad - if g is None: - continue - assert g.ndim == 2, "Muon only supports 2D parameters." - - flops = self._calc_flops(g, group["ns_steps"]) - param_to_flops[id(p)] = flops - total_flops += flops - - if self.debug: - logger.debug("Total TFLOPs for Muon: %.2f TFLOPs", - total_flops / 1e12) - - paired = list(zip(names, params)) - - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) - - round_robin = 0 - mesh = ordered_params[0].device_mesh - placements = ordered_params[0].placements - - shard_mesh, shard_pg, shard_placements = self.get_shard_mesh( - ordered_params[0]) - shard_mesh_flattened = shard_mesh.mesh.flatten() - num_ranks = dist.get_world_size(group=shard_pg) - - for n, p in zip(ordered_names, ordered_params): - if mesh != p.device_mesh: - raise ValueError("All parameters must be on the same mesh.") - if placements != p.placements: - raise ValueError("All parameters must have same placements.") - - worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks - round_robin = (round_robin + 1) % len(shard_mesh_flattened) - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - # Precompute per-rank indices and numels for all-to-all. - rank_indices: dict[int, tuple] = {} - rank_numels: dict[int, int] = {} - for r in range(num_ranks): - indices = get_slices_of_dtensor(p, r, shard_mesh, - shard_placements) - rank_indices[r] = indices - numel = 1 - for idx, dim_size in zip(indices, p.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) - rank_numels[r] = numel - - param_to_state[id(p)] = _muon_state( - worker_rank=worker_rank, - process_group=shard_pg, - rank_indices=rank_indices, - rank_numels=rank_numels, - name=n, - qk_clip_state=qk_clip_state, - ) - - return param_to_state, ordered_params - - def base(self, names, params, group, lr, weight_decay, qk_logits): - # Momentum is already applied by _step_muon before this method. - for n, p in zip(names, params): - g = p.grad - if g is None: - continue - - u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) - - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - def distributed_muon( - self, - names: list[str], - params: list[torch.nn.Parameter], - group: dict[str, Any], - lr: float, - weight_decay: float, - qk_logits: list[torch.Tensor | DTensor] | None, - ): - """Batched Distributed Muon — for testing/correctness verification only. - - Uses all-gather to reconstruct full tensors, computes Newton-Schulz on - the full grad, then slices back to local shards. This is simpler but - slower than the parallel pipeline (all2all) path, so it serves as a - reference implementation for verifying correctness. - """ - with record_function("distributed_muon"): - # Momentum is already applied by _step_muon before this method. - ns_steps = group["ns_steps"] - - # Separate plain tensors (no communication) from DTensors. - plain_names, plain_params = [], [] - dtensor_names, dtensor_params = [], [] - for n, p in zip(names, params): - if p.grad is None: - continue - if isinstance(p.data, DTensor): - dtensor_names.append(n) - dtensor_params.append(p) - else: - plain_names.append(n) - plain_params.append(p) - - # Process plain tensors per-param (no communication). - for n, p in zip(plain_names, plain_params): - u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE), - steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, - qk_logits) - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - if not dtensor_params: - return - - # Group DTensors by (placements, mesh) for batched all-gather. - placement_groups: dict[tuple, - tuple[list, - list]] = defaultdict(lambda: ([], [])) - for n, p in zip(dtensor_names, dtensor_params): - key = (p.placements, p.device_mesh) - placement_groups[key][0].append(n) - placement_groups[key][1].append(p) - - logger.info( - "distributed_muon: %d placement groups, %d total dtensors", - len(placement_groups), len(dtensor_params)) - - for (placements, mesh), (grp_names, - grp_params) in placement_groups.items(): - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - placements, mesh) - rank = dist.get_rank(shard_pg) - world_size = dist.get_world_size(shard_pg) - - logger.info(" group: %d params, placements=%s, world_size=%d", - len(grp_params), placements, world_size) - - # Separate params that can be batched (all shard dims evenly - # divisible) from those needing per-param full_tensor - # (e.g. MoE gate weights with fewer rows than shard ranks). - # all_gather_into_tensor requires equal buffer sizes across - # ranks, so uneven splits must use DTensor full_tensor(). - batch_names, batch_params = [], [] - single_names, single_params = [], [] - for n, p in zip(grp_names, grp_params): - even = all(p.shape[pl.dim] % - shard_mesh.mesh.shape[dim_idx] == 0 - for dim_idx, pl in enumerate(shard_placements)) - if even: - batch_names.append(n) - batch_params.append(p) - else: - single_names.append(n) - single_params.append(p) - - # Process uneven-split params per-param via full_tensor(). - for n, p in zip(single_names, single_params): - with record_function("distributed_muon::newton_schulz"): - g_full = p.grad.full_tensor().to(COMM_DTYPE) - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - if not batch_params: - continue - - logger.info(" batched=%d, single=%d", len(batch_params), - len(single_params)) - - # Concat all local grad shards into a single flat buffer. - with record_function("distributed_muon::gather"): - grad_locals = [ - p.grad.to_local().to(COMM_DTYPE).flatten() - for p in batch_params - ] - numels = [g.numel() for g in grad_locals] - grad_concat = torch.cat(grad_locals) - del grad_locals - - # Single all-gather (replaces N separate full_tensor). - grad_gathered = torch.empty( - grad_concat.numel() * world_size, - dtype=COMM_DTYPE, - device="cuda", - ) - dist.all_gather_into_tensor(grad_gathered, - grad_concat, - group=shard_pg) - - total_numel = grad_concat.numel() - del grad_concat - - # Precompute per-param offsets within the concat buffer. - offsets = [] - off = 0 - for ne in numels: - offsets.append(off) - off += ne - - # Per-param: reconstruct full grad → NS → local update. - for i, (n, p) in enumerate(zip(batch_names, batch_params)): - with record_function("distributed_muon::newton_schulz"): - g_full = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - for r in range(world_size): - r_start = r * total_numel + offsets[i] - shard = grad_gathered[r_start:r_start + numels[i]] - indices = get_slices_of_dtensor( - p, r, shard_mesh, shard_placements) - g_full[indices] = shard.reshape( - g_full[indices].shape) - - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - def _setup_parallel(self, names, params, group, qk_logits): - """Compute (or retrieve cached) parallel pipeline metadata. - - Returns: - (ordered_params, param_to_state, rank, chunk_size) - """ - cache_key = tuple(names) - - if cache_key not in self._parallel_cache: - # First call: compute metadata and populate cache. - param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - shard_pg = param_to_state[id(ordered_params[0])].process_group - rank = dist.get_rank(group=shard_pg) - - if self.chunk_size == -1: - shard_ranks = dist.get_world_size(shard_pg) - chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO - elif self.chunk_size > 0: - chunk_size = self.chunk_size - else: - raise ValueError( - "chunk_size must be -1 or a positive integer.") - - ordered_names = [ - param_to_state[id(p)].name for p in ordered_params - ] - name_to_state = { - param_to_state[id(p)].name: param_to_state[id(p)] - for p in ordered_params - } - self._parallel_cache[cache_key] = { - 'ordered_names': ordered_names, - 'name_to_state': name_to_state, - 'rank': rank, - 'chunk_size': chunk_size, - } - else: - # Cached path: rebuild param_to_state with current id(p) keys. - cache = self._parallel_cache[cache_key] - rank = cache['rank'] - chunk_size = cache['chunk_size'] - - name_to_param = dict(zip(names, params)) - ordered_params = [name_to_param[n] for n in cache['ordered_names']] - - param_to_state = {} - for p, n in zip(ordered_params, cache['ordered_names']): - cached_state = cache['name_to_state'][n] - param_to_state[id(p)] = _muon_state( - worker_rank=cached_state.worker_rank, - process_group=cached_state.process_group, - rank_indices=cached_state.rank_indices, - rank_numels=cached_state.rank_numels, - name=n, - qk_clip_state=get_qk_clip_info(self.clip_config, n, - qk_logits), - ) - - return ordered_params, param_to_state, rank, chunk_size - - def parallel(self, - names, - params, - group, - lr, - weight_decay, - qk_logits, - prelaunch_gather=None): - """ - Perform a parallel optimization step using Muon. - - Parameters are chunked and each chunk is processed by a - :func:`muon_chunk_pipeline` generator. :func:`run_pipeline` - interleaves multiple chunks so that communication and computation - overlap across chunks (the same overlap previously achieved by the - warmup + main-loop index scheduling). - - If ``prelaunch_gather`` is provided, it is passed to the first - chunk's generator to skip re-launching the already in-flight - A2A gather. - """ - - # Momentum is already applied by _step_muon before this method. - - ordered_params, param_to_state, rank, chunk_size = ( - self._setup_parallel(names, params, group, qk_logits)) - - def pipelines(): - first = True - for start in range(0, len(ordered_params), chunk_size): - chunk = ordered_params[start:start + chunk_size] - if chunk: - kwargs = dict( - params=chunk, - param_to_state=param_to_state, - rank=rank, - ns_steps=group["ns_steps"], - lr=lr, - weight_decay=weight_decay, - none_grad=group["none_grad"], - ) - if first and prelaunch_gather is not None: - kwargs['prelaunch_gather'] = prelaunch_gather - first = False - yield muon_chunk_pipeline(**kwargs) - - with record_function("muon::pipeline"): - run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1) - - def _step_muon(self, group, qk_logits=None): - params = group["params"] - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - # Apply momentum to all params before routing/expansion. - # Batched using _foreach_* ops (compiled, fullgraph=True). - with record_function("muon::momentum"): - active_params = [p for p in params if p.grad is not None] - if active_params: - # Ensure momentum buffers exist (avoid zeros_like when already present). - for p in active_params: - if "momentum_buffer" not in self.state[p]: - self.state[p]["momentum_buffer"] = torch.zeros_like( - p.grad) - - # Extract local tensors for compiled batch function. - local_grads = [ - p.grad._local_tensor - if isinstance(p.grad, DTensor) else p.grad - for p in active_params - ] - local_bufs = [ - self.state[p]["momentum_buffer"]._local_tensor - if isinstance(self.state[p]["momentum_buffer"], DTensor) - else self.state[p]["momentum_buffer"] - for p in active_params - ] - - # Wrap momentum as tensor for torch.compile. - batch_pre_ortho(local_grads, local_bufs, - torch.tensor(momentum), group["nesterov"]) - - # For non-nesterov, the result is the momentum buffer. - if not group["nesterov"]: - for p in active_params: - p.grad = self.state[p]["momentum_buffer"] - - # Identify batched experts for deferred NS. - # Detection is cheap (condition checks only); actual NS compute is - # deferred so it can overlap with the first chunk's A2A gather. - deferred_expert_work = [] - if self.expert_keys: - batched_expert_indices = [] - for i, (n, p) in enumerate(zip(names, params)): - if not (is_expert_param(n, self.expert_keys) - and p.grad is not None): - continue - # Eligible: plain tensor, or DTensor with no non-dim-0 shards. - if isinstance(p.data, DTensor): - has_tp = any( - _is_shard(pl) and pl.dim != 0 for pl in p.placements) - if has_tp: - continue - batched_expert_indices.append(i) - - if batched_expert_indices: - # Save refs for deferred NS; free grads from param list. - for i in batched_expert_indices: - p = params[i] - g = p.grad - local_g = (g._local_tensor - if isinstance(g, DTensor) else g) - local_data = (p.data._local_tensor if isinstance( - p.data, DTensor) else p.data) - deferred_expert_work.append((local_data, local_g)) - p.grad = None - - # Remove batched experts from lists before expansion. - keep = sorted( - set(range(len(params))) - set(batched_expert_indices)) - names = [names[i] for i in keep] - params = [params[i] for i in keep] - - def _run_deferred_expert_ns(): - """Execute deferred batched expert NS.""" - if not deferred_expert_work: - return - with record_function("muon::batched_expert_ns"): - ns_steps = group["ns_steps"] - for local_data, local_g in deferred_expert_work: - u = zeropower_via_newtonschulz5_batched( - local_g.to(COMM_DTYPE), steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:]) - local_data.mul_(1 - lr * weight_decay) - local_data.add_(u, alpha=-adjusted_lr) - - # Expand expert params by splitting on dim 0. - logger.debug("[_step_muon] before expand: %d params, expert_keys=%s", - len(params), self.expert_keys) - if self.expert_keys: - cache_key = tuple(id(p) for p in params) - cache = self._expert_expand_cache.get(cache_key) - - if cache is None: - # Cold path: full expansion + build cache metadata. - exp_names, exp_params = _expand_expert_params( - names, params, self.expert_keys) - - # Build per-expert-group info for hot-path grad updates. - grad_info = [] - exp_idx = 0 - for orig_idx, (n, p) in enumerate(zip(names, params)): - if not is_expert_param(n, self.expert_keys): - exp_idx += 1 - continue - - is_dt = isinstance(p.data, DTensor) - num_experts = (p.to_local() if is_dt else p.data).shape[0] - - # Detect TP mesh from the first expanded expert param. - tp_mesh = None - tp_pls = None - sample = exp_params[exp_idx] - if isinstance(sample.data, DTensor): - tp_mesh = sample.data.device_mesh - tp_pls = list(sample.data.placements) - - grad_info.append((orig_idx, num_experts, exp_idx, is_dt, - tp_mesh, tp_pls)) - exp_idx += num_experts - - self._expert_expand_cache[cache_key] = { - 'names': exp_names, - 'params': exp_params, - 'grad_info': grad_info, - } - names, params = exp_names, exp_params - else: - # Hot path: reuse cached params, only update expert grads. - for (orig_idx, num_experts, exp_start, is_dt, tp_mesh, - tp_pls) in cache['grad_info']: - p = params[orig_idx] - g = p.grad - local_grad = (g.to_local() - if is_dt and isinstance(g, DTensor) else g) - for i in range(num_experts): - expert_p = cache['params'][exp_start + i] - sg = local_grad[i] - if tp_mesh is not None: - expert_p.grad = DTensor.from_local( - sg, device_mesh=tp_mesh, placements=tp_pls) - else: - expert_p.grad = sg - p.grad = None - - names = cache['names'] - params = cache['params'] - else: - names, params = _expand_expert_params(names, params, - self.expert_keys) - logger.debug("[_step_muon] after expand: %d params", len(params)) - - param_dtensors = [] - name_dtensors = [] - - param_tensors = [] - name_tensors = [] - - # distributed_muon is a reference implementation for testing only. - # The parallel pipeline (all2all) path below is the production path. - if self.use_distributed_muon: - _run_deferred_expert_ns() - self.distributed_muon(names=names, - params=params, - group=group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits) - return - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - logger.debug( - "[route] %s → base (DTensor all-Replicate), " - "shape=%s, placements=%s", n, p.shape, p.placements) - param_tensors.append(p) - name_tensors.append(n) - else: - logger.debug( - "[route] %s → parallel (DTensor), shape=%s, " - "placements=%s, mesh=%s", n, p.shape, p.placements, - p.device_mesh.mesh_dim_names) - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - logger.debug("[route] %s → base (plain tensor), shape=%s", n, - p.data.shape) - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError(f"Unsupported parameter type: {type(p.data)}") - - logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, " - f"{len(param_tensors)} Tensors → base") - - def group_dtensors(dtensors, names): - # To support different placements, we group parameters by placements - # and run parallel Muon on each group. - - placement_to_params = defaultdict(lambda: ([], [])) - - assert len(dtensors) == len(names) - for p, n in zip(dtensors, names): - placement_to_params[tuple([p.placements, - p.device_mesh])][0].append(n) - placement_to_params[tuple([p.placements, - p.device_mesh])][1].append(p) - return placement_to_params - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - dtensor_group = group_dtensors(param_dtensors, name_dtensors) - - # Pre-launch the first chunk's A2A gather so that the NCCL - # communication overlaps with the (deferred) batched expert NS - # compute on the default CUDA stream. - prelaunch = None - if deferred_expert_work: - first_names, first_params = next(iter(dtensor_group.values())) - ordered, pts, rnk, csz = self._setup_parallel( - first_names, first_params, group, qk_logits) - first_chunk = ordered[:csz] - if first_chunk: - prelaunch = prelaunch_first_gather(first_chunk, pts, rnk, - group["none_grad"]) - - _run_deferred_expert_ns() - - first_group = True - for _, (names, params) in dtensor_group.items(): - pg = prelaunch if first_group else None - first_group = False - self.parallel( - names, - params, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - prelaunch_gather=pg, - ) - else: - _run_deferred_expert_ns() - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - ) - - def _register_states_for_offload(self): - """Register all optimizer state tensors with the CPU offload pool. - - Called once after the first step when states have been lazily created. - Offloads all param states (momentum buffers for Muon, moment1/moment2 - for AdamW) to free GPU memory between steps. - """ - pool = self._cpu_offload_pool - tracked = 0 - for group in self.param_groups: - for p in group["params"]: - if p not in self.state: - continue - state = self.state[p] - if group.get("use_muon", False): - if "momentum_buffer" in state: - pool.track(state["momentum_buffer"]) - tracked += 1 - else: - if "moment1" in state: - pool.track(state["moment1"]) - if "moment2" in state: - pool.track(state["moment2"]) - tracked += 1 - logger.info("[CPUOffload] Registered %d param states for offload", - tracked) - - @torch.no_grad - def step(self, closure=None, qk_logits=None): - """Perform a single optimization step. - - Args: - closure (Callable, optional): A closure that reevaluates the model - and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). - """ - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # H2D: reload optimizer states from CPU before computation. - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload() - - logger.debug("[Muon.step] expert_keys=%s, %d param groups", - self.expert_keys, len(self.param_groups)) - - for i, group in enumerate(self.param_groups): - if group["use_muon"]: - logger.debug("[Muon.step] group %d: use_muon=True, %d params", - i, len(group["params"])) - self._step_muon(group, qk_logits=qk_logits) - else: - logger.debug( - "[Muon.step] group %d: use_muon=False (AdamW), %d params", - i, len(group["params"])) - step_adamw(self.state, group) - - # D2H: offload optimizer states to CPU after computation. - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - return loss - - # ------------------------------------------------------------------ - # CPU offload public helpers - # ------------------------------------------------------------------ - - def turn_on_cpu_offload(self): - """Enable CPU offload for optimizer states.""" - if self.cpu_offload: - return - logger.info("[Muon] turn_on_cpu_offload") - self.cpu_offload = True - if not self.state: - return - self._cpu_offload_pool = CPUOffloadPool() - self._offload_initialized = False - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_off_cpu_offload(self): - """Disable CPU offload and keep optimizer states resident on GPU.""" - if not self.cpu_offload: - return - logger.info("[Muon] turn_off_cpu_offload") - if self._offload_initialized: - self._cpu_offload_pool.reload() - torch.cuda.current_stream().synchronize() - self._cpu_offload_pool = None - self._offload_initialized = False - self.cpu_offload = False - - # ------------------------------------------------------------------ - # Checkpoint support for cpu_offload - # ------------------------------------------------------------------ - - def state_dict(self) -> dict: - if self.cpu_offload: - raise RuntimeError( - "Muon.state_dict() requires turn_off_cpu_offload() before checkpoint save." - ) - return super().state_dict() - - def load_state_dict(self, state_dict: dict) -> None: - if self.cpu_offload: - raise RuntimeError( - "Muon.load_state_dict() requires turn_off_cpu_offload() before checkpoint load." - ) - super().load_state_dict(state_dict) - - # Invalidate adamw.py's module-level tensor caches so that - # the next step rebuilds them with the newly loaded state tensors. - _placement_cache.clear() - _tensor_cache.clear() diff --git a/build/torch210-cxx11-cu126-x86_64-linux/newton_schulz.py b/build/torch210-cxx11-cu126-x86_64-linux/newton_schulz.py deleted file mode 100644 index d939264b69a34e7a3fa78859f34dc265a1159d59..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-x86_64-linux/newton_schulz.py +++ /dev/null @@ -1,240 +0,0 @@ -from itertools import repeat -from math import inf, sqrt - -import numpy as np -import torch - -from .matmul_transpose_triton import matmul_transpose_assign - -COMM_DTYPE = torch.bfloat16 -DEFAULT_CHUNK_SIZE_RATIO = 4 - - -def _optimal_quintic(l, u, max_iter=1000): - """ - Use the simplified Remez algorithm to find the optimal odd quintic approximant - to the constant function x -> 1 over the interval [l, u]. - - Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum - approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the - two interior equioscillation nodes q, r until convergence. Returns the - closed-form equioscillating solution when l ≈ u. - - Raises ValueError if any intermediate value (a, b, c, E, q, r) is non-finite - (NaN or inf). Raises RuntimeError if convergence is not reached within - max_iter iterations. - """ - assert 0 <= l <= u - if 1 - 5e-6 <= l / u: - return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5) - q = (3 * l + u) / 4 - r = (l + 3 * u) / 4 - E = inf - for _ in range(max_iter): - old_E = E - LHS = np.array( - [ - [l, l**3, l**5, 1], - [q, q**3, q**5, -1], - [r, r**3, r**5, 1], - [u, u**3, u**5, -1], - ] - ) - a, b, c, E = np.linalg.solve(LHS, np.ones(4)) - if not np.all(np.isfinite([a, b, c, E])): - raise ValueError( - f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}" - ) - q, r = np.sqrt( - (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / (10 * c) - ) - if not np.all(np.isfinite([q, r])): - raise ValueError(f"_optimal_quintic: non-finite node update q={q}, r={r}") - if abs(old_E - E) <= 1e-15: - break - else: - raise RuntimeError( - f"_optimal_quintic: did not converge after {max_iter} iterations" - ) - return float(a), float(b), float(c) - - -def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0): - """ - Compute the Polar Express coefficient series for `num_iters` quintic iterations. - - Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that - compose to map singular values from [l, 1] toward 1. At each step: - 1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion` - prevents near-zero singular values from stalling by raising the effective - lower bound; if it is active (cushion*u > l), the coefficients are - rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u]. - 2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the - last iteration, providing numerical headroom at the cost of a slightly slower - final convergence step. - 3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1). - - Returns a list of (a, b, c) tuples, one per iteration. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - u = 1 - assert 0 <= l <= u - safety_factor = 1 + safety_factor_eps - coefficients = [] - for iter in range(num_iters): - a, b, c = _optimal_quintic(max(l, cushion * u), u) - if cushion * u > l: - pl = a * l + b * l**3 + c * l**5 - pu = a * u + b * u**3 + c * u**5 - rescaler = 2 / (pl + pu) - a *= rescaler - b *= rescaler - c *= rescaler - if iter < num_iters - 1: - a /= safety_factor - b /= safety_factor**3 - c /= safety_factor**5 - coefficients.append((a, b, c)) - l = a * l + b * l**3 + c * l**5 - u = 2 - l - return coefficients - - -# Precomputed Polar Express coefficients (a, b, c) for 10 quintic Newton-Schulz -# iterations. Each tuple is the minimax-optimal (Remez/equioscillation) odd quintic -# approximant to x->1 over the current singular-value interval, computed once at -# import time and reused across all optimizer steps. -# -# Contrast with the former hardcoded NS coefficients (5 fixed tuples): -# - Former: empirically tuned to maximize slope at zero; did not converge -# singular values to 1, yielding US'V^T with S' ~ Uniform(0.5, 1.5) instead -# of the true polar factor UV^T. -# - Polar Express: analytically optimal per step, adapting to the shrinking -# singular-value interval [l, u] as iterations progress; converges all -# singular values to 1, producing the exact polar factor UV^T. -_coeffs_list = _optimal_composition( - l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02 -) - - -# This code is adapted from: -# KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py) -# NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress) -# matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon) -@torch.no_grad() -def _zeropower_via_newtonschulz5(G, steps): - """ - Compute the polar factor of G via the Polar Express method. - - Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c) - are the Polar Express coefficients from `_coeffs_list`. Each step is the - optimal odd quintic approximant to x -> 1 over the current singular-value - interval, minimizing the maximum approximation error (Remez / minimax criterion). - The composition maps singular values from [l, 1] to near 1, producing the - polar factor (orthogonal factor in the polar decomposition G = UP). - - `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2, - cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE - X = G # no manual typecast - - if G.size(0) > G.size(1): - X = X.T - - X = X / (X.norm() + 1e-7) - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - # Perform the NS iterations - for a, b, c in hs: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(0) > G.size(1): - X = X.T - - return X - - -@torch.no_grad() -def _zeropower_via_newtonschulz5_batched(G, steps): - """Batched polar factor computation for 3D (E, out, in) tensors. - - Same algorithm as ``_zeropower_via_newtonschulz5`` but uses - ``torch.bmm`` / ``torch.baddbmm`` instead of the 2D Triton kernel, - processing all E expert matrices in a single batched call. - """ - assert len(G.shape) == 3 - assert G.dtype == COMM_DTYPE - X = G - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - # Per-expert Frobenius norm. - X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) - - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - for a, b, c in hs: - buf1 = torch.bmm(X, X.transpose(-2, -1)) - buf2 = torch.bmm(buf1, buf1.transpose(-2, -1)) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.baddbmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - return X - - -_ns_per_shape: dict[tuple[int, ...], callable] = {} -_use_compile = True - - -def set_ns_compile(enabled: bool): - """Toggle torch.compile for Newton-Schulz iteration.""" - global _use_compile - _use_compile = enabled - - -def zeropower_via_newtonschulz5(G, steps=5): - if not _use_compile: - return _zeropower_via_newtonschulz5(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() - - -def zeropower_via_newtonschulz5_batched(G, steps=5): - """Compile-cached batched Newton-Schulz for 3D expert tensors.""" - if not _use_compile: - return _zeropower_via_newtonschulz5_batched(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile( - _zeropower_via_newtonschulz5_batched, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() diff --git a/build/torch210-cxx11-cu126-x86_64-linux/optimizer/__init__.py b/build/torch210-cxx11-cu126-x86_64-linux/optimizer/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-x86_64-linux/optimizer/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch210-cxx11-cu126-x86_64-linux/pipeline.py b/build/torch210-cxx11-cu126-x86_64-linux/pipeline.py deleted file mode 100644 index c0c2d515856182d8d15ad27dd4e4e093b29397d6..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-x86_64-linux/pipeline.py +++ /dev/null @@ -1,468 +0,0 @@ -import logging -from typing import Generator - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -from .core import _muon_state, adjust_lr_for_muon -from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5 -from .qk_clip import compute_scales - -logger = logging.getLogger(__name__) - -# ====================================================================== -# Stage helpers -# ====================================================================== - - -def _launch_gather( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Allocate gather buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_gather``). - gathered_grads: ``{id(p): empty_tensor}`` for owned params, - ``None`` for non-owned. - recv_counts: Per-source-rank element counts. - """ - # Allocate gathered-grad buffers - gathered_grads: dict[int, torch.Tensor | None] = {} - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - gathered_grads[id(p)] = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - else: - gathered_grads[id(p)] = None - - # Build send buffer – batch grad copies via torch.cat - # (1-2 fused kernels vs N individual narrow().copy_() calls). - send_counts = [0] * num_ranks - for p in params: - state = param_to_state[id(p)] - send_counts[state.worker_rank] += state.rank_numels[rank] - - total_send = sum(send_counts) - if total_send > 0: - # Group grad slices by destination rank in a single pass. - dst_to_grads = [[] for _ in range(num_ranks)] - for p in params: - state = param_to_state[id(p)] - n = state.rank_numels[rank] - if n > 0: - g = p.grad.to_local() - dst_to_grads[state.worker_rank].append(g.reshape(-1)) - - # Flatten in dst order and cat once. - all_slices = [] - for dst in range(num_ranks): - all_slices.extend(dst_to_grads[dst]) - send_buf = torch.cat(all_slices) - if send_buf.dtype != COMM_DTYPE: - send_buf = send_buf.to(COMM_DTYPE) - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += state.rank_numels[src] - recv_counts[src] = total - - recv_buf = torch.empty(sum(recv_counts), dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - logger.debug(f"send_buf size: {send_buf.numel()}, " - f"recv_buf size: {recv_buf.numel()}, " - f"recv_counts: {recv_counts}, " - f"send_counts: {send_counts}, " - f"process_group: {str(process_group)}") - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, gathered_grads, recv_counts - - -def _complete_gather( - recv_buf: torch.Tensor, - recv_counts: list[int], - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - param_to_state: dict[int, _muon_state], - rank: int, -) -> None: - """Reconstruct gathered grads from the recv buffer (in-place).""" - off = 0 - for src in range(len(recv_counts)): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - - indices = state.rank_indices[src] - - shard_view = gathered_grads[id(p)][indices] - n = shard_view.numel() - if n == 0: - continue - - sg = recv_buf.narrow(0, off + inner_off, n) - sg = sg.reshape(shard_view.shape) - gathered_grads[id(p)][indices] = sg - - inner_off += n - assert inner_off == block - off += block - - -def _compute_ns( - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - ns_steps: int, -) -> dict[int, torch.Tensor | None]: - """Run Newton-Schulz orthogonalization on owned parameters. - - Returns: - computed_us: ``{id(p): orthogonalized_update}`` for owned params. - """ - computed_us: dict[int, torch.Tensor | None] = {} - for p in owned_params: - u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps) - gathered_grads[id(p)] = None # free gathered grad - computed_us[id(p)] = u - return computed_us - - -def _launch_scatter( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, - computed_us: dict[int, torch.Tensor | None], -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor], list[int]]: - """Allocate scatter buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_scatter``). - scattered_us: Empty dict, populated by ``_complete_scatter`` with - zero-copy views into ``recv_buf``. - recv_counts: Per-source-rank element counts. - """ - # scattered_us is populated by _complete_scatter with zero-copy views - # into recv_buf, avoiding N empty_like allocations + N copy_ calls. - # Pre-seed entries for params whose local shard is empty (rank_numels == 0) - # so _update_params can iterate all params without KeyError. - scattered_us: dict[int, torch.Tensor] = {} - for p in params: - if param_to_state[id(p)].rank_numels[rank] == 0: - scattered_us[id(p)] = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - # Build send buffer – batch via torch.cat - # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls). - send_counts = [0] * num_ranks - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - for dst_rank in range(num_ranks): - send_counts[dst_rank] += state.rank_numels[dst_rank] - - total_send = sum(send_counts) - if total_send > 0: - # Cache u_full conversions to avoid redundant .to() per dst_rank. - u_fulls = {} - for p in owned_params: - u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous() - - # Collect slices in dst order (matches all-to-all send layout). - all_slices = [] - for dst_rank in range(num_ranks): - for p in owned_params: - state = param_to_state[id(p)] - su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten() - if su.numel() > 0: - all_slices.append(su) - - send_buf = torch.cat(all_slices) if all_slices else torch.empty( - 0, dtype=COMM_DTYPE, device="cuda") - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += state.rank_numels[rank] - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, scattered_us, recv_counts - - -def _complete_scatter( - recv_buf: torch.Tensor, - recv_counts: list[int], - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], -) -> None: - """Populate scattered_us with zero-copy views into recv_buf. - - Instead of pre-allocating tensors and copying, we assign views directly - from ``recv_buf``. This eliminates N ``empty_like`` + N ``copy_`` calls. - The underlying storage of ``recv_buf`` is kept alive through the views - until ``scattered_us`` is cleared after ``_update_params``. - """ - off = 0 - for src in range(len(recv_counts)): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = state.rank_numels[rank] - if n == 0: - continue - - scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - - inner_off += n - - assert inner_off == block - off += block - - -def _update_params( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], - lr: float, - weight_decay: float, -) -> None: - """Apply weight decay, Muon update, and optional QK clipping. - - Uses batched ``_foreach_mul_`` for weight decay and batched - ``_foreach_add_`` for the Muon update, grouping parameters by - adjusted_lr to minimize kernel launches while preserving float32 - precision for the alpha scaling. - """ - if not params: - return - - # Batched weight decay: p *= (1 - lr * wd) — single fused kernel. - p_locals = [p._local_tensor for p in params] - torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay) - - # Group params by adjusted_lr so _foreach_add_ can use a single - # alpha per group (preserves float32 precision for alpha scaling). - lr_groups: dict[float, tuple[list, list]] = {} - for p in params: - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - if adjusted_lr not in lr_groups: - lr_groups[adjusted_lr] = ([], []) - lr_groups[adjusted_lr][0].append(p._local_tensor) - lr_groups[adjusted_lr][1].append(scattered_us[id(p)]) - - for adjusted_lr, (p_group, u_group) in lr_groups.items(): - torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr) - - # QK clipping – applied directly on the local tensor to - # avoid DTensor sharding-propagation issues with _StridedShard. - for p in params: - state = param_to_state[id(p)] - if state.qk_clip_state is None: - continue - scales_full = compute_scales(p, state.qk_clip_state) - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = state.rank_indices[rank][0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - -# ====================================================================== -# Pre-launch helper for overlapping first chunk's gather with other work. -# ====================================================================== - - -@torch.no_grad() -def prelaunch_first_gather( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - none_grad: bool, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Launch the first chunk's A2A gather early for overlap with other compute. - - Call this *before* expensive GPU work (e.g. batched expert NS) so that - the NCCL all-to-all runs concurrently on the NCCL stream while the - default stream executes compute. - - Returns the same 4-tuple that ``_launch_gather`` produces, which should - be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - with record_function("muon::prelaunch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - return work, recv_buf, gathered_grads, recv_counts - - -# ====================================================================== -# Main generator – thin orchestrator that wires stages together. -# ====================================================================== - - -@torch.no_grad() -def muon_chunk_pipeline( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - ns_steps: int, - lr: float, - weight_decay: float, - none_grad: bool, - prelaunch_gather: tuple | None = None, -) -> Generator[None, None, None]: - """Process one chunk of parameters through the full Muon pipeline. - - Stages: gather -> compute (Newton-Schulz) -> scatter -> update. - - Each ``yield`` lets :func:`run_pipeline` interleave other chunks so - that communication and computation overlap across chunks. Async - communication is launched via ``async_op=True`` and completed after - the yield with ``work.wait()``. - - Overlap happens because :func:`run_pipeline` admits one new chunk - per iteration (staggered admission). While chunk *N* does NS - compute on the default CUDA stream, chunk *N+1*'s async all-to-all - runs concurrently on the NCCL stream — no separate ``comm_stream`` - is required. - - If ``prelaunch_gather`` is provided, the gather was already launched - by :func:`prelaunch_first_gather` and we skip launching it again. - - Yields exactly **2** times: - - 1. After launching async all-to-all gather (or immediately if pre-launched). - 2. After launching async all-to-all scatter. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - if prelaunch_gather is not None: - # Gather was pre-launched; none_grad already handled by caller. - work, recv_buf, gathered_grads, recv_counts = prelaunch_gather - else: - # Normal path: launch async gather. - with record_function("muon::launch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - yield # --- YIELD 1: other chunks can launch their gather --- - - with record_function("muon::wait_gather"): - work.wait() - _complete_gather(recv_buf, recv_counts, owned_params, gathered_grads, - param_to_state, rank) - del recv_buf - - # Stage 3: Newton-Schulz orthogonalization. - with record_function("muon::newton_schulz"): - computed_us = _compute_ns(owned_params, gathered_grads, ns_steps) - gathered_grads.clear() - - # Stages 4-5: launch async scatter. - with record_function("muon::launch_scatter"): - work, recv_buf, scattered_us, recv_counts = _launch_scatter( - params, owned_params, param_to_state, rank, num_ranks, - process_group, computed_us) - computed_us.clear() - - yield # --- YIELD 2: other chunks can launch their scatter --- - - with record_function("muon::wait_scatter"): - work.wait() - _complete_scatter(recv_buf, recv_counts, params, param_to_state, rank, - scattered_us) - del recv_buf - - # Stage 6: apply parameter updates. - with record_function("muon::update_params"): - _update_params(params, param_to_state, rank, scattered_us, lr, - weight_decay) - scattered_us.clear() diff --git a/build/torch210-cxx11-cu126-x86_64-linux/qk_clip.py b/build/torch210-cxx11-cu126-x86_64-linux/qk_clip.py deleted file mode 100644 index 2aba711b3004b7f09e7141da7ef834bd61cc2430..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-x86_64-linux/qk_clip.py +++ /dev/null @@ -1,198 +0,0 @@ -import logging -import math -from dataclasses import dataclass - -import torch -from torch.distributed.tensor import DTensor - -from .core import normalize_fqn - -logger = logging.getLogger(__name__) - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - and return (kind, layer_index). - - Supported kinds: - MHA/GQA: 'wq', 'wk', 'q_proj', 'k_proj' - MLA: 'wq_b' (Q up-proj), 'wkv_b' (KV up-proj) - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.1.attn.wq_b.weight' -> ('wq_b', 1) - 'model.0.attn.wkv_b.weight' -> ('wkv_b', 0) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = normalize_fqn(name).split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: str | None # 'wq'/'q_proj'/'wq_b' or 'wk'/'k_proj'/'wkv_b' or None - indices: list[int] # which heads to consider for clipping - head_dim: int # from config (qk_head_dim for MLA wq_b) - threshold: float # from config - logit: torch.Tensor | None - - # MLA-specific fields - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 - - -def get_qk_clip_info(clip_config, n, qk_logits): - """Extract QK clipping info for a named parameter. - - Args: - clip_config: QK clipping configuration dict (or None). - MHA/GQA keys: head_dim, threshold, q_indices, k_indices - MLA extra keys: is_mla=True, qk_nope_head_dim, qk_rope_head_dim, v_head_dim - n: Parameter name string. - qk_logits: Dict mapping layer indices to logit tensors (or None). - - Returns: - QKClipInfo instance with clipping configuration for this parameter. - """ - if clip_config is None: - return None - - head_dim = clip_config.get('head_dim') - threshold = clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - # In TP settings, qk_logits may be DTensor - # We convert it to full tensor here for simplicity - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - -def compute_scales(p, qk_clip_state): - """Compute per-head scaling factors for QK clipping. - - Returns scales tensor (√γ per head) if any head exceeds threshold, else None. - For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim. - """ - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - # Check if any head exceeds threshold before allocating. - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - - if not head_scales: - return None - - # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows - if qk_clip_state.is_mla and kind == 'wkv_b': - effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim - else: - effective_head_dim = head_dim - - H_global = p.shape[0] // effective_head_dim - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale - return scales_full - - -def qk_clip(p, scales, info): - """Apply per-head scaling to a Q/K projection weight matrix. - - Args: - p: Parameter (nn.Parameter or raw tensor). - scales: [n_heads] tensor, each element = √γ_h. - info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions. - - MLA sub-region scaling per Algorithm 1 (MuonClip): - wq_b: q_nope rows → √γ, q_pe rows → γ - wkv_b: k_nope rows → √γ, v rows → unchanged - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not info.is_mla: - # MHA/GQA: uniform √γ applied to all rows in each head - W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: vectorized sub-region scaling within each head - if info.kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_head_dim = qk_nope + info.qk_rope_head_dim - W_3d = W.view(-1, qk_head_dim, W.shape[1]) # [H, qk_head_dim, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # q_nope → √γ - W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1, - 1)) # q_pe → γ - - elif info.kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - W_3d = W.view(-1, kv_stride, W.shape[1]) # [H, kv_stride, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # k_nope → √γ - # v rows: not touched (k_R shared rotary unchanged) diff --git a/build/torch210-cxx11-cu128-x86_64-linux/_ops.py b/build/torch210-cxx11-cu128-x86_64-linux/_ops.py deleted file mode 100644 index 034bff088659b5df6f6d401feb18c89dc5f33b29..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _optimizer_8d53b78_dirty -ops = torch.ops._optimizer_8d53b78_dirty - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_optimizer_8d53b78_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch210-cxx11-cu128-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so b/build/torch210-cxx11-cu128-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so deleted file mode 100755 index 2926579ac6973002c2d4067df7122743ecd1567d..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2af397ae01c8c01ee0e879f6812bd9df55d152afbcc6713f5c1987d5bce7793b -size 2004144 diff --git a/build/torch210-cxx11-cu128-x86_64-linux/adamw.py b/build/torch210-cxx11-cu128-x86_64-linux/adamw.py deleted file mode 100644 index b5a95816a9f5b9e1889eaadae65373bfbced809a..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-x86_64-linux/adamw.py +++ /dev/null @@ -1,271 +0,0 @@ -import logging -from collections import defaultdict -from typing import cast - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -def fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: float | torch.Tensor, - weight_decay: float, - eps: float, - maximize: bool, -) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: dict | None = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) - - -def _to_local(t): - """Unwrap DTensor to local tensor for fused ops.""" - return t._local_tensor if isinstance(t, DTensor) else t - - -# --------------------------------------------------------------------------- -# Caches for eliminating per-step Python overhead. -# -# Placement grouping and tensor list assembly are identical every step -# (params don't change placement, moment/step tensors are the same objects -# after initialisation). We cache them keyed by id() of the param list -# stored in param_groups (stable across steps). -# -# Only gradients change each step and must be collected fresh. -# --------------------------------------------------------------------------- - -# id(group["params"]) → dict[placement_key, list[param]] -_placement_cache: dict[int, dict[tuple, list]] = {} - -# id(placement_group_list) → (params_local, moment1, moment2, state_steps) -_tensor_cache: dict[int, tuple[list, list, list, list]] = {} - - -def _step_adamw_params_slow(optimizer_state, params, group): - """Uncached fallback for the rare case where some params lack grads.""" - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - g = p.grad - if g is None: - continue - state = optimizer_state[p] - params_with_grads.append(_to_local(p)) - grads.append(_to_local(g)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - if not params_with_grads: - return - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw_params(optimizer_state, params, group): - """Run fused AdamW on a list of parameters sharing the same placement. - - After the first call, cached tensor lists (params_local, moment1, - moment2, state_steps) are reused — only gradients are collected fresh. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - params: List of parameters to update. - group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay. - """ - # Collect grads — the only thing that changes each step. - with record_function("adamw::collect_grads"): - grads = [] - for p in params: - g = p.grad - if g is None: - # Rare: fall back to slow path that filters per-param. - _step_adamw_params_slow(optimizer_state, params, group) - return - grads.append(_to_local(g)) - - tensor_key = id(params) - if tensor_key not in _tensor_cache: - with record_function("adamw::init_tensor_cache"): - params_local = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - state = optimizer_state[p] - params_local.append(_to_local(p)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(p.grad) - state["moment2"] = torch.zeros_like(p.grad) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - _tensor_cache[tensor_key] = (params_local, moment1, moment2, - state_steps) - - params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key] - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - with record_function("adamw::fused_adamw"): - fused_adamw( - params_local, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw(optimizer_state, group): - """Dispatch AdamW step, grouping parameters by type and placement. - - Placement grouping is cached after the first call since params never - change their placement between steps. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - group: Parameter group dict. - """ - params = group["params"] - placement_key = id(params) - - if placement_key not in _placement_cache: - with record_function("adamw::group_by_placement"): - placement_to_params: dict[tuple, - list[torch.Tensor]] = defaultdict(list) - for p in params: - match p: - case DTensor(): - logger.debug( - "[AdamW] DTensor param: shape=%s, placements=%s, " - "mesh=%s, grad=%s", p.shape, p.placements, - p.device_mesh.mesh_dim_names, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple( - [p.placements, p.device_mesh])].append(p) - case torch.Tensor(): - logger.debug( - "[AdamW] plain param: shape=%s, grad=%s", p.shape, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple([torch.Tensor, - None])].append(p) - - logger.debug("[AdamW] %d placement groups, %d total params", - len(placement_to_params), len(params)) - - _placement_cache[placement_key] = dict(placement_to_params) - - for group_params in _placement_cache[placement_key].values(): - step_adamw_params(optimizer_state, group_params, group) diff --git a/build/torch210-cxx11-cu128-x86_64-linux/async_utils.py b/build/torch210-cxx11-cu128-x86_64-linux/async_utils.py deleted file mode 100644 index a45c530ac9cad88e3555ec1047a6aa59f225347e..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-x86_64-linux/async_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import logging -from typing import Generator - -logger = logging.getLogger(__name__) - - -class _Task: - """Internal: wraps a generator, advances one yield at a time.""" - - def __init__(self, generator: Generator[None, None, None], index: int): - self._generator = generator - self._index = index - self._steps_completed = 0 - self.step() # run to first yield - - def step(self) -> bool: - try: - next(self._generator) - self._steps_completed += 1 - logger.debug("pipeline[%d] completed stage %d", self._index, - self._steps_completed) - return True - except StopIteration: - logger.debug("pipeline[%d] finished after %d stages", self._index, - self._steps_completed) - return False - - def close(self): - self._generator.close() - - -def run_pipeline( - pipelines: Generator[Generator[None, None, None], None, None], - max_concurrent: int, -) -> None: - """Run generator-based pipelines with bounded concurrency. - - Each pipeline is a generator that yields at stage boundaries. - The runtime interleaves pipelines so communication and computation - overlap across chunks. - """ - if max_concurrent <= 0: - raise ValueError(f"max_concurrent must be > 0, got {max_concurrent}") - - have_new = True - task_index = 0 - previous_tasks: list[_Task] = [] - - try: - while have_new or previous_tasks: - running_tasks: list[_Task] = [] - - # Admit one new pipeline per iteration (staggered admission). - # Admitting one at a time ensures that while chunk N does NS - # compute on the default stream, chunk N+1's NCCL all-to-all - # runs concurrently on the NCCL stream — creating real - # communication/computation overlap on the GPU. - if have_new and len(previous_tasks) < max_concurrent: - try: - gen = next(pipelines) - task = _Task(gen, task_index) - task_index += 1 - running_tasks.append(task) - except StopIteration: - have_new = False - - # Advance every previously-yielded task by one step. - for task in previous_tasks: - if task.step(): - running_tasks.append(task) - - previous_tasks = running_tasks - except BaseException: - # Clean up all in-flight generators to release GPU resources. - for task in previous_tasks: - task.close() - raise diff --git a/build/torch210-cxx11-cu128-x86_64-linux/core.py b/build/torch210-cxx11-cu128-x86_64-linux/core.py deleted file mode 100644 index c69d515afef305ad0ed66374095fa2d2468d99cc..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-x86_64-linux/core.py +++ /dev/null @@ -1,219 +0,0 @@ -import logging -import math -from dataclasses import dataclass -from typing import List - -import torch -from torch.distributed import ProcessGroup -from torch.distributed.tensor import DTensor - -# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into -# parameter FQNs. Activation checkpointing similarly inserts -# "_checkpoint_wrapped_module". Strip these so name-based matching (skip_keys, -# expert_keys, QK layer parsing) works regardless of wrapper nesting. -_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) - -logger = logging.getLogger(__name__) - - -def normalize_fqn(name: str) -> str: - """Strip torch.compile / checkpoint wrapper components from a parameter FQN.""" - return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS) - - -@dataclass -class _muon_state: - worker_rank: int - process_group: ProcessGroup - rank_indices: dict[int, tuple] # local_rank -> per-dim indices - rank_numels: dict[int, int] # local_rank -> numel - name: str - qk_clip_state: torch.Tensor | None = None - - -def _batch_momentum( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update (no nesterov).""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - - -def _batch_momentum_nesterov( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update with nesterov correction.""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - nesterov_terms = torch._foreach_mul(momentum_bufs, momentum) - torch._foreach_add_(grads, nesterov_terms) - - -_compiled_momentum: dict[bool, callable] = {} -_use_momentum_compile = True - - -def set_momentum_compile(enabled: bool): - """Toggle torch.compile for batched momentum.""" - global _use_momentum_compile - _use_momentum_compile = enabled - - -def batch_pre_ortho( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, - nesterov: bool, -) -> None: - """Batched momentum update on lists of plain tensors. - - Mirrors dion's ``muon_update_pre_orthogonalize``. - Inputs must be plain CUDA tensors (not DTensor). - Modifies ``momentum_bufs`` and (for nesterov) ``grads`` in-place. - - When compile is enabled, uses separately compiled functions for - nesterov=True/False to avoid graph breaks from the branch. - """ - fn = _batch_momentum_nesterov if nesterov else _batch_momentum - if _use_momentum_compile: - if nesterov not in _compiled_momentum: - _compiled_momentum[nesterov] = torch.compile(fn) - fn = _compiled_momentum[nesterov] - fn(grads, momentum_bufs, momentum) - - -def _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay): - """Weight-decay + update on plain tensors. - - Not compiled: per-param @torch.compile caused ~0.25ms TorchDynamo cache - lookup per call × 256+ params = massive overhead. The pipeline path uses - batched _foreach_* ops instead; this function remains for base() and - distributed_muon(). - """ - p_data.mul_(1 - lr * weight_decay) - p_data.add_(u_data, alpha=-adjusted_lr) - - -def update_p(p, u, lr, adjusted_lr, weight_decay): - """Apply weight decay and orthogonalized update to parameter. - - Args: - p: Parameter (torch.nn.Parameter or DTensor). - u: Orthogonalized update tensor. - lr: Base learning rate. - adjusted_lr: Size-adjusted learning rate. - weight_decay: Weight decay coefficient. - """ - # Unwrap Parameter -> underlying data tensor. - p_data = p.data if isinstance(p, torch.nn.Parameter) else p - # Unwrap DTensor -> local CUDA tensor for compiled kernel. - if isinstance(p_data, DTensor): - p_data = p_data._local_tensor - u_data = u._local_tensor if isinstance(u, DTensor) else u - _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay) - - -def adjust_lr_for_muon(lr, param_shape): - """Scale learning rate based on parameter matrix dimensions. - - Args: - lr: Base learning rate. - param_shape: Shape of the parameter tensor. - - Returns: - Adjusted learning rate. - """ - A, B = param_shape[:2] - # We adjust the learning rate and weight decay based on the size of the parameter matrix - # as described in the paper - adjusted_ratio = 0.2 * math.sqrt(max(A, B)) - adjusted_lr = lr * adjusted_ratio - return adjusted_lr - - -def _match_key(parts, key): - """Check if key matches as contiguous components in parts. - - Single-component keys (e.g. "experts") match any single component. - Multi-component keys (e.g. "experts.w1") match as a contiguous subsequence. - """ - key_parts = key.split(".") - key_len = len(key_parts) - if key_len == 1: - return key in parts - return any(parts[i:i + key_len] == key_parts - for i in range(len(parts) - key_len + 1)) - - -def is_expert_param(name, expert_keys): - """Check if a parameter name matches any expert key (component-level).""" - if not expert_keys: - return False - parts = normalize_fqn(name).split(".") - return any(_match_key(parts, key) for key in expert_keys) - - -def default_is_muon(name, x, expert_keys=None): - normalized = normalize_fqn(name) - parts = normalized.split(".") - skip_keys = [ - "embed_tokens", - "lm_head", - "tok_embeddings", - "output", - "mhc_attn", - "mhc_ffn", - "lambda_proj", - ] - if any(key in parts for key in skip_keys): - logger.info( - "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d", - normalized, name, x.ndim) - return False - effective_ndim = x.ndim - is_expert = is_expert_param(name, expert_keys) - if is_expert: - effective_ndim -= 1 - result = effective_ndim >= 2 - logger.info( - "[is_muon] %s (orig: %s): ndim=%d, expert=%s, effective_ndim=%d → %s", - normalized, name, x.ndim, is_expert, effective_ndim, - "Muon" if result else "AdamW") - return result - - -def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None): - if is_muon_func is None: - is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys) - - muon_params, muon_names = [], [] - non_muon_params, non_muon_names = [], [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - non_muon_names.append(n) - - logger.info("[param_groups] expert_keys=%s, Muon=%d, AdamW=%d", - expert_keys, len(muon_names), len(non_muon_names)) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] diff --git a/build/torch210-cxx11-cu128-x86_64-linux/cpu_offload.py b/build/torch210-cxx11-cu128-x86_64-linux/cpu_offload.py deleted file mode 100644 index fb5e69154a1d4a6c884491413a37a9acf0f66c80..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-x86_64-linux/cpu_offload.py +++ /dev/null @@ -1,206 +0,0 @@ -"""CPU offloading for optimizer states. - -Manages a pinned CPU memory pool and async CUDA streams to offload -optimizer state tensors (momentum buffers, Adam moments) to CPU between -optimizer steps, freeing GPU memory. - -All tracked tensors are packed into a single flat pinned CPU buffer -(per dtype). D2H and H2D copies are performed per-tensor directly -between individual GPU tensors and their slice of the CPU flat buffer -— no GPU staging buffer is allocated, so there is **no temporary GPU -memory spike** during offload or reload. - -Individual tensor storages are freed after offload via -``untyped_storage().resize_(0)``, preserving tensor identity so -downstream caches remain valid. -""" - -import logging -from collections import defaultdict - -import torch -from torch.distributed.tensor import DTensor - -logger = logging.getLogger(__name__) - - -class CPUOffloadPool: - """Pinned CPU memory pool for async optimizer state offloading. - - Tracked tensors are grouped by dtype. Each group gets a single flat - pinned CPU buffer. D2H / H2D copies are per-tensor (into slices of - the flat buffer) to avoid allocating a GPU staging buffer. - """ - - def __init__(self): - self._managed: list[torch.Tensor] = [] - self._storage_nbytes: dict[int, int] = {} # id(t) → bytes - - # Per-dtype group: populated on first offload. - # dtype → dict with keys: - # "indices" : list[int] managed-list indices - # "offsets" : list[tuple[int,int]] (start, numel) in flat buf - # "total" : int total numel - # "cpu_flat" : Tensor pinned CPU buffer - self._groups: dict[torch.dtype, dict] = {} - - self._offload_stream: torch.cuda.Stream | None = None - self._device: torch.device | None = None - self._initialized: bool = False - self._logged: bool = False - - # ------------------------------------------------------------------ - @staticmethod - def _local(t: torch.Tensor) -> torch.Tensor: - """Unwrap DTensor to its local CUDA tensor.""" - return t._local_tensor if isinstance(t, DTensor) else t - - def _ensure_stream(self): - if self._offload_stream is None: - self._offload_stream = torch.cuda.Stream(device=self._device) - - # ------------------------------------------------------------------ - def track(self, tensor: torch.Tensor): - """Register a GPU tensor for CPU offloading. Idempotent.""" - tid = id(tensor) - if tid in self._storage_nbytes: - return - local = self._local(tensor) - if self._device is None: - self._device = local.device - storage = local.untyped_storage() - # Skip tensors with empty storage (e.g. empty FSDP shards) - if storage.size() == 0: - return - self._storage_nbytes[tid] = storage.size() - self._managed.append(tensor) - - # ------------------------------------------------------------------ - def _init_buffers(self): - """Build per-dtype flat buffers on first offload.""" - # Group managed tensors by dtype. - dtype_map: dict[torch.dtype, list[tuple[int, int]]] = defaultdict(list) - for idx, t in enumerate(self._managed): - local = self._local(t) - dtype_map[local.dtype].append((idx, local.numel())) - - total_cpu_bytes = 0 - for dtype, entries in dtype_map.items(): - offsets: list[tuple[int, int]] = [] - indices: list[int] = [] - off = 0 - for idx, n in entries: - indices.append(idx) - offsets.append((off, n)) - off += n - cpu_flat = torch.empty(off, dtype=dtype, device="cpu", pin_memory=True) - self._groups[dtype] = { - "indices": indices, - "offsets": offsets, - "total": off, - "cpu_flat": cpu_flat, - } - total_cpu_bytes += off * cpu_flat.element_size() - - self._initialized = True - logger.info( - "[CPUOffload] Pool initialized: %d tensors, %d dtype group(s), " - "%.2f MB pinned CPU memory", - len(self._managed), - len(self._groups), - total_cpu_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def offload(self): - """Per-tensor async D2H into CPU flat buffer, then free GPU storage.""" - if not self._managed: - return - if not self._initialized: - self._init_buffers() - self._ensure_stream() - - # Offload stream waits for compute to finish. - compute_event = torch.cuda.current_stream(self._device).record_event() - self._offload_stream.wait_event(compute_event) - - offloaded_bytes = 0 - - # Per-tensor D2H copies directly into CPU flat buffer slices. - # No GPU staging buffer → no temporary GPU memory spike. - with torch.cuda.stream(self._offload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - cpu_flat[off : off + n].copy_(local.reshape(-1), non_blocking=True) - - offloaded_bytes += grp["total"] * cpu_flat.element_size() - - # Wait for all D2H copies to land, then free GPU storage. - self._offload_stream.synchronize() - for t in self._managed: - storage = self._local(t).untyped_storage() - if storage.size() != 0: - storage.resize_(0) - else: - raise RuntimeError( - f"Tensor storage is already freed (size=0) before offload. " - f"This indicates a double-free or external interference. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - - if not self._logged: - logger.info( - "[CPUOffload] Offloaded %.2f MB (GPU → CPU)", - offloaded_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def reload(self): - """Per-tensor H2D from CPU flat buffer on the default stream. - - Runs on the current (default) CUDA stream to avoid stream - interaction issues with the parallel Muon pipeline. Since - pinned CPU memory is the source, the copies overlap with - GPU idle time between steps. - """ - if not self._managed or not self._initialized: - return - - reloaded_bytes = 0 - - # Re-allocate all GPU storages first. - for t in self._managed: - local = self._local(t) - storage = local.untyped_storage() - if storage.size() != 0: - raise RuntimeError( - f"Storage should have been freed (size=0) before reload, " - f"but got size={storage.size()}. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - storage.resize_(self._storage_nbytes[id(t)]) - - # Per-tensor H2D copies from CPU flat buffer slices. - # non_blocking=True with pinned source allows DMA overlap. - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off : off + n], non_blocking=True) - - reloaded_bytes += grp["total"] * cpu_flat.element_size() - - if not self._logged: - logger.info( - "[CPUOffload] Reloaded %.2f MB (CPU → GPU)", reloaded_bytes / (1024**2) - ) diff --git a/build/torch210-cxx11-cu128-x86_64-linux/distributed/utils.py b/build/torch210-cxx11-cu128-x86_64-linux/distributed/utils.py deleted file mode 100644 index 683096bcd50489e7a09959753fdbc29ba0f3b6a6..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-x86_64-linux/distributed/utils.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import (Placement, Shard, - _StridedShard) - - -def _is_shard(placement: Placement) -> bool: - """Check if a placement is a shard type (Shard or _StridedShard). - - In PyTorch 2.10+, _StridedShard no longer inherits from Shard, so - ``placement.is_shard()`` returns False for _StridedShard. This helper - handles both old and new hierarchies. - """ - return isinstance(placement, (Shard, _StridedShard)) - - -def get_slices_of_dtensor( - target: DTensor | torch.Tensor, - local_rank: int, - shard_mesh: DeviceMesh, - shard_placements: tuple[Placement], -) -> tuple[slice | torch.Tensor, ...]: - """ - Get per-dimension indices for a given rank's shard of the target tensor. - - Uses ``Shard.local_shard_size_and_offset`` and - ``_StridedShard.local_shard_size_and_offset`` for correct handling of - both contiguous and strided (non-contiguous) sharding. - - Args: - target (DTensor | torch.Tensor): The target tensor (for its shape). - local_rank (int): The local rank within the shard group. - shard_mesh (DeviceMesh): The shard mesh (only shard dimensions). - shard_placements (tuple[Placement]): The shard placements. - - Returns: - A tuple of indices (one per tensor dim). Each element is either: - - A ``slice`` (for contiguous or unsharded dims) - - A 1-D ``torch.LongTensor`` of indices (for strided sharding) - """ - - # find the global rank of the local rank in the shard mesh - rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank] - - rank_coords = (shard_mesh.mesh == rank).nonzero() - - assert len(rank_coords) == 1 - rank_coords = tuple(rank_coords[0].tolist()) - - assert len(rank_coords) == len(shard_placements) - - # Track per-shard-dim indices. - # None means "not yet sharded on this dim". - dim_indices: dict[int, torch.Tensor] = {} - - # Caution: Assuming replicate-to-shard of the shard mesh goes with - # left-to-right sharding. This is ensured by the sorting logic of - # construct_shard_mesh function. - for mesh_dim_idx, (rank_coord, placement) in enumerate( - zip(rank_coords, shard_placements)): - assert _is_shard(placement) - - num_chunks = shard_mesh.mesh.shape[mesh_dim_idx] - shard_dim = placement.dim - - # Current effective size on this dim (may already be sub-sharded) - if shard_dim in dim_indices: - curr_size = len(dim_indices[shard_dim]) - else: - curr_size = target.size()[shard_dim] - - # Compute indices for this level of sharding - if isinstance(placement, _StridedShard): - _shard_size, offsets = _StridedShard.local_shard_size_and_offset( - placement, - curr_size, - num_chunks, - rank_coord, - return_first_offset=False) - new_indices = torch.tensor(offsets, dtype=torch.long) - else: - shard_size, offset = Shard.local_shard_size_and_offset( - curr_size, num_chunks, rank_coord) - new_indices = torch.arange(offset, - offset + shard_size, - dtype=torch.long) - - # Compose with previous indices on this dim - if shard_dim in dim_indices: - dim_indices[shard_dim] = dim_indices[shard_dim][new_indices] - else: - dim_indices[shard_dim] = new_indices - - # Build result tuple - result: list[slice | torch.Tensor] = [] - for d in range(len(target.size())): - if d not in dim_indices: - result.append(slice(None)) - else: - indices = dim_indices[d] - # Convert contiguous indices to slice for efficiency - if len(indices) > 0: - start = indices[0].item() - expected = torch.arange(start, - start + len(indices), - dtype=torch.long) - if torch.equal(indices, expected): - result.append(slice(start, start + len(indices))) - else: - result.append(indices) - else: - result.append(slice(0, 0)) - - return tuple(result) - - -_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, - ProcessGroup]] = dict() - - -def construct_shard_mesh( - placements: tuple[Placement], - mesh: DeviceMesh, -) -> tuple[DeviceMesh, ProcessGroup, tuple[Placement, ...]]: - """Construct shard sub-mesh and ProcessGroup for all-to-all communication. - - Given a DTensor's placements and device mesh, extracts the "shard group" - — the set of ranks that together hold all shards of the same replica — - and creates a ProcessGroup for all-to-all among them. - - Steps: - 1. Sort placements: Replicate first, then Shard by (dim, granularity). - 2. Permute the mesh tensor to match the sorted order. - 3. Collapse Replicate dims → list of shard sub-meshes (one per replica). - 4. Create/retrieve a cached ProcessGroup for the current rank's sub-mesh. - - Example — 8 GPUs, mesh shape (2, 2, 2), - placements ``[Shard(0), Replicate, _StridedShard(0)]``:: - - Step 1 — Sort: [Replicate, _StridedShard(0), Shard(0)] - Permutation: [1, 2, 0] - - Step 2 — Permute mesh dims by [1, 2, 0]: - Original: Permuted: - [[[0,1],[2,3]], [[[0,2],[1,3]], - [[4,5],[6,7]]] [[4,6],[5,7]]] - - Step 3 — Unbind replicate dim (dim 0), giving 2 shard sub-meshes: - sub-mesh 0 = [[0,2],[1,3]] (replica group 0) - sub-mesh 1 = [[4,6],[5,7]] (replica group 1) - shard_placements = (_StridedShard(0), Shard(0)) - - Step 4 — Rank 0 → ProcessGroup([0,1,4,5]) - Rank 2 → ProcessGroup([2,3,6,7]) - - Returns: - ``(shard_mesh, process_group, shard_placements)`` - """ - my_rank = dist.get_rank() - assert mesh.mesh.device.type == 'cpu' - - # -- Fast path: 1D all-shard mesh → reuse existing PG. ---------------- - # Reuses the mesh's existing ProcessGroup directly, avoiding the - # overhead of dist.new_group(). The standard path below also handles - # subset calls safely via use_local_synchronization=True, but this - # fast path is still beneficial for the common 1D shard case. - if mesh.ndim == 1 and len(placements) == 1 and _is_shard(placements[0]): - key = (*mesh.mesh.shape, *mesh.mesh.flatten().tolist()) - if key not in _ranks_to_dist_cache: - _ranks_to_dist_cache[key] = (mesh, mesh.get_group()) - return (*_ranks_to_dist_cache[key], tuple(placements)) - - mesh_tensor = mesh.mesh.clone() - - # -- Step 1: Sort placements (Replicate first, then Shard by dim). ------ - # _StridedShard comes BEFORE regular Shard on the same dim so that - # get_slices_of_dtensor applies the outer sharding first, matching - # DTensor's left-to-right (outer-to-inner) composition order. - def _sort_key(item): - index, placement = item - assert not placement.is_partial(), "Partial placement not supported" - if placement.is_replicate(): - return (-1, 0, index) - assert _is_shard(placement), f"Unsupported: {type(placement)}" - split = (-1 / placement.split_factor if isinstance( - placement, _StridedShard) else 0) - return (placement.dim, split, index) - - indexed = sorted(enumerate(placements), key=_sort_key) - perm, sorted_placements = zip(*indexed) - - # -- Step 2: Permute mesh to match sorted placement order. -------------- - sorted_mesh = mesh_tensor.permute(perm) - - # -- Step 3: Collapse replicate dims → list of shard sub-meshes. -------- - # E.g. mesh (2, 3, 4, 4) with [R, R, S(0), S(1)] → 6 sub-meshes of (4, 4) - num_rep = sum(1 for p in sorted_placements if p.is_replicate()) - if num_rep > 0: - if num_rep > 1: - sorted_mesh = sorted_mesh.flatten(0, num_rep - 1) - shard_meshes = list(torch.unbind(sorted_mesh, dim=0)) - else: - shard_meshes = [sorted_mesh] - shard_placements = sorted_placements[num_rep:] - assert len(shard_placements) == len(set(shard_placements)) - - # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. -- - # Each rank only creates the group it belongs to, using - # use_local_synchronization=True so that only group members need to - # coordinate. This avoids deadlocks when different PP stages call - # construct_shard_mesh for different parameters. - def _cache_key(t: torch.Tensor) -> tuple: - return (*t.shape, *t.flatten().tolist()) - - my_key = None - for sm in shard_meshes: - if (my_rank == sm).any().item(): - key = _cache_key(sm) - assert my_key is None, "Rank appears in multiple shard groups" - my_key = key - if key not in _ranks_to_dist_cache: - pg = dist.new_group(sm.flatten().tolist(), - use_local_synchronization=True) - _ranks_to_dist_cache[key] = ( - DeviceMesh(device_type="cuda", mesh=sm), - pg, - ) - - return (*_ranks_to_dist_cache[my_key], shard_placements) diff --git a/build/torch210-cxx11-cu128-x86_64-linux/matmul_transpose_triton.py b/build/torch210-cxx11-cu128-x86_64-linux/matmul_transpose_triton.py deleted file mode 100644 index 792de23d82c3fb45fe33d397ab9b76a0787259d0..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-x86_64-linux/matmul_transpose_triton.py +++ /dev/null @@ -1,122 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], - restore_value=['y'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -@torch.library.custom_op("muon::matmul_transpose_assign", - mutates_args=("d_out", )) -def matmul_transpose_assign(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """Compute d_out = d_in @ d_in.T using an optimized Triton kernel.""" - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -@matmul_transpose_assign.register_fake -def _(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """FakeTensor impl: d_out is already allocated, mutation is declared.""" - pass diff --git a/build/torch210-cxx11-cu128-x86_64-linux/metadata.json b/build/torch210-cxx11-cu128-x86_64-linux/metadata.json deleted file mode 100644 index c55a35717622f1dd5c8ba376ea3a814cbcc10d78..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-x86_64-linux/metadata.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch210-cxx11-cu128-x86_64-linux/muon.py b/build/torch210-cxx11-cu128-x86_64-linux/muon.py deleted file mode 100644 index 14c0e22471fa6d47a51ed95e0e0c341dc18d5194..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-x86_64-linux/muon.py +++ /dev/null @@ -1,1068 +0,0 @@ -import logging -import types -from collections import defaultdict -from typing import Any - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Replicate, Shard -from torch.profiler import record_function - -from .adamw import _placement_cache, _tensor_cache, step_adamw -from .async_utils import run_pipeline -from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho, - get_default_muon_param_groups, is_expert_param, update_p) -from .cpu_offload import CPUOffloadPool -from .distributed.utils import (_is_shard, construct_shard_mesh, - get_slices_of_dtensor) -from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO, - _zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5_batched) -from .pipeline import muon_chunk_pipeline, prelaunch_first_gather -from .qk_clip import compute_scales, get_qk_clip_info, qk_clip - -logger = logging.getLogger(__name__) - - -def _expand_expert_params(names, params, expert_keys): - """Expand expert params by splitting on dim 0 (expert dimension). - - Params whose name matches any key in ``expert_keys`` are treated as - expert-parallel tensors. Their outermost dimension is the expert - dimension: an ``(E, out, in)`` tensor becomes ``E`` separate 2D - ``nn.Parameter`` views so that in-place updates propagate back to - the original storage. - - Non-expert params with ``ndim > 2`` trigger an ``AssertionError`` — - if they are expert params, their key must be added to ``expert_keys``. - - The grad must already be set on each expert param (e.g. after momentum). - - For DTensor expert params, placements that shard on dim 0 (expert dim) - are consumed by the split. Non-dim-0 shard placements (e.g. TP) are - preserved: each 2D slice is wrapped as a DTensor on the corresponding - submesh so the parallel pipeline handles the TP communication. - """ - expanded_names = [] - expanded_params = [] - - for n, p in zip(names, params): - is_expert = is_expert_param(n, expert_keys) - is_dtensor = isinstance(p.data, DTensor) - - if is_expert: - if is_dtensor: - logger.debug( - "[expand_expert] %s: expert DTensor, shape=%s, " - "placements=%s, mesh=%s, local_shape=%s", n, p.shape, - p.placements, p.device_mesh.mesh_dim_names, - p.to_local().shape) - else: - logger.debug( - "[expand_expert] %s: expert plain tensor, shape=%s", n, - p.data.shape) - - if not is_expert: - assert p.data.ndim <= 2, ( - f"Param {n} has ndim={p.data.ndim} but does not match " - f"expert_keys={expert_keys}. If this is an expert param, " - f"add its key to expert_keys.") - expanded_names.append(n) - expanded_params.append(p) - continue - - g = p.grad - assert g is not None, ( - f"Expert param {n} must have grad set before expansion") - - tp_mesh = None - tp_placements_2d = None - - if is_dtensor: - local_data = p.to_local() - local_grad = g.to_local() if isinstance(g, DTensor) else g - - # Find non-dim-0 shard placements (e.g. TP sharding). - # After splitting on dim 0, Shard(k) becomes Shard(k-1). - tp_dim_indices = [] - tp_placements_2d = [] - for i, pl in enumerate(p.placements): - if _is_shard(pl) and pl.dim != 0: - tp_dim_indices.append(i) - tp_placements_2d.append(Shard(pl.dim - 1)) - - if tp_dim_indices: - tp_dim_names = tuple(p.device_mesh.mesh_dim_names[i] - for i in tp_dim_indices) - if len(tp_dim_names) == 1: - tp_mesh = p.device_mesh[tp_dim_names[0]] - else: - tp_mesh = p.device_mesh[tp_dim_names] - else: - local_data = p.data - local_grad = g - - # Expand: split dim 0, reshape each slice to 2D. - num_local_experts = local_data.shape[0] - for i in range(num_local_experts): - slice_data = local_data[i] - slice_grad = local_grad[i] - - if tp_mesh is not None: - # Wrap as DTensor on TP submesh so the pipeline handles - # TP communication (gather/scatter across TP ranks). - dt_data = DTensor.from_local(slice_data, - device_mesh=tp_mesh, - placements=tp_placements_2d) - dt_grad = DTensor.from_local(slice_grad, - device_mesh=tp_mesh, - placements=tp_placements_2d) - expert_param = torch.nn.Parameter(dt_data, requires_grad=False) - expert_param.grad = dt_grad - else: - expert_param = torch.nn.Parameter(slice_data, - requires_grad=False) - expert_param.grad = slice_grad - - expanded_names.append(f"{n}[{i}]") - expanded_params.append(expert_param) - - p.grad = None # allow expert grad storage to be freed after pipeline - - return expanded_names, expanded_params - - -class Muon(torch.optim.Optimizer): - """ - Muon - MomentUm Orthogonalized by Newton-schulz - - Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- - processing step, in which each 2D parameter's update is replaced with the nearest orthogonal - matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has - the advantage that it can be stably run in bfloat16 on the GPU. - - Some warnings: - - We believe this optimizer is unlikely to work well for training with small batch size. - - We believe it may not work well for finetuning pretrained models, but we haven't tested this. - - Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. - lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) - momentum: The momentum used by the internal SGD. (0.95 is a good default) - nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) - ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. - Parameters that are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW instead. - adamw_lr: The learning rate for the internal AdamW. - adamw_betas: The betas for the internal AdamW. - adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - warmup_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher warmup_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. - chunk_size : Batch size of parameters to process in each - all2all gather/compute/scatter step. - Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified. - use_distributed_muon: Use distributed muon by Liu et al. (2024). - For testing purpose only. - expert_keys: List of strings to identify expert-parallel parameters. - If any key appears in a parameter's name, its outermost - dimension is treated as the expert dimension and expanded - into per-expert 2D params for Muon. For example, - ``expert_keys=["experts"]`` matches any param whose name - contains "experts". 3D+ params not matched by any key - will raise an error. - """ - - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config=None, - warmup_step=5, - chunk_size=-1, - use_distributed_muon=False, - expert_keys=None): - defaults = dict( - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - nesterov=nesterov, - ns_steps=ns_steps, - adamw_betas=adamw_betas, - adamw_eps=adamw_eps, - none_grad=none_grad, - use_muon=True, - ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) - super().__init__(params, defaults) - - self.debug = debug - self.clip_config = clip_config if clip_config is not None else { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100, - } - self.warmup_step = warmup_step - self.chunk_size = chunk_size - self.use_distributed_muon = use_distributed_muon - self.expert_keys = expert_keys - self.cpu_offload = False - self._cpu_offload_pool: CPUOffloadPool | None = None - self._offload_initialized = False - self._parallel_cache: dict[tuple[str, ...], dict] = {} - self._expert_expand_cache: dict[tuple[int, ...], dict] = {} - - def _calc_flops(self, G, steps): - assert len(G.shape) == 2 - M, N = G.shape - if M > N: - M, N = N, M - - return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) - - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - p.placements, p.device_mesh) - - return shard_mesh, shard_pg, shard_placements - - def init_state_and_assign_params(self, names, params, group, qk_logits): - param_to_state = {} - param_to_flops = {} - - total_flops = 0 - for p in params: - g = p.grad - if g is None: - continue - assert g.ndim == 2, "Muon only supports 2D parameters." - - flops = self._calc_flops(g, group["ns_steps"]) - param_to_flops[id(p)] = flops - total_flops += flops - - if self.debug: - logger.debug("Total TFLOPs for Muon: %.2f TFLOPs", - total_flops / 1e12) - - paired = list(zip(names, params)) - - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) - - round_robin = 0 - mesh = ordered_params[0].device_mesh - placements = ordered_params[0].placements - - shard_mesh, shard_pg, shard_placements = self.get_shard_mesh( - ordered_params[0]) - shard_mesh_flattened = shard_mesh.mesh.flatten() - num_ranks = dist.get_world_size(group=shard_pg) - - for n, p in zip(ordered_names, ordered_params): - if mesh != p.device_mesh: - raise ValueError("All parameters must be on the same mesh.") - if placements != p.placements: - raise ValueError("All parameters must have same placements.") - - worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks - round_robin = (round_robin + 1) % len(shard_mesh_flattened) - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - # Precompute per-rank indices and numels for all-to-all. - rank_indices: dict[int, tuple] = {} - rank_numels: dict[int, int] = {} - for r in range(num_ranks): - indices = get_slices_of_dtensor(p, r, shard_mesh, - shard_placements) - rank_indices[r] = indices - numel = 1 - for idx, dim_size in zip(indices, p.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) - rank_numels[r] = numel - - param_to_state[id(p)] = _muon_state( - worker_rank=worker_rank, - process_group=shard_pg, - rank_indices=rank_indices, - rank_numels=rank_numels, - name=n, - qk_clip_state=qk_clip_state, - ) - - return param_to_state, ordered_params - - def base(self, names, params, group, lr, weight_decay, qk_logits): - # Momentum is already applied by _step_muon before this method. - for n, p in zip(names, params): - g = p.grad - if g is None: - continue - - u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) - - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - def distributed_muon( - self, - names: list[str], - params: list[torch.nn.Parameter], - group: dict[str, Any], - lr: float, - weight_decay: float, - qk_logits: list[torch.Tensor | DTensor] | None, - ): - """Batched Distributed Muon — for testing/correctness verification only. - - Uses all-gather to reconstruct full tensors, computes Newton-Schulz on - the full grad, then slices back to local shards. This is simpler but - slower than the parallel pipeline (all2all) path, so it serves as a - reference implementation for verifying correctness. - """ - with record_function("distributed_muon"): - # Momentum is already applied by _step_muon before this method. - ns_steps = group["ns_steps"] - - # Separate plain tensors (no communication) from DTensors. - plain_names, plain_params = [], [] - dtensor_names, dtensor_params = [], [] - for n, p in zip(names, params): - if p.grad is None: - continue - if isinstance(p.data, DTensor): - dtensor_names.append(n) - dtensor_params.append(p) - else: - plain_names.append(n) - plain_params.append(p) - - # Process plain tensors per-param (no communication). - for n, p in zip(plain_names, plain_params): - u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE), - steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, - qk_logits) - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - if not dtensor_params: - return - - # Group DTensors by (placements, mesh) for batched all-gather. - placement_groups: dict[tuple, - tuple[list, - list]] = defaultdict(lambda: ([], [])) - for n, p in zip(dtensor_names, dtensor_params): - key = (p.placements, p.device_mesh) - placement_groups[key][0].append(n) - placement_groups[key][1].append(p) - - logger.info( - "distributed_muon: %d placement groups, %d total dtensors", - len(placement_groups), len(dtensor_params)) - - for (placements, mesh), (grp_names, - grp_params) in placement_groups.items(): - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - placements, mesh) - rank = dist.get_rank(shard_pg) - world_size = dist.get_world_size(shard_pg) - - logger.info(" group: %d params, placements=%s, world_size=%d", - len(grp_params), placements, world_size) - - # Separate params that can be batched (all shard dims evenly - # divisible) from those needing per-param full_tensor - # (e.g. MoE gate weights with fewer rows than shard ranks). - # all_gather_into_tensor requires equal buffer sizes across - # ranks, so uneven splits must use DTensor full_tensor(). - batch_names, batch_params = [], [] - single_names, single_params = [], [] - for n, p in zip(grp_names, grp_params): - even = all(p.shape[pl.dim] % - shard_mesh.mesh.shape[dim_idx] == 0 - for dim_idx, pl in enumerate(shard_placements)) - if even: - batch_names.append(n) - batch_params.append(p) - else: - single_names.append(n) - single_params.append(p) - - # Process uneven-split params per-param via full_tensor(). - for n, p in zip(single_names, single_params): - with record_function("distributed_muon::newton_schulz"): - g_full = p.grad.full_tensor().to(COMM_DTYPE) - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - if not batch_params: - continue - - logger.info(" batched=%d, single=%d", len(batch_params), - len(single_params)) - - # Concat all local grad shards into a single flat buffer. - with record_function("distributed_muon::gather"): - grad_locals = [ - p.grad.to_local().to(COMM_DTYPE).flatten() - for p in batch_params - ] - numels = [g.numel() for g in grad_locals] - grad_concat = torch.cat(grad_locals) - del grad_locals - - # Single all-gather (replaces N separate full_tensor). - grad_gathered = torch.empty( - grad_concat.numel() * world_size, - dtype=COMM_DTYPE, - device="cuda", - ) - dist.all_gather_into_tensor(grad_gathered, - grad_concat, - group=shard_pg) - - total_numel = grad_concat.numel() - del grad_concat - - # Precompute per-param offsets within the concat buffer. - offsets = [] - off = 0 - for ne in numels: - offsets.append(off) - off += ne - - # Per-param: reconstruct full grad → NS → local update. - for i, (n, p) in enumerate(zip(batch_names, batch_params)): - with record_function("distributed_muon::newton_schulz"): - g_full = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - for r in range(world_size): - r_start = r * total_numel + offsets[i] - shard = grad_gathered[r_start:r_start + numels[i]] - indices = get_slices_of_dtensor( - p, r, shard_mesh, shard_placements) - g_full[indices] = shard.reshape( - g_full[indices].shape) - - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - def _setup_parallel(self, names, params, group, qk_logits): - """Compute (or retrieve cached) parallel pipeline metadata. - - Returns: - (ordered_params, param_to_state, rank, chunk_size) - """ - cache_key = tuple(names) - - if cache_key not in self._parallel_cache: - # First call: compute metadata and populate cache. - param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - shard_pg = param_to_state[id(ordered_params[0])].process_group - rank = dist.get_rank(group=shard_pg) - - if self.chunk_size == -1: - shard_ranks = dist.get_world_size(shard_pg) - chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO - elif self.chunk_size > 0: - chunk_size = self.chunk_size - else: - raise ValueError( - "chunk_size must be -1 or a positive integer.") - - ordered_names = [ - param_to_state[id(p)].name for p in ordered_params - ] - name_to_state = { - param_to_state[id(p)].name: param_to_state[id(p)] - for p in ordered_params - } - self._parallel_cache[cache_key] = { - 'ordered_names': ordered_names, - 'name_to_state': name_to_state, - 'rank': rank, - 'chunk_size': chunk_size, - } - else: - # Cached path: rebuild param_to_state with current id(p) keys. - cache = self._parallel_cache[cache_key] - rank = cache['rank'] - chunk_size = cache['chunk_size'] - - name_to_param = dict(zip(names, params)) - ordered_params = [name_to_param[n] for n in cache['ordered_names']] - - param_to_state = {} - for p, n in zip(ordered_params, cache['ordered_names']): - cached_state = cache['name_to_state'][n] - param_to_state[id(p)] = _muon_state( - worker_rank=cached_state.worker_rank, - process_group=cached_state.process_group, - rank_indices=cached_state.rank_indices, - rank_numels=cached_state.rank_numels, - name=n, - qk_clip_state=get_qk_clip_info(self.clip_config, n, - qk_logits), - ) - - return ordered_params, param_to_state, rank, chunk_size - - def parallel(self, - names, - params, - group, - lr, - weight_decay, - qk_logits, - prelaunch_gather=None): - """ - Perform a parallel optimization step using Muon. - - Parameters are chunked and each chunk is processed by a - :func:`muon_chunk_pipeline` generator. :func:`run_pipeline` - interleaves multiple chunks so that communication and computation - overlap across chunks (the same overlap previously achieved by the - warmup + main-loop index scheduling). - - If ``prelaunch_gather`` is provided, it is passed to the first - chunk's generator to skip re-launching the already in-flight - A2A gather. - """ - - # Momentum is already applied by _step_muon before this method. - - ordered_params, param_to_state, rank, chunk_size = ( - self._setup_parallel(names, params, group, qk_logits)) - - def pipelines(): - first = True - for start in range(0, len(ordered_params), chunk_size): - chunk = ordered_params[start:start + chunk_size] - if chunk: - kwargs = dict( - params=chunk, - param_to_state=param_to_state, - rank=rank, - ns_steps=group["ns_steps"], - lr=lr, - weight_decay=weight_decay, - none_grad=group["none_grad"], - ) - if first and prelaunch_gather is not None: - kwargs['prelaunch_gather'] = prelaunch_gather - first = False - yield muon_chunk_pipeline(**kwargs) - - with record_function("muon::pipeline"): - run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1) - - def _step_muon(self, group, qk_logits=None): - params = group["params"] - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - # Apply momentum to all params before routing/expansion. - # Batched using _foreach_* ops (compiled, fullgraph=True). - with record_function("muon::momentum"): - active_params = [p for p in params if p.grad is not None] - if active_params: - # Ensure momentum buffers exist (avoid zeros_like when already present). - for p in active_params: - if "momentum_buffer" not in self.state[p]: - self.state[p]["momentum_buffer"] = torch.zeros_like( - p.grad) - - # Extract local tensors for compiled batch function. - local_grads = [ - p.grad._local_tensor - if isinstance(p.grad, DTensor) else p.grad - for p in active_params - ] - local_bufs = [ - self.state[p]["momentum_buffer"]._local_tensor - if isinstance(self.state[p]["momentum_buffer"], DTensor) - else self.state[p]["momentum_buffer"] - for p in active_params - ] - - # Wrap momentum as tensor for torch.compile. - batch_pre_ortho(local_grads, local_bufs, - torch.tensor(momentum), group["nesterov"]) - - # For non-nesterov, the result is the momentum buffer. - if not group["nesterov"]: - for p in active_params: - p.grad = self.state[p]["momentum_buffer"] - - # Identify batched experts for deferred NS. - # Detection is cheap (condition checks only); actual NS compute is - # deferred so it can overlap with the first chunk's A2A gather. - deferred_expert_work = [] - if self.expert_keys: - batched_expert_indices = [] - for i, (n, p) in enumerate(zip(names, params)): - if not (is_expert_param(n, self.expert_keys) - and p.grad is not None): - continue - # Eligible: plain tensor, or DTensor with no non-dim-0 shards. - if isinstance(p.data, DTensor): - has_tp = any( - _is_shard(pl) and pl.dim != 0 for pl in p.placements) - if has_tp: - continue - batched_expert_indices.append(i) - - if batched_expert_indices: - # Save refs for deferred NS; free grads from param list. - for i in batched_expert_indices: - p = params[i] - g = p.grad - local_g = (g._local_tensor - if isinstance(g, DTensor) else g) - local_data = (p.data._local_tensor if isinstance( - p.data, DTensor) else p.data) - deferred_expert_work.append((local_data, local_g)) - p.grad = None - - # Remove batched experts from lists before expansion. - keep = sorted( - set(range(len(params))) - set(batched_expert_indices)) - names = [names[i] for i in keep] - params = [params[i] for i in keep] - - def _run_deferred_expert_ns(): - """Execute deferred batched expert NS.""" - if not deferred_expert_work: - return - with record_function("muon::batched_expert_ns"): - ns_steps = group["ns_steps"] - for local_data, local_g in deferred_expert_work: - u = zeropower_via_newtonschulz5_batched( - local_g.to(COMM_DTYPE), steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:]) - local_data.mul_(1 - lr * weight_decay) - local_data.add_(u, alpha=-adjusted_lr) - - # Expand expert params by splitting on dim 0. - logger.debug("[_step_muon] before expand: %d params, expert_keys=%s", - len(params), self.expert_keys) - if self.expert_keys: - cache_key = tuple(id(p) for p in params) - cache = self._expert_expand_cache.get(cache_key) - - if cache is None: - # Cold path: full expansion + build cache metadata. - exp_names, exp_params = _expand_expert_params( - names, params, self.expert_keys) - - # Build per-expert-group info for hot-path grad updates. - grad_info = [] - exp_idx = 0 - for orig_idx, (n, p) in enumerate(zip(names, params)): - if not is_expert_param(n, self.expert_keys): - exp_idx += 1 - continue - - is_dt = isinstance(p.data, DTensor) - num_experts = (p.to_local() if is_dt else p.data).shape[0] - - # Detect TP mesh from the first expanded expert param. - tp_mesh = None - tp_pls = None - sample = exp_params[exp_idx] - if isinstance(sample.data, DTensor): - tp_mesh = sample.data.device_mesh - tp_pls = list(sample.data.placements) - - grad_info.append((orig_idx, num_experts, exp_idx, is_dt, - tp_mesh, tp_pls)) - exp_idx += num_experts - - self._expert_expand_cache[cache_key] = { - 'names': exp_names, - 'params': exp_params, - 'grad_info': grad_info, - } - names, params = exp_names, exp_params - else: - # Hot path: reuse cached params, only update expert grads. - for (orig_idx, num_experts, exp_start, is_dt, tp_mesh, - tp_pls) in cache['grad_info']: - p = params[orig_idx] - g = p.grad - local_grad = (g.to_local() - if is_dt and isinstance(g, DTensor) else g) - for i in range(num_experts): - expert_p = cache['params'][exp_start + i] - sg = local_grad[i] - if tp_mesh is not None: - expert_p.grad = DTensor.from_local( - sg, device_mesh=tp_mesh, placements=tp_pls) - else: - expert_p.grad = sg - p.grad = None - - names = cache['names'] - params = cache['params'] - else: - names, params = _expand_expert_params(names, params, - self.expert_keys) - logger.debug("[_step_muon] after expand: %d params", len(params)) - - param_dtensors = [] - name_dtensors = [] - - param_tensors = [] - name_tensors = [] - - # distributed_muon is a reference implementation for testing only. - # The parallel pipeline (all2all) path below is the production path. - if self.use_distributed_muon: - _run_deferred_expert_ns() - self.distributed_muon(names=names, - params=params, - group=group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits) - return - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - logger.debug( - "[route] %s → base (DTensor all-Replicate), " - "shape=%s, placements=%s", n, p.shape, p.placements) - param_tensors.append(p) - name_tensors.append(n) - else: - logger.debug( - "[route] %s → parallel (DTensor), shape=%s, " - "placements=%s, mesh=%s", n, p.shape, p.placements, - p.device_mesh.mesh_dim_names) - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - logger.debug("[route] %s → base (plain tensor), shape=%s", n, - p.data.shape) - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError(f"Unsupported parameter type: {type(p.data)}") - - logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, " - f"{len(param_tensors)} Tensors → base") - - def group_dtensors(dtensors, names): - # To support different placements, we group parameters by placements - # and run parallel Muon on each group. - - placement_to_params = defaultdict(lambda: ([], [])) - - assert len(dtensors) == len(names) - for p, n in zip(dtensors, names): - placement_to_params[tuple([p.placements, - p.device_mesh])][0].append(n) - placement_to_params[tuple([p.placements, - p.device_mesh])][1].append(p) - return placement_to_params - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - dtensor_group = group_dtensors(param_dtensors, name_dtensors) - - # Pre-launch the first chunk's A2A gather so that the NCCL - # communication overlaps with the (deferred) batched expert NS - # compute on the default CUDA stream. - prelaunch = None - if deferred_expert_work: - first_names, first_params = next(iter(dtensor_group.values())) - ordered, pts, rnk, csz = self._setup_parallel( - first_names, first_params, group, qk_logits) - first_chunk = ordered[:csz] - if first_chunk: - prelaunch = prelaunch_first_gather(first_chunk, pts, rnk, - group["none_grad"]) - - _run_deferred_expert_ns() - - first_group = True - for _, (names, params) in dtensor_group.items(): - pg = prelaunch if first_group else None - first_group = False - self.parallel( - names, - params, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - prelaunch_gather=pg, - ) - else: - _run_deferred_expert_ns() - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - ) - - def _register_states_for_offload(self): - """Register all optimizer state tensors with the CPU offload pool. - - Called once after the first step when states have been lazily created. - Offloads all param states (momentum buffers for Muon, moment1/moment2 - for AdamW) to free GPU memory between steps. - """ - pool = self._cpu_offload_pool - tracked = 0 - for group in self.param_groups: - for p in group["params"]: - if p not in self.state: - continue - state = self.state[p] - if group.get("use_muon", False): - if "momentum_buffer" in state: - pool.track(state["momentum_buffer"]) - tracked += 1 - else: - if "moment1" in state: - pool.track(state["moment1"]) - if "moment2" in state: - pool.track(state["moment2"]) - tracked += 1 - logger.info("[CPUOffload] Registered %d param states for offload", - tracked) - - @torch.no_grad - def step(self, closure=None, qk_logits=None): - """Perform a single optimization step. - - Args: - closure (Callable, optional): A closure that reevaluates the model - and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). - """ - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # H2D: reload optimizer states from CPU before computation. - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload() - - logger.debug("[Muon.step] expert_keys=%s, %d param groups", - self.expert_keys, len(self.param_groups)) - - for i, group in enumerate(self.param_groups): - if group["use_muon"]: - logger.debug("[Muon.step] group %d: use_muon=True, %d params", - i, len(group["params"])) - self._step_muon(group, qk_logits=qk_logits) - else: - logger.debug( - "[Muon.step] group %d: use_muon=False (AdamW), %d params", - i, len(group["params"])) - step_adamw(self.state, group) - - # D2H: offload optimizer states to CPU after computation. - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - return loss - - # ------------------------------------------------------------------ - # CPU offload public helpers - # ------------------------------------------------------------------ - - def turn_on_cpu_offload(self): - """Enable CPU offload for optimizer states.""" - if self.cpu_offload: - return - logger.info("[Muon] turn_on_cpu_offload") - self.cpu_offload = True - if not self.state: - return - self._cpu_offload_pool = CPUOffloadPool() - self._offload_initialized = False - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_off_cpu_offload(self): - """Disable CPU offload and keep optimizer states resident on GPU.""" - if not self.cpu_offload: - return - logger.info("[Muon] turn_off_cpu_offload") - if self._offload_initialized: - self._cpu_offload_pool.reload() - torch.cuda.current_stream().synchronize() - self._cpu_offload_pool = None - self._offload_initialized = False - self.cpu_offload = False - - # ------------------------------------------------------------------ - # Checkpoint support for cpu_offload - # ------------------------------------------------------------------ - - def state_dict(self) -> dict: - if self.cpu_offload: - raise RuntimeError( - "Muon.state_dict() requires turn_off_cpu_offload() before checkpoint save." - ) - return super().state_dict() - - def load_state_dict(self, state_dict: dict) -> None: - if self.cpu_offload: - raise RuntimeError( - "Muon.load_state_dict() requires turn_off_cpu_offload() before checkpoint load." - ) - super().load_state_dict(state_dict) - - # Invalidate adamw.py's module-level tensor caches so that - # the next step rebuilds them with the newly loaded state tensors. - _placement_cache.clear() - _tensor_cache.clear() diff --git a/build/torch210-cxx11-cu128-x86_64-linux/newton_schulz.py b/build/torch210-cxx11-cu128-x86_64-linux/newton_schulz.py deleted file mode 100644 index d939264b69a34e7a3fa78859f34dc265a1159d59..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-x86_64-linux/newton_schulz.py +++ /dev/null @@ -1,240 +0,0 @@ -from itertools import repeat -from math import inf, sqrt - -import numpy as np -import torch - -from .matmul_transpose_triton import matmul_transpose_assign - -COMM_DTYPE = torch.bfloat16 -DEFAULT_CHUNK_SIZE_RATIO = 4 - - -def _optimal_quintic(l, u, max_iter=1000): - """ - Use the simplified Remez algorithm to find the optimal odd quintic approximant - to the constant function x -> 1 over the interval [l, u]. - - Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum - approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the - two interior equioscillation nodes q, r until convergence. Returns the - closed-form equioscillating solution when l ≈ u. - - Raises ValueError if any intermediate value (a, b, c, E, q, r) is non-finite - (NaN or inf). Raises RuntimeError if convergence is not reached within - max_iter iterations. - """ - assert 0 <= l <= u - if 1 - 5e-6 <= l / u: - return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5) - q = (3 * l + u) / 4 - r = (l + 3 * u) / 4 - E = inf - for _ in range(max_iter): - old_E = E - LHS = np.array( - [ - [l, l**3, l**5, 1], - [q, q**3, q**5, -1], - [r, r**3, r**5, 1], - [u, u**3, u**5, -1], - ] - ) - a, b, c, E = np.linalg.solve(LHS, np.ones(4)) - if not np.all(np.isfinite([a, b, c, E])): - raise ValueError( - f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}" - ) - q, r = np.sqrt( - (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / (10 * c) - ) - if not np.all(np.isfinite([q, r])): - raise ValueError(f"_optimal_quintic: non-finite node update q={q}, r={r}") - if abs(old_E - E) <= 1e-15: - break - else: - raise RuntimeError( - f"_optimal_quintic: did not converge after {max_iter} iterations" - ) - return float(a), float(b), float(c) - - -def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0): - """ - Compute the Polar Express coefficient series for `num_iters` quintic iterations. - - Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that - compose to map singular values from [l, 1] toward 1. At each step: - 1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion` - prevents near-zero singular values from stalling by raising the effective - lower bound; if it is active (cushion*u > l), the coefficients are - rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u]. - 2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the - last iteration, providing numerical headroom at the cost of a slightly slower - final convergence step. - 3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1). - - Returns a list of (a, b, c) tuples, one per iteration. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - u = 1 - assert 0 <= l <= u - safety_factor = 1 + safety_factor_eps - coefficients = [] - for iter in range(num_iters): - a, b, c = _optimal_quintic(max(l, cushion * u), u) - if cushion * u > l: - pl = a * l + b * l**3 + c * l**5 - pu = a * u + b * u**3 + c * u**5 - rescaler = 2 / (pl + pu) - a *= rescaler - b *= rescaler - c *= rescaler - if iter < num_iters - 1: - a /= safety_factor - b /= safety_factor**3 - c /= safety_factor**5 - coefficients.append((a, b, c)) - l = a * l + b * l**3 + c * l**5 - u = 2 - l - return coefficients - - -# Precomputed Polar Express coefficients (a, b, c) for 10 quintic Newton-Schulz -# iterations. Each tuple is the minimax-optimal (Remez/equioscillation) odd quintic -# approximant to x->1 over the current singular-value interval, computed once at -# import time and reused across all optimizer steps. -# -# Contrast with the former hardcoded NS coefficients (5 fixed tuples): -# - Former: empirically tuned to maximize slope at zero; did not converge -# singular values to 1, yielding US'V^T with S' ~ Uniform(0.5, 1.5) instead -# of the true polar factor UV^T. -# - Polar Express: analytically optimal per step, adapting to the shrinking -# singular-value interval [l, u] as iterations progress; converges all -# singular values to 1, producing the exact polar factor UV^T. -_coeffs_list = _optimal_composition( - l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02 -) - - -# This code is adapted from: -# KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py) -# NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress) -# matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon) -@torch.no_grad() -def _zeropower_via_newtonschulz5(G, steps): - """ - Compute the polar factor of G via the Polar Express method. - - Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c) - are the Polar Express coefficients from `_coeffs_list`. Each step is the - optimal odd quintic approximant to x -> 1 over the current singular-value - interval, minimizing the maximum approximation error (Remez / minimax criterion). - The composition maps singular values from [l, 1] to near 1, producing the - polar factor (orthogonal factor in the polar decomposition G = UP). - - `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2, - cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE - X = G # no manual typecast - - if G.size(0) > G.size(1): - X = X.T - - X = X / (X.norm() + 1e-7) - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - # Perform the NS iterations - for a, b, c in hs: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(0) > G.size(1): - X = X.T - - return X - - -@torch.no_grad() -def _zeropower_via_newtonschulz5_batched(G, steps): - """Batched polar factor computation for 3D (E, out, in) tensors. - - Same algorithm as ``_zeropower_via_newtonschulz5`` but uses - ``torch.bmm`` / ``torch.baddbmm`` instead of the 2D Triton kernel, - processing all E expert matrices in a single batched call. - """ - assert len(G.shape) == 3 - assert G.dtype == COMM_DTYPE - X = G - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - # Per-expert Frobenius norm. - X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) - - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - for a, b, c in hs: - buf1 = torch.bmm(X, X.transpose(-2, -1)) - buf2 = torch.bmm(buf1, buf1.transpose(-2, -1)) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.baddbmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - return X - - -_ns_per_shape: dict[tuple[int, ...], callable] = {} -_use_compile = True - - -def set_ns_compile(enabled: bool): - """Toggle torch.compile for Newton-Schulz iteration.""" - global _use_compile - _use_compile = enabled - - -def zeropower_via_newtonschulz5(G, steps=5): - if not _use_compile: - return _zeropower_via_newtonschulz5(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() - - -def zeropower_via_newtonschulz5_batched(G, steps=5): - """Compile-cached batched Newton-Schulz for 3D expert tensors.""" - if not _use_compile: - return _zeropower_via_newtonschulz5_batched(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile( - _zeropower_via_newtonschulz5_batched, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() diff --git a/build/torch210-cxx11-cu128-x86_64-linux/optimizer/__init__.py b/build/torch210-cxx11-cu128-x86_64-linux/optimizer/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-x86_64-linux/optimizer/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch210-cxx11-cu128-x86_64-linux/pipeline.py b/build/torch210-cxx11-cu128-x86_64-linux/pipeline.py deleted file mode 100644 index c0c2d515856182d8d15ad27dd4e4e093b29397d6..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-x86_64-linux/pipeline.py +++ /dev/null @@ -1,468 +0,0 @@ -import logging -from typing import Generator - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -from .core import _muon_state, adjust_lr_for_muon -from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5 -from .qk_clip import compute_scales - -logger = logging.getLogger(__name__) - -# ====================================================================== -# Stage helpers -# ====================================================================== - - -def _launch_gather( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Allocate gather buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_gather``). - gathered_grads: ``{id(p): empty_tensor}`` for owned params, - ``None`` for non-owned. - recv_counts: Per-source-rank element counts. - """ - # Allocate gathered-grad buffers - gathered_grads: dict[int, torch.Tensor | None] = {} - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - gathered_grads[id(p)] = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - else: - gathered_grads[id(p)] = None - - # Build send buffer – batch grad copies via torch.cat - # (1-2 fused kernels vs N individual narrow().copy_() calls). - send_counts = [0] * num_ranks - for p in params: - state = param_to_state[id(p)] - send_counts[state.worker_rank] += state.rank_numels[rank] - - total_send = sum(send_counts) - if total_send > 0: - # Group grad slices by destination rank in a single pass. - dst_to_grads = [[] for _ in range(num_ranks)] - for p in params: - state = param_to_state[id(p)] - n = state.rank_numels[rank] - if n > 0: - g = p.grad.to_local() - dst_to_grads[state.worker_rank].append(g.reshape(-1)) - - # Flatten in dst order and cat once. - all_slices = [] - for dst in range(num_ranks): - all_slices.extend(dst_to_grads[dst]) - send_buf = torch.cat(all_slices) - if send_buf.dtype != COMM_DTYPE: - send_buf = send_buf.to(COMM_DTYPE) - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += state.rank_numels[src] - recv_counts[src] = total - - recv_buf = torch.empty(sum(recv_counts), dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - logger.debug(f"send_buf size: {send_buf.numel()}, " - f"recv_buf size: {recv_buf.numel()}, " - f"recv_counts: {recv_counts}, " - f"send_counts: {send_counts}, " - f"process_group: {str(process_group)}") - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, gathered_grads, recv_counts - - -def _complete_gather( - recv_buf: torch.Tensor, - recv_counts: list[int], - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - param_to_state: dict[int, _muon_state], - rank: int, -) -> None: - """Reconstruct gathered grads from the recv buffer (in-place).""" - off = 0 - for src in range(len(recv_counts)): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - - indices = state.rank_indices[src] - - shard_view = gathered_grads[id(p)][indices] - n = shard_view.numel() - if n == 0: - continue - - sg = recv_buf.narrow(0, off + inner_off, n) - sg = sg.reshape(shard_view.shape) - gathered_grads[id(p)][indices] = sg - - inner_off += n - assert inner_off == block - off += block - - -def _compute_ns( - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - ns_steps: int, -) -> dict[int, torch.Tensor | None]: - """Run Newton-Schulz orthogonalization on owned parameters. - - Returns: - computed_us: ``{id(p): orthogonalized_update}`` for owned params. - """ - computed_us: dict[int, torch.Tensor | None] = {} - for p in owned_params: - u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps) - gathered_grads[id(p)] = None # free gathered grad - computed_us[id(p)] = u - return computed_us - - -def _launch_scatter( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, - computed_us: dict[int, torch.Tensor | None], -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor], list[int]]: - """Allocate scatter buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_scatter``). - scattered_us: Empty dict, populated by ``_complete_scatter`` with - zero-copy views into ``recv_buf``. - recv_counts: Per-source-rank element counts. - """ - # scattered_us is populated by _complete_scatter with zero-copy views - # into recv_buf, avoiding N empty_like allocations + N copy_ calls. - # Pre-seed entries for params whose local shard is empty (rank_numels == 0) - # so _update_params can iterate all params without KeyError. - scattered_us: dict[int, torch.Tensor] = {} - for p in params: - if param_to_state[id(p)].rank_numels[rank] == 0: - scattered_us[id(p)] = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - # Build send buffer – batch via torch.cat - # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls). - send_counts = [0] * num_ranks - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - for dst_rank in range(num_ranks): - send_counts[dst_rank] += state.rank_numels[dst_rank] - - total_send = sum(send_counts) - if total_send > 0: - # Cache u_full conversions to avoid redundant .to() per dst_rank. - u_fulls = {} - for p in owned_params: - u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous() - - # Collect slices in dst order (matches all-to-all send layout). - all_slices = [] - for dst_rank in range(num_ranks): - for p in owned_params: - state = param_to_state[id(p)] - su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten() - if su.numel() > 0: - all_slices.append(su) - - send_buf = torch.cat(all_slices) if all_slices else torch.empty( - 0, dtype=COMM_DTYPE, device="cuda") - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += state.rank_numels[rank] - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, scattered_us, recv_counts - - -def _complete_scatter( - recv_buf: torch.Tensor, - recv_counts: list[int], - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], -) -> None: - """Populate scattered_us with zero-copy views into recv_buf. - - Instead of pre-allocating tensors and copying, we assign views directly - from ``recv_buf``. This eliminates N ``empty_like`` + N ``copy_`` calls. - The underlying storage of ``recv_buf`` is kept alive through the views - until ``scattered_us`` is cleared after ``_update_params``. - """ - off = 0 - for src in range(len(recv_counts)): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = state.rank_numels[rank] - if n == 0: - continue - - scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - - inner_off += n - - assert inner_off == block - off += block - - -def _update_params( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], - lr: float, - weight_decay: float, -) -> None: - """Apply weight decay, Muon update, and optional QK clipping. - - Uses batched ``_foreach_mul_`` for weight decay and batched - ``_foreach_add_`` for the Muon update, grouping parameters by - adjusted_lr to minimize kernel launches while preserving float32 - precision for the alpha scaling. - """ - if not params: - return - - # Batched weight decay: p *= (1 - lr * wd) — single fused kernel. - p_locals = [p._local_tensor for p in params] - torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay) - - # Group params by adjusted_lr so _foreach_add_ can use a single - # alpha per group (preserves float32 precision for alpha scaling). - lr_groups: dict[float, tuple[list, list]] = {} - for p in params: - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - if adjusted_lr not in lr_groups: - lr_groups[adjusted_lr] = ([], []) - lr_groups[adjusted_lr][0].append(p._local_tensor) - lr_groups[adjusted_lr][1].append(scattered_us[id(p)]) - - for adjusted_lr, (p_group, u_group) in lr_groups.items(): - torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr) - - # QK clipping – applied directly on the local tensor to - # avoid DTensor sharding-propagation issues with _StridedShard. - for p in params: - state = param_to_state[id(p)] - if state.qk_clip_state is None: - continue - scales_full = compute_scales(p, state.qk_clip_state) - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = state.rank_indices[rank][0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - -# ====================================================================== -# Pre-launch helper for overlapping first chunk's gather with other work. -# ====================================================================== - - -@torch.no_grad() -def prelaunch_first_gather( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - none_grad: bool, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Launch the first chunk's A2A gather early for overlap with other compute. - - Call this *before* expensive GPU work (e.g. batched expert NS) so that - the NCCL all-to-all runs concurrently on the NCCL stream while the - default stream executes compute. - - Returns the same 4-tuple that ``_launch_gather`` produces, which should - be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - with record_function("muon::prelaunch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - return work, recv_buf, gathered_grads, recv_counts - - -# ====================================================================== -# Main generator – thin orchestrator that wires stages together. -# ====================================================================== - - -@torch.no_grad() -def muon_chunk_pipeline( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - ns_steps: int, - lr: float, - weight_decay: float, - none_grad: bool, - prelaunch_gather: tuple | None = None, -) -> Generator[None, None, None]: - """Process one chunk of parameters through the full Muon pipeline. - - Stages: gather -> compute (Newton-Schulz) -> scatter -> update. - - Each ``yield`` lets :func:`run_pipeline` interleave other chunks so - that communication and computation overlap across chunks. Async - communication is launched via ``async_op=True`` and completed after - the yield with ``work.wait()``. - - Overlap happens because :func:`run_pipeline` admits one new chunk - per iteration (staggered admission). While chunk *N* does NS - compute on the default CUDA stream, chunk *N+1*'s async all-to-all - runs concurrently on the NCCL stream — no separate ``comm_stream`` - is required. - - If ``prelaunch_gather`` is provided, the gather was already launched - by :func:`prelaunch_first_gather` and we skip launching it again. - - Yields exactly **2** times: - - 1. After launching async all-to-all gather (or immediately if pre-launched). - 2. After launching async all-to-all scatter. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - if prelaunch_gather is not None: - # Gather was pre-launched; none_grad already handled by caller. - work, recv_buf, gathered_grads, recv_counts = prelaunch_gather - else: - # Normal path: launch async gather. - with record_function("muon::launch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - yield # --- YIELD 1: other chunks can launch their gather --- - - with record_function("muon::wait_gather"): - work.wait() - _complete_gather(recv_buf, recv_counts, owned_params, gathered_grads, - param_to_state, rank) - del recv_buf - - # Stage 3: Newton-Schulz orthogonalization. - with record_function("muon::newton_schulz"): - computed_us = _compute_ns(owned_params, gathered_grads, ns_steps) - gathered_grads.clear() - - # Stages 4-5: launch async scatter. - with record_function("muon::launch_scatter"): - work, recv_buf, scattered_us, recv_counts = _launch_scatter( - params, owned_params, param_to_state, rank, num_ranks, - process_group, computed_us) - computed_us.clear() - - yield # --- YIELD 2: other chunks can launch their scatter --- - - with record_function("muon::wait_scatter"): - work.wait() - _complete_scatter(recv_buf, recv_counts, params, param_to_state, rank, - scattered_us) - del recv_buf - - # Stage 6: apply parameter updates. - with record_function("muon::update_params"): - _update_params(params, param_to_state, rank, scattered_us, lr, - weight_decay) - scattered_us.clear() diff --git a/build/torch210-cxx11-cu128-x86_64-linux/qk_clip.py b/build/torch210-cxx11-cu128-x86_64-linux/qk_clip.py deleted file mode 100644 index 2aba711b3004b7f09e7141da7ef834bd61cc2430..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-x86_64-linux/qk_clip.py +++ /dev/null @@ -1,198 +0,0 @@ -import logging -import math -from dataclasses import dataclass - -import torch -from torch.distributed.tensor import DTensor - -from .core import normalize_fqn - -logger = logging.getLogger(__name__) - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - and return (kind, layer_index). - - Supported kinds: - MHA/GQA: 'wq', 'wk', 'q_proj', 'k_proj' - MLA: 'wq_b' (Q up-proj), 'wkv_b' (KV up-proj) - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.1.attn.wq_b.weight' -> ('wq_b', 1) - 'model.0.attn.wkv_b.weight' -> ('wkv_b', 0) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = normalize_fqn(name).split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: str | None # 'wq'/'q_proj'/'wq_b' or 'wk'/'k_proj'/'wkv_b' or None - indices: list[int] # which heads to consider for clipping - head_dim: int # from config (qk_head_dim for MLA wq_b) - threshold: float # from config - logit: torch.Tensor | None - - # MLA-specific fields - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 - - -def get_qk_clip_info(clip_config, n, qk_logits): - """Extract QK clipping info for a named parameter. - - Args: - clip_config: QK clipping configuration dict (or None). - MHA/GQA keys: head_dim, threshold, q_indices, k_indices - MLA extra keys: is_mla=True, qk_nope_head_dim, qk_rope_head_dim, v_head_dim - n: Parameter name string. - qk_logits: Dict mapping layer indices to logit tensors (or None). - - Returns: - QKClipInfo instance with clipping configuration for this parameter. - """ - if clip_config is None: - return None - - head_dim = clip_config.get('head_dim') - threshold = clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - # In TP settings, qk_logits may be DTensor - # We convert it to full tensor here for simplicity - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - -def compute_scales(p, qk_clip_state): - """Compute per-head scaling factors for QK clipping. - - Returns scales tensor (√γ per head) if any head exceeds threshold, else None. - For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim. - """ - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - # Check if any head exceeds threshold before allocating. - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - - if not head_scales: - return None - - # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows - if qk_clip_state.is_mla and kind == 'wkv_b': - effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim - else: - effective_head_dim = head_dim - - H_global = p.shape[0] // effective_head_dim - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale - return scales_full - - -def qk_clip(p, scales, info): - """Apply per-head scaling to a Q/K projection weight matrix. - - Args: - p: Parameter (nn.Parameter or raw tensor). - scales: [n_heads] tensor, each element = √γ_h. - info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions. - - MLA sub-region scaling per Algorithm 1 (MuonClip): - wq_b: q_nope rows → √γ, q_pe rows → γ - wkv_b: k_nope rows → √γ, v rows → unchanged - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not info.is_mla: - # MHA/GQA: uniform √γ applied to all rows in each head - W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: vectorized sub-region scaling within each head - if info.kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_head_dim = qk_nope + info.qk_rope_head_dim - W_3d = W.view(-1, qk_head_dim, W.shape[1]) # [H, qk_head_dim, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # q_nope → √γ - W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1, - 1)) # q_pe → γ - - elif info.kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - W_3d = W.view(-1, kv_stride, W.shape[1]) # [H, kv_stride, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # k_nope → √γ - # v rows: not touched (k_R shared rotary unchanged) diff --git a/build/torch210-cxx11-cu130-x86_64-linux/_ops.py b/build/torch210-cxx11-cu130-x86_64-linux/_ops.py deleted file mode 100644 index 034bff088659b5df6f6d401feb18c89dc5f33b29..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _optimizer_8d53b78_dirty -ops = torch.ops._optimizer_8d53b78_dirty - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_optimizer_8d53b78_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch210-cxx11-cu130-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so b/build/torch210-cxx11-cu130-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so deleted file mode 100755 index 570449600e33b32b83585ec10fe0593b4c4318bc..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:45eef069a7caa85678cd1e05f0c60c5cfbc676dc93a1bcb31e55eb34730aa469 -size 2004728 diff --git a/build/torch210-cxx11-cu130-x86_64-linux/adamw.py b/build/torch210-cxx11-cu130-x86_64-linux/adamw.py deleted file mode 100644 index b5a95816a9f5b9e1889eaadae65373bfbced809a..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-x86_64-linux/adamw.py +++ /dev/null @@ -1,271 +0,0 @@ -import logging -from collections import defaultdict -from typing import cast - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -def fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: float | torch.Tensor, - weight_decay: float, - eps: float, - maximize: bool, -) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: dict | None = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) - - -def _to_local(t): - """Unwrap DTensor to local tensor for fused ops.""" - return t._local_tensor if isinstance(t, DTensor) else t - - -# --------------------------------------------------------------------------- -# Caches for eliminating per-step Python overhead. -# -# Placement grouping and tensor list assembly are identical every step -# (params don't change placement, moment/step tensors are the same objects -# after initialisation). We cache them keyed by id() of the param list -# stored in param_groups (stable across steps). -# -# Only gradients change each step and must be collected fresh. -# --------------------------------------------------------------------------- - -# id(group["params"]) → dict[placement_key, list[param]] -_placement_cache: dict[int, dict[tuple, list]] = {} - -# id(placement_group_list) → (params_local, moment1, moment2, state_steps) -_tensor_cache: dict[int, tuple[list, list, list, list]] = {} - - -def _step_adamw_params_slow(optimizer_state, params, group): - """Uncached fallback for the rare case where some params lack grads.""" - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - g = p.grad - if g is None: - continue - state = optimizer_state[p] - params_with_grads.append(_to_local(p)) - grads.append(_to_local(g)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - if not params_with_grads: - return - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw_params(optimizer_state, params, group): - """Run fused AdamW on a list of parameters sharing the same placement. - - After the first call, cached tensor lists (params_local, moment1, - moment2, state_steps) are reused — only gradients are collected fresh. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - params: List of parameters to update. - group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay. - """ - # Collect grads — the only thing that changes each step. - with record_function("adamw::collect_grads"): - grads = [] - for p in params: - g = p.grad - if g is None: - # Rare: fall back to slow path that filters per-param. - _step_adamw_params_slow(optimizer_state, params, group) - return - grads.append(_to_local(g)) - - tensor_key = id(params) - if tensor_key not in _tensor_cache: - with record_function("adamw::init_tensor_cache"): - params_local = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - state = optimizer_state[p] - params_local.append(_to_local(p)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(p.grad) - state["moment2"] = torch.zeros_like(p.grad) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - _tensor_cache[tensor_key] = (params_local, moment1, moment2, - state_steps) - - params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key] - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - with record_function("adamw::fused_adamw"): - fused_adamw( - params_local, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw(optimizer_state, group): - """Dispatch AdamW step, grouping parameters by type and placement. - - Placement grouping is cached after the first call since params never - change their placement between steps. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - group: Parameter group dict. - """ - params = group["params"] - placement_key = id(params) - - if placement_key not in _placement_cache: - with record_function("adamw::group_by_placement"): - placement_to_params: dict[tuple, - list[torch.Tensor]] = defaultdict(list) - for p in params: - match p: - case DTensor(): - logger.debug( - "[AdamW] DTensor param: shape=%s, placements=%s, " - "mesh=%s, grad=%s", p.shape, p.placements, - p.device_mesh.mesh_dim_names, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple( - [p.placements, p.device_mesh])].append(p) - case torch.Tensor(): - logger.debug( - "[AdamW] plain param: shape=%s, grad=%s", p.shape, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple([torch.Tensor, - None])].append(p) - - logger.debug("[AdamW] %d placement groups, %d total params", - len(placement_to_params), len(params)) - - _placement_cache[placement_key] = dict(placement_to_params) - - for group_params in _placement_cache[placement_key].values(): - step_adamw_params(optimizer_state, group_params, group) diff --git a/build/torch210-cxx11-cu130-x86_64-linux/async_utils.py b/build/torch210-cxx11-cu130-x86_64-linux/async_utils.py deleted file mode 100644 index a45c530ac9cad88e3555ec1047a6aa59f225347e..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-x86_64-linux/async_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import logging -from typing import Generator - -logger = logging.getLogger(__name__) - - -class _Task: - """Internal: wraps a generator, advances one yield at a time.""" - - def __init__(self, generator: Generator[None, None, None], index: int): - self._generator = generator - self._index = index - self._steps_completed = 0 - self.step() # run to first yield - - def step(self) -> bool: - try: - next(self._generator) - self._steps_completed += 1 - logger.debug("pipeline[%d] completed stage %d", self._index, - self._steps_completed) - return True - except StopIteration: - logger.debug("pipeline[%d] finished after %d stages", self._index, - self._steps_completed) - return False - - def close(self): - self._generator.close() - - -def run_pipeline( - pipelines: Generator[Generator[None, None, None], None, None], - max_concurrent: int, -) -> None: - """Run generator-based pipelines with bounded concurrency. - - Each pipeline is a generator that yields at stage boundaries. - The runtime interleaves pipelines so communication and computation - overlap across chunks. - """ - if max_concurrent <= 0: - raise ValueError(f"max_concurrent must be > 0, got {max_concurrent}") - - have_new = True - task_index = 0 - previous_tasks: list[_Task] = [] - - try: - while have_new or previous_tasks: - running_tasks: list[_Task] = [] - - # Admit one new pipeline per iteration (staggered admission). - # Admitting one at a time ensures that while chunk N does NS - # compute on the default stream, chunk N+1's NCCL all-to-all - # runs concurrently on the NCCL stream — creating real - # communication/computation overlap on the GPU. - if have_new and len(previous_tasks) < max_concurrent: - try: - gen = next(pipelines) - task = _Task(gen, task_index) - task_index += 1 - running_tasks.append(task) - except StopIteration: - have_new = False - - # Advance every previously-yielded task by one step. - for task in previous_tasks: - if task.step(): - running_tasks.append(task) - - previous_tasks = running_tasks - except BaseException: - # Clean up all in-flight generators to release GPU resources. - for task in previous_tasks: - task.close() - raise diff --git a/build/torch210-cxx11-cu130-x86_64-linux/core.py b/build/torch210-cxx11-cu130-x86_64-linux/core.py deleted file mode 100644 index c69d515afef305ad0ed66374095fa2d2468d99cc..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-x86_64-linux/core.py +++ /dev/null @@ -1,219 +0,0 @@ -import logging -import math -from dataclasses import dataclass -from typing import List - -import torch -from torch.distributed import ProcessGroup -from torch.distributed.tensor import DTensor - -# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into -# parameter FQNs. Activation checkpointing similarly inserts -# "_checkpoint_wrapped_module". Strip these so name-based matching (skip_keys, -# expert_keys, QK layer parsing) works regardless of wrapper nesting. -_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) - -logger = logging.getLogger(__name__) - - -def normalize_fqn(name: str) -> str: - """Strip torch.compile / checkpoint wrapper components from a parameter FQN.""" - return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS) - - -@dataclass -class _muon_state: - worker_rank: int - process_group: ProcessGroup - rank_indices: dict[int, tuple] # local_rank -> per-dim indices - rank_numels: dict[int, int] # local_rank -> numel - name: str - qk_clip_state: torch.Tensor | None = None - - -def _batch_momentum( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update (no nesterov).""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - - -def _batch_momentum_nesterov( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update with nesterov correction.""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - nesterov_terms = torch._foreach_mul(momentum_bufs, momentum) - torch._foreach_add_(grads, nesterov_terms) - - -_compiled_momentum: dict[bool, callable] = {} -_use_momentum_compile = True - - -def set_momentum_compile(enabled: bool): - """Toggle torch.compile for batched momentum.""" - global _use_momentum_compile - _use_momentum_compile = enabled - - -def batch_pre_ortho( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, - nesterov: bool, -) -> None: - """Batched momentum update on lists of plain tensors. - - Mirrors dion's ``muon_update_pre_orthogonalize``. - Inputs must be plain CUDA tensors (not DTensor). - Modifies ``momentum_bufs`` and (for nesterov) ``grads`` in-place. - - When compile is enabled, uses separately compiled functions for - nesterov=True/False to avoid graph breaks from the branch. - """ - fn = _batch_momentum_nesterov if nesterov else _batch_momentum - if _use_momentum_compile: - if nesterov not in _compiled_momentum: - _compiled_momentum[nesterov] = torch.compile(fn) - fn = _compiled_momentum[nesterov] - fn(grads, momentum_bufs, momentum) - - -def _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay): - """Weight-decay + update on plain tensors. - - Not compiled: per-param @torch.compile caused ~0.25ms TorchDynamo cache - lookup per call × 256+ params = massive overhead. The pipeline path uses - batched _foreach_* ops instead; this function remains for base() and - distributed_muon(). - """ - p_data.mul_(1 - lr * weight_decay) - p_data.add_(u_data, alpha=-adjusted_lr) - - -def update_p(p, u, lr, adjusted_lr, weight_decay): - """Apply weight decay and orthogonalized update to parameter. - - Args: - p: Parameter (torch.nn.Parameter or DTensor). - u: Orthogonalized update tensor. - lr: Base learning rate. - adjusted_lr: Size-adjusted learning rate. - weight_decay: Weight decay coefficient. - """ - # Unwrap Parameter -> underlying data tensor. - p_data = p.data if isinstance(p, torch.nn.Parameter) else p - # Unwrap DTensor -> local CUDA tensor for compiled kernel. - if isinstance(p_data, DTensor): - p_data = p_data._local_tensor - u_data = u._local_tensor if isinstance(u, DTensor) else u - _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay) - - -def adjust_lr_for_muon(lr, param_shape): - """Scale learning rate based on parameter matrix dimensions. - - Args: - lr: Base learning rate. - param_shape: Shape of the parameter tensor. - - Returns: - Adjusted learning rate. - """ - A, B = param_shape[:2] - # We adjust the learning rate and weight decay based on the size of the parameter matrix - # as described in the paper - adjusted_ratio = 0.2 * math.sqrt(max(A, B)) - adjusted_lr = lr * adjusted_ratio - return adjusted_lr - - -def _match_key(parts, key): - """Check if key matches as contiguous components in parts. - - Single-component keys (e.g. "experts") match any single component. - Multi-component keys (e.g. "experts.w1") match as a contiguous subsequence. - """ - key_parts = key.split(".") - key_len = len(key_parts) - if key_len == 1: - return key in parts - return any(parts[i:i + key_len] == key_parts - for i in range(len(parts) - key_len + 1)) - - -def is_expert_param(name, expert_keys): - """Check if a parameter name matches any expert key (component-level).""" - if not expert_keys: - return False - parts = normalize_fqn(name).split(".") - return any(_match_key(parts, key) for key in expert_keys) - - -def default_is_muon(name, x, expert_keys=None): - normalized = normalize_fqn(name) - parts = normalized.split(".") - skip_keys = [ - "embed_tokens", - "lm_head", - "tok_embeddings", - "output", - "mhc_attn", - "mhc_ffn", - "lambda_proj", - ] - if any(key in parts for key in skip_keys): - logger.info( - "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d", - normalized, name, x.ndim) - return False - effective_ndim = x.ndim - is_expert = is_expert_param(name, expert_keys) - if is_expert: - effective_ndim -= 1 - result = effective_ndim >= 2 - logger.info( - "[is_muon] %s (orig: %s): ndim=%d, expert=%s, effective_ndim=%d → %s", - normalized, name, x.ndim, is_expert, effective_ndim, - "Muon" if result else "AdamW") - return result - - -def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None): - if is_muon_func is None: - is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys) - - muon_params, muon_names = [], [] - non_muon_params, non_muon_names = [], [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - non_muon_names.append(n) - - logger.info("[param_groups] expert_keys=%s, Muon=%d, AdamW=%d", - expert_keys, len(muon_names), len(non_muon_names)) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] diff --git a/build/torch210-cxx11-cu130-x86_64-linux/cpu_offload.py b/build/torch210-cxx11-cu130-x86_64-linux/cpu_offload.py deleted file mode 100644 index fb5e69154a1d4a6c884491413a37a9acf0f66c80..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-x86_64-linux/cpu_offload.py +++ /dev/null @@ -1,206 +0,0 @@ -"""CPU offloading for optimizer states. - -Manages a pinned CPU memory pool and async CUDA streams to offload -optimizer state tensors (momentum buffers, Adam moments) to CPU between -optimizer steps, freeing GPU memory. - -All tracked tensors are packed into a single flat pinned CPU buffer -(per dtype). D2H and H2D copies are performed per-tensor directly -between individual GPU tensors and their slice of the CPU flat buffer -— no GPU staging buffer is allocated, so there is **no temporary GPU -memory spike** during offload or reload. - -Individual tensor storages are freed after offload via -``untyped_storage().resize_(0)``, preserving tensor identity so -downstream caches remain valid. -""" - -import logging -from collections import defaultdict - -import torch -from torch.distributed.tensor import DTensor - -logger = logging.getLogger(__name__) - - -class CPUOffloadPool: - """Pinned CPU memory pool for async optimizer state offloading. - - Tracked tensors are grouped by dtype. Each group gets a single flat - pinned CPU buffer. D2H / H2D copies are per-tensor (into slices of - the flat buffer) to avoid allocating a GPU staging buffer. - """ - - def __init__(self): - self._managed: list[torch.Tensor] = [] - self._storage_nbytes: dict[int, int] = {} # id(t) → bytes - - # Per-dtype group: populated on first offload. - # dtype → dict with keys: - # "indices" : list[int] managed-list indices - # "offsets" : list[tuple[int,int]] (start, numel) in flat buf - # "total" : int total numel - # "cpu_flat" : Tensor pinned CPU buffer - self._groups: dict[torch.dtype, dict] = {} - - self._offload_stream: torch.cuda.Stream | None = None - self._device: torch.device | None = None - self._initialized: bool = False - self._logged: bool = False - - # ------------------------------------------------------------------ - @staticmethod - def _local(t: torch.Tensor) -> torch.Tensor: - """Unwrap DTensor to its local CUDA tensor.""" - return t._local_tensor if isinstance(t, DTensor) else t - - def _ensure_stream(self): - if self._offload_stream is None: - self._offload_stream = torch.cuda.Stream(device=self._device) - - # ------------------------------------------------------------------ - def track(self, tensor: torch.Tensor): - """Register a GPU tensor for CPU offloading. Idempotent.""" - tid = id(tensor) - if tid in self._storage_nbytes: - return - local = self._local(tensor) - if self._device is None: - self._device = local.device - storage = local.untyped_storage() - # Skip tensors with empty storage (e.g. empty FSDP shards) - if storage.size() == 0: - return - self._storage_nbytes[tid] = storage.size() - self._managed.append(tensor) - - # ------------------------------------------------------------------ - def _init_buffers(self): - """Build per-dtype flat buffers on first offload.""" - # Group managed tensors by dtype. - dtype_map: dict[torch.dtype, list[tuple[int, int]]] = defaultdict(list) - for idx, t in enumerate(self._managed): - local = self._local(t) - dtype_map[local.dtype].append((idx, local.numel())) - - total_cpu_bytes = 0 - for dtype, entries in dtype_map.items(): - offsets: list[tuple[int, int]] = [] - indices: list[int] = [] - off = 0 - for idx, n in entries: - indices.append(idx) - offsets.append((off, n)) - off += n - cpu_flat = torch.empty(off, dtype=dtype, device="cpu", pin_memory=True) - self._groups[dtype] = { - "indices": indices, - "offsets": offsets, - "total": off, - "cpu_flat": cpu_flat, - } - total_cpu_bytes += off * cpu_flat.element_size() - - self._initialized = True - logger.info( - "[CPUOffload] Pool initialized: %d tensors, %d dtype group(s), " - "%.2f MB pinned CPU memory", - len(self._managed), - len(self._groups), - total_cpu_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def offload(self): - """Per-tensor async D2H into CPU flat buffer, then free GPU storage.""" - if not self._managed: - return - if not self._initialized: - self._init_buffers() - self._ensure_stream() - - # Offload stream waits for compute to finish. - compute_event = torch.cuda.current_stream(self._device).record_event() - self._offload_stream.wait_event(compute_event) - - offloaded_bytes = 0 - - # Per-tensor D2H copies directly into CPU flat buffer slices. - # No GPU staging buffer → no temporary GPU memory spike. - with torch.cuda.stream(self._offload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - cpu_flat[off : off + n].copy_(local.reshape(-1), non_blocking=True) - - offloaded_bytes += grp["total"] * cpu_flat.element_size() - - # Wait for all D2H copies to land, then free GPU storage. - self._offload_stream.synchronize() - for t in self._managed: - storage = self._local(t).untyped_storage() - if storage.size() != 0: - storage.resize_(0) - else: - raise RuntimeError( - f"Tensor storage is already freed (size=0) before offload. " - f"This indicates a double-free or external interference. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - - if not self._logged: - logger.info( - "[CPUOffload] Offloaded %.2f MB (GPU → CPU)", - offloaded_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def reload(self): - """Per-tensor H2D from CPU flat buffer on the default stream. - - Runs on the current (default) CUDA stream to avoid stream - interaction issues with the parallel Muon pipeline. Since - pinned CPU memory is the source, the copies overlap with - GPU idle time between steps. - """ - if not self._managed or not self._initialized: - return - - reloaded_bytes = 0 - - # Re-allocate all GPU storages first. - for t in self._managed: - local = self._local(t) - storage = local.untyped_storage() - if storage.size() != 0: - raise RuntimeError( - f"Storage should have been freed (size=0) before reload, " - f"but got size={storage.size()}. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - storage.resize_(self._storage_nbytes[id(t)]) - - # Per-tensor H2D copies from CPU flat buffer slices. - # non_blocking=True with pinned source allows DMA overlap. - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off : off + n], non_blocking=True) - - reloaded_bytes += grp["total"] * cpu_flat.element_size() - - if not self._logged: - logger.info( - "[CPUOffload] Reloaded %.2f MB (CPU → GPU)", reloaded_bytes / (1024**2) - ) diff --git a/build/torch210-cxx11-cu130-x86_64-linux/distributed/utils.py b/build/torch210-cxx11-cu130-x86_64-linux/distributed/utils.py deleted file mode 100644 index 683096bcd50489e7a09959753fdbc29ba0f3b6a6..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-x86_64-linux/distributed/utils.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import (Placement, Shard, - _StridedShard) - - -def _is_shard(placement: Placement) -> bool: - """Check if a placement is a shard type (Shard or _StridedShard). - - In PyTorch 2.10+, _StridedShard no longer inherits from Shard, so - ``placement.is_shard()`` returns False for _StridedShard. This helper - handles both old and new hierarchies. - """ - return isinstance(placement, (Shard, _StridedShard)) - - -def get_slices_of_dtensor( - target: DTensor | torch.Tensor, - local_rank: int, - shard_mesh: DeviceMesh, - shard_placements: tuple[Placement], -) -> tuple[slice | torch.Tensor, ...]: - """ - Get per-dimension indices for a given rank's shard of the target tensor. - - Uses ``Shard.local_shard_size_and_offset`` and - ``_StridedShard.local_shard_size_and_offset`` for correct handling of - both contiguous and strided (non-contiguous) sharding. - - Args: - target (DTensor | torch.Tensor): The target tensor (for its shape). - local_rank (int): The local rank within the shard group. - shard_mesh (DeviceMesh): The shard mesh (only shard dimensions). - shard_placements (tuple[Placement]): The shard placements. - - Returns: - A tuple of indices (one per tensor dim). Each element is either: - - A ``slice`` (for contiguous or unsharded dims) - - A 1-D ``torch.LongTensor`` of indices (for strided sharding) - """ - - # find the global rank of the local rank in the shard mesh - rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank] - - rank_coords = (shard_mesh.mesh == rank).nonzero() - - assert len(rank_coords) == 1 - rank_coords = tuple(rank_coords[0].tolist()) - - assert len(rank_coords) == len(shard_placements) - - # Track per-shard-dim indices. - # None means "not yet sharded on this dim". - dim_indices: dict[int, torch.Tensor] = {} - - # Caution: Assuming replicate-to-shard of the shard mesh goes with - # left-to-right sharding. This is ensured by the sorting logic of - # construct_shard_mesh function. - for mesh_dim_idx, (rank_coord, placement) in enumerate( - zip(rank_coords, shard_placements)): - assert _is_shard(placement) - - num_chunks = shard_mesh.mesh.shape[mesh_dim_idx] - shard_dim = placement.dim - - # Current effective size on this dim (may already be sub-sharded) - if shard_dim in dim_indices: - curr_size = len(dim_indices[shard_dim]) - else: - curr_size = target.size()[shard_dim] - - # Compute indices for this level of sharding - if isinstance(placement, _StridedShard): - _shard_size, offsets = _StridedShard.local_shard_size_and_offset( - placement, - curr_size, - num_chunks, - rank_coord, - return_first_offset=False) - new_indices = torch.tensor(offsets, dtype=torch.long) - else: - shard_size, offset = Shard.local_shard_size_and_offset( - curr_size, num_chunks, rank_coord) - new_indices = torch.arange(offset, - offset + shard_size, - dtype=torch.long) - - # Compose with previous indices on this dim - if shard_dim in dim_indices: - dim_indices[shard_dim] = dim_indices[shard_dim][new_indices] - else: - dim_indices[shard_dim] = new_indices - - # Build result tuple - result: list[slice | torch.Tensor] = [] - for d in range(len(target.size())): - if d not in dim_indices: - result.append(slice(None)) - else: - indices = dim_indices[d] - # Convert contiguous indices to slice for efficiency - if len(indices) > 0: - start = indices[0].item() - expected = torch.arange(start, - start + len(indices), - dtype=torch.long) - if torch.equal(indices, expected): - result.append(slice(start, start + len(indices))) - else: - result.append(indices) - else: - result.append(slice(0, 0)) - - return tuple(result) - - -_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, - ProcessGroup]] = dict() - - -def construct_shard_mesh( - placements: tuple[Placement], - mesh: DeviceMesh, -) -> tuple[DeviceMesh, ProcessGroup, tuple[Placement, ...]]: - """Construct shard sub-mesh and ProcessGroup for all-to-all communication. - - Given a DTensor's placements and device mesh, extracts the "shard group" - — the set of ranks that together hold all shards of the same replica — - and creates a ProcessGroup for all-to-all among them. - - Steps: - 1. Sort placements: Replicate first, then Shard by (dim, granularity). - 2. Permute the mesh tensor to match the sorted order. - 3. Collapse Replicate dims → list of shard sub-meshes (one per replica). - 4. Create/retrieve a cached ProcessGroup for the current rank's sub-mesh. - - Example — 8 GPUs, mesh shape (2, 2, 2), - placements ``[Shard(0), Replicate, _StridedShard(0)]``:: - - Step 1 — Sort: [Replicate, _StridedShard(0), Shard(0)] - Permutation: [1, 2, 0] - - Step 2 — Permute mesh dims by [1, 2, 0]: - Original: Permuted: - [[[0,1],[2,3]], [[[0,2],[1,3]], - [[4,5],[6,7]]] [[4,6],[5,7]]] - - Step 3 — Unbind replicate dim (dim 0), giving 2 shard sub-meshes: - sub-mesh 0 = [[0,2],[1,3]] (replica group 0) - sub-mesh 1 = [[4,6],[5,7]] (replica group 1) - shard_placements = (_StridedShard(0), Shard(0)) - - Step 4 — Rank 0 → ProcessGroup([0,1,4,5]) - Rank 2 → ProcessGroup([2,3,6,7]) - - Returns: - ``(shard_mesh, process_group, shard_placements)`` - """ - my_rank = dist.get_rank() - assert mesh.mesh.device.type == 'cpu' - - # -- Fast path: 1D all-shard mesh → reuse existing PG. ---------------- - # Reuses the mesh's existing ProcessGroup directly, avoiding the - # overhead of dist.new_group(). The standard path below also handles - # subset calls safely via use_local_synchronization=True, but this - # fast path is still beneficial for the common 1D shard case. - if mesh.ndim == 1 and len(placements) == 1 and _is_shard(placements[0]): - key = (*mesh.mesh.shape, *mesh.mesh.flatten().tolist()) - if key not in _ranks_to_dist_cache: - _ranks_to_dist_cache[key] = (mesh, mesh.get_group()) - return (*_ranks_to_dist_cache[key], tuple(placements)) - - mesh_tensor = mesh.mesh.clone() - - # -- Step 1: Sort placements (Replicate first, then Shard by dim). ------ - # _StridedShard comes BEFORE regular Shard on the same dim so that - # get_slices_of_dtensor applies the outer sharding first, matching - # DTensor's left-to-right (outer-to-inner) composition order. - def _sort_key(item): - index, placement = item - assert not placement.is_partial(), "Partial placement not supported" - if placement.is_replicate(): - return (-1, 0, index) - assert _is_shard(placement), f"Unsupported: {type(placement)}" - split = (-1 / placement.split_factor if isinstance( - placement, _StridedShard) else 0) - return (placement.dim, split, index) - - indexed = sorted(enumerate(placements), key=_sort_key) - perm, sorted_placements = zip(*indexed) - - # -- Step 2: Permute mesh to match sorted placement order. -------------- - sorted_mesh = mesh_tensor.permute(perm) - - # -- Step 3: Collapse replicate dims → list of shard sub-meshes. -------- - # E.g. mesh (2, 3, 4, 4) with [R, R, S(0), S(1)] → 6 sub-meshes of (4, 4) - num_rep = sum(1 for p in sorted_placements if p.is_replicate()) - if num_rep > 0: - if num_rep > 1: - sorted_mesh = sorted_mesh.flatten(0, num_rep - 1) - shard_meshes = list(torch.unbind(sorted_mesh, dim=0)) - else: - shard_meshes = [sorted_mesh] - shard_placements = sorted_placements[num_rep:] - assert len(shard_placements) == len(set(shard_placements)) - - # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. -- - # Each rank only creates the group it belongs to, using - # use_local_synchronization=True so that only group members need to - # coordinate. This avoids deadlocks when different PP stages call - # construct_shard_mesh for different parameters. - def _cache_key(t: torch.Tensor) -> tuple: - return (*t.shape, *t.flatten().tolist()) - - my_key = None - for sm in shard_meshes: - if (my_rank == sm).any().item(): - key = _cache_key(sm) - assert my_key is None, "Rank appears in multiple shard groups" - my_key = key - if key not in _ranks_to_dist_cache: - pg = dist.new_group(sm.flatten().tolist(), - use_local_synchronization=True) - _ranks_to_dist_cache[key] = ( - DeviceMesh(device_type="cuda", mesh=sm), - pg, - ) - - return (*_ranks_to_dist_cache[my_key], shard_placements) diff --git a/build/torch210-cxx11-cu130-x86_64-linux/matmul_transpose_triton.py b/build/torch210-cxx11-cu130-x86_64-linux/matmul_transpose_triton.py deleted file mode 100644 index 792de23d82c3fb45fe33d397ab9b76a0787259d0..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-x86_64-linux/matmul_transpose_triton.py +++ /dev/null @@ -1,122 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], - restore_value=['y'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -@torch.library.custom_op("muon::matmul_transpose_assign", - mutates_args=("d_out", )) -def matmul_transpose_assign(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """Compute d_out = d_in @ d_in.T using an optimized Triton kernel.""" - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -@matmul_transpose_assign.register_fake -def _(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """FakeTensor impl: d_out is already allocated, mutation is declared.""" - pass diff --git a/build/torch210-cxx11-cu130-x86_64-linux/metadata.json b/build/torch210-cxx11-cu130-x86_64-linux/metadata.json deleted file mode 100644 index c55a35717622f1dd5c8ba376ea3a814cbcc10d78..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-x86_64-linux/metadata.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch210-cxx11-cu130-x86_64-linux/muon.py b/build/torch210-cxx11-cu130-x86_64-linux/muon.py deleted file mode 100644 index 14c0e22471fa6d47a51ed95e0e0c341dc18d5194..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-x86_64-linux/muon.py +++ /dev/null @@ -1,1068 +0,0 @@ -import logging -import types -from collections import defaultdict -from typing import Any - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Replicate, Shard -from torch.profiler import record_function - -from .adamw import _placement_cache, _tensor_cache, step_adamw -from .async_utils import run_pipeline -from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho, - get_default_muon_param_groups, is_expert_param, update_p) -from .cpu_offload import CPUOffloadPool -from .distributed.utils import (_is_shard, construct_shard_mesh, - get_slices_of_dtensor) -from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO, - _zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5_batched) -from .pipeline import muon_chunk_pipeline, prelaunch_first_gather -from .qk_clip import compute_scales, get_qk_clip_info, qk_clip - -logger = logging.getLogger(__name__) - - -def _expand_expert_params(names, params, expert_keys): - """Expand expert params by splitting on dim 0 (expert dimension). - - Params whose name matches any key in ``expert_keys`` are treated as - expert-parallel tensors. Their outermost dimension is the expert - dimension: an ``(E, out, in)`` tensor becomes ``E`` separate 2D - ``nn.Parameter`` views so that in-place updates propagate back to - the original storage. - - Non-expert params with ``ndim > 2`` trigger an ``AssertionError`` — - if they are expert params, their key must be added to ``expert_keys``. - - The grad must already be set on each expert param (e.g. after momentum). - - For DTensor expert params, placements that shard on dim 0 (expert dim) - are consumed by the split. Non-dim-0 shard placements (e.g. TP) are - preserved: each 2D slice is wrapped as a DTensor on the corresponding - submesh so the parallel pipeline handles the TP communication. - """ - expanded_names = [] - expanded_params = [] - - for n, p in zip(names, params): - is_expert = is_expert_param(n, expert_keys) - is_dtensor = isinstance(p.data, DTensor) - - if is_expert: - if is_dtensor: - logger.debug( - "[expand_expert] %s: expert DTensor, shape=%s, " - "placements=%s, mesh=%s, local_shape=%s", n, p.shape, - p.placements, p.device_mesh.mesh_dim_names, - p.to_local().shape) - else: - logger.debug( - "[expand_expert] %s: expert plain tensor, shape=%s", n, - p.data.shape) - - if not is_expert: - assert p.data.ndim <= 2, ( - f"Param {n} has ndim={p.data.ndim} but does not match " - f"expert_keys={expert_keys}. If this is an expert param, " - f"add its key to expert_keys.") - expanded_names.append(n) - expanded_params.append(p) - continue - - g = p.grad - assert g is not None, ( - f"Expert param {n} must have grad set before expansion") - - tp_mesh = None - tp_placements_2d = None - - if is_dtensor: - local_data = p.to_local() - local_grad = g.to_local() if isinstance(g, DTensor) else g - - # Find non-dim-0 shard placements (e.g. TP sharding). - # After splitting on dim 0, Shard(k) becomes Shard(k-1). - tp_dim_indices = [] - tp_placements_2d = [] - for i, pl in enumerate(p.placements): - if _is_shard(pl) and pl.dim != 0: - tp_dim_indices.append(i) - tp_placements_2d.append(Shard(pl.dim - 1)) - - if tp_dim_indices: - tp_dim_names = tuple(p.device_mesh.mesh_dim_names[i] - for i in tp_dim_indices) - if len(tp_dim_names) == 1: - tp_mesh = p.device_mesh[tp_dim_names[0]] - else: - tp_mesh = p.device_mesh[tp_dim_names] - else: - local_data = p.data - local_grad = g - - # Expand: split dim 0, reshape each slice to 2D. - num_local_experts = local_data.shape[0] - for i in range(num_local_experts): - slice_data = local_data[i] - slice_grad = local_grad[i] - - if tp_mesh is not None: - # Wrap as DTensor on TP submesh so the pipeline handles - # TP communication (gather/scatter across TP ranks). - dt_data = DTensor.from_local(slice_data, - device_mesh=tp_mesh, - placements=tp_placements_2d) - dt_grad = DTensor.from_local(slice_grad, - device_mesh=tp_mesh, - placements=tp_placements_2d) - expert_param = torch.nn.Parameter(dt_data, requires_grad=False) - expert_param.grad = dt_grad - else: - expert_param = torch.nn.Parameter(slice_data, - requires_grad=False) - expert_param.grad = slice_grad - - expanded_names.append(f"{n}[{i}]") - expanded_params.append(expert_param) - - p.grad = None # allow expert grad storage to be freed after pipeline - - return expanded_names, expanded_params - - -class Muon(torch.optim.Optimizer): - """ - Muon - MomentUm Orthogonalized by Newton-schulz - - Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- - processing step, in which each 2D parameter's update is replaced with the nearest orthogonal - matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has - the advantage that it can be stably run in bfloat16 on the GPU. - - Some warnings: - - We believe this optimizer is unlikely to work well for training with small batch size. - - We believe it may not work well for finetuning pretrained models, but we haven't tested this. - - Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. - lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) - momentum: The momentum used by the internal SGD. (0.95 is a good default) - nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) - ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. - Parameters that are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW instead. - adamw_lr: The learning rate for the internal AdamW. - adamw_betas: The betas for the internal AdamW. - adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - warmup_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher warmup_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. - chunk_size : Batch size of parameters to process in each - all2all gather/compute/scatter step. - Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified. - use_distributed_muon: Use distributed muon by Liu et al. (2024). - For testing purpose only. - expert_keys: List of strings to identify expert-parallel parameters. - If any key appears in a parameter's name, its outermost - dimension is treated as the expert dimension and expanded - into per-expert 2D params for Muon. For example, - ``expert_keys=["experts"]`` matches any param whose name - contains "experts". 3D+ params not matched by any key - will raise an error. - """ - - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config=None, - warmup_step=5, - chunk_size=-1, - use_distributed_muon=False, - expert_keys=None): - defaults = dict( - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - nesterov=nesterov, - ns_steps=ns_steps, - adamw_betas=adamw_betas, - adamw_eps=adamw_eps, - none_grad=none_grad, - use_muon=True, - ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) - super().__init__(params, defaults) - - self.debug = debug - self.clip_config = clip_config if clip_config is not None else { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100, - } - self.warmup_step = warmup_step - self.chunk_size = chunk_size - self.use_distributed_muon = use_distributed_muon - self.expert_keys = expert_keys - self.cpu_offload = False - self._cpu_offload_pool: CPUOffloadPool | None = None - self._offload_initialized = False - self._parallel_cache: dict[tuple[str, ...], dict] = {} - self._expert_expand_cache: dict[tuple[int, ...], dict] = {} - - def _calc_flops(self, G, steps): - assert len(G.shape) == 2 - M, N = G.shape - if M > N: - M, N = N, M - - return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) - - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - p.placements, p.device_mesh) - - return shard_mesh, shard_pg, shard_placements - - def init_state_and_assign_params(self, names, params, group, qk_logits): - param_to_state = {} - param_to_flops = {} - - total_flops = 0 - for p in params: - g = p.grad - if g is None: - continue - assert g.ndim == 2, "Muon only supports 2D parameters." - - flops = self._calc_flops(g, group["ns_steps"]) - param_to_flops[id(p)] = flops - total_flops += flops - - if self.debug: - logger.debug("Total TFLOPs for Muon: %.2f TFLOPs", - total_flops / 1e12) - - paired = list(zip(names, params)) - - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) - - round_robin = 0 - mesh = ordered_params[0].device_mesh - placements = ordered_params[0].placements - - shard_mesh, shard_pg, shard_placements = self.get_shard_mesh( - ordered_params[0]) - shard_mesh_flattened = shard_mesh.mesh.flatten() - num_ranks = dist.get_world_size(group=shard_pg) - - for n, p in zip(ordered_names, ordered_params): - if mesh != p.device_mesh: - raise ValueError("All parameters must be on the same mesh.") - if placements != p.placements: - raise ValueError("All parameters must have same placements.") - - worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks - round_robin = (round_robin + 1) % len(shard_mesh_flattened) - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - # Precompute per-rank indices and numels for all-to-all. - rank_indices: dict[int, tuple] = {} - rank_numels: dict[int, int] = {} - for r in range(num_ranks): - indices = get_slices_of_dtensor(p, r, shard_mesh, - shard_placements) - rank_indices[r] = indices - numel = 1 - for idx, dim_size in zip(indices, p.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) - rank_numels[r] = numel - - param_to_state[id(p)] = _muon_state( - worker_rank=worker_rank, - process_group=shard_pg, - rank_indices=rank_indices, - rank_numels=rank_numels, - name=n, - qk_clip_state=qk_clip_state, - ) - - return param_to_state, ordered_params - - def base(self, names, params, group, lr, weight_decay, qk_logits): - # Momentum is already applied by _step_muon before this method. - for n, p in zip(names, params): - g = p.grad - if g is None: - continue - - u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) - - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - def distributed_muon( - self, - names: list[str], - params: list[torch.nn.Parameter], - group: dict[str, Any], - lr: float, - weight_decay: float, - qk_logits: list[torch.Tensor | DTensor] | None, - ): - """Batched Distributed Muon — for testing/correctness verification only. - - Uses all-gather to reconstruct full tensors, computes Newton-Schulz on - the full grad, then slices back to local shards. This is simpler but - slower than the parallel pipeline (all2all) path, so it serves as a - reference implementation for verifying correctness. - """ - with record_function("distributed_muon"): - # Momentum is already applied by _step_muon before this method. - ns_steps = group["ns_steps"] - - # Separate plain tensors (no communication) from DTensors. - plain_names, plain_params = [], [] - dtensor_names, dtensor_params = [], [] - for n, p in zip(names, params): - if p.grad is None: - continue - if isinstance(p.data, DTensor): - dtensor_names.append(n) - dtensor_params.append(p) - else: - plain_names.append(n) - plain_params.append(p) - - # Process plain tensors per-param (no communication). - for n, p in zip(plain_names, plain_params): - u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE), - steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, - qk_logits) - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - if not dtensor_params: - return - - # Group DTensors by (placements, mesh) for batched all-gather. - placement_groups: dict[tuple, - tuple[list, - list]] = defaultdict(lambda: ([], [])) - for n, p in zip(dtensor_names, dtensor_params): - key = (p.placements, p.device_mesh) - placement_groups[key][0].append(n) - placement_groups[key][1].append(p) - - logger.info( - "distributed_muon: %d placement groups, %d total dtensors", - len(placement_groups), len(dtensor_params)) - - for (placements, mesh), (grp_names, - grp_params) in placement_groups.items(): - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - placements, mesh) - rank = dist.get_rank(shard_pg) - world_size = dist.get_world_size(shard_pg) - - logger.info(" group: %d params, placements=%s, world_size=%d", - len(grp_params), placements, world_size) - - # Separate params that can be batched (all shard dims evenly - # divisible) from those needing per-param full_tensor - # (e.g. MoE gate weights with fewer rows than shard ranks). - # all_gather_into_tensor requires equal buffer sizes across - # ranks, so uneven splits must use DTensor full_tensor(). - batch_names, batch_params = [], [] - single_names, single_params = [], [] - for n, p in zip(grp_names, grp_params): - even = all(p.shape[pl.dim] % - shard_mesh.mesh.shape[dim_idx] == 0 - for dim_idx, pl in enumerate(shard_placements)) - if even: - batch_names.append(n) - batch_params.append(p) - else: - single_names.append(n) - single_params.append(p) - - # Process uneven-split params per-param via full_tensor(). - for n, p in zip(single_names, single_params): - with record_function("distributed_muon::newton_schulz"): - g_full = p.grad.full_tensor().to(COMM_DTYPE) - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - if not batch_params: - continue - - logger.info(" batched=%d, single=%d", len(batch_params), - len(single_params)) - - # Concat all local grad shards into a single flat buffer. - with record_function("distributed_muon::gather"): - grad_locals = [ - p.grad.to_local().to(COMM_DTYPE).flatten() - for p in batch_params - ] - numels = [g.numel() for g in grad_locals] - grad_concat = torch.cat(grad_locals) - del grad_locals - - # Single all-gather (replaces N separate full_tensor). - grad_gathered = torch.empty( - grad_concat.numel() * world_size, - dtype=COMM_DTYPE, - device="cuda", - ) - dist.all_gather_into_tensor(grad_gathered, - grad_concat, - group=shard_pg) - - total_numel = grad_concat.numel() - del grad_concat - - # Precompute per-param offsets within the concat buffer. - offsets = [] - off = 0 - for ne in numels: - offsets.append(off) - off += ne - - # Per-param: reconstruct full grad → NS → local update. - for i, (n, p) in enumerate(zip(batch_names, batch_params)): - with record_function("distributed_muon::newton_schulz"): - g_full = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - for r in range(world_size): - r_start = r * total_numel + offsets[i] - shard = grad_gathered[r_start:r_start + numels[i]] - indices = get_slices_of_dtensor( - p, r, shard_mesh, shard_placements) - g_full[indices] = shard.reshape( - g_full[indices].shape) - - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - def _setup_parallel(self, names, params, group, qk_logits): - """Compute (or retrieve cached) parallel pipeline metadata. - - Returns: - (ordered_params, param_to_state, rank, chunk_size) - """ - cache_key = tuple(names) - - if cache_key not in self._parallel_cache: - # First call: compute metadata and populate cache. - param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - shard_pg = param_to_state[id(ordered_params[0])].process_group - rank = dist.get_rank(group=shard_pg) - - if self.chunk_size == -1: - shard_ranks = dist.get_world_size(shard_pg) - chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO - elif self.chunk_size > 0: - chunk_size = self.chunk_size - else: - raise ValueError( - "chunk_size must be -1 or a positive integer.") - - ordered_names = [ - param_to_state[id(p)].name for p in ordered_params - ] - name_to_state = { - param_to_state[id(p)].name: param_to_state[id(p)] - for p in ordered_params - } - self._parallel_cache[cache_key] = { - 'ordered_names': ordered_names, - 'name_to_state': name_to_state, - 'rank': rank, - 'chunk_size': chunk_size, - } - else: - # Cached path: rebuild param_to_state with current id(p) keys. - cache = self._parallel_cache[cache_key] - rank = cache['rank'] - chunk_size = cache['chunk_size'] - - name_to_param = dict(zip(names, params)) - ordered_params = [name_to_param[n] for n in cache['ordered_names']] - - param_to_state = {} - for p, n in zip(ordered_params, cache['ordered_names']): - cached_state = cache['name_to_state'][n] - param_to_state[id(p)] = _muon_state( - worker_rank=cached_state.worker_rank, - process_group=cached_state.process_group, - rank_indices=cached_state.rank_indices, - rank_numels=cached_state.rank_numels, - name=n, - qk_clip_state=get_qk_clip_info(self.clip_config, n, - qk_logits), - ) - - return ordered_params, param_to_state, rank, chunk_size - - def parallel(self, - names, - params, - group, - lr, - weight_decay, - qk_logits, - prelaunch_gather=None): - """ - Perform a parallel optimization step using Muon. - - Parameters are chunked and each chunk is processed by a - :func:`muon_chunk_pipeline` generator. :func:`run_pipeline` - interleaves multiple chunks so that communication and computation - overlap across chunks (the same overlap previously achieved by the - warmup + main-loop index scheduling). - - If ``prelaunch_gather`` is provided, it is passed to the first - chunk's generator to skip re-launching the already in-flight - A2A gather. - """ - - # Momentum is already applied by _step_muon before this method. - - ordered_params, param_to_state, rank, chunk_size = ( - self._setup_parallel(names, params, group, qk_logits)) - - def pipelines(): - first = True - for start in range(0, len(ordered_params), chunk_size): - chunk = ordered_params[start:start + chunk_size] - if chunk: - kwargs = dict( - params=chunk, - param_to_state=param_to_state, - rank=rank, - ns_steps=group["ns_steps"], - lr=lr, - weight_decay=weight_decay, - none_grad=group["none_grad"], - ) - if first and prelaunch_gather is not None: - kwargs['prelaunch_gather'] = prelaunch_gather - first = False - yield muon_chunk_pipeline(**kwargs) - - with record_function("muon::pipeline"): - run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1) - - def _step_muon(self, group, qk_logits=None): - params = group["params"] - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - # Apply momentum to all params before routing/expansion. - # Batched using _foreach_* ops (compiled, fullgraph=True). - with record_function("muon::momentum"): - active_params = [p for p in params if p.grad is not None] - if active_params: - # Ensure momentum buffers exist (avoid zeros_like when already present). - for p in active_params: - if "momentum_buffer" not in self.state[p]: - self.state[p]["momentum_buffer"] = torch.zeros_like( - p.grad) - - # Extract local tensors for compiled batch function. - local_grads = [ - p.grad._local_tensor - if isinstance(p.grad, DTensor) else p.grad - for p in active_params - ] - local_bufs = [ - self.state[p]["momentum_buffer"]._local_tensor - if isinstance(self.state[p]["momentum_buffer"], DTensor) - else self.state[p]["momentum_buffer"] - for p in active_params - ] - - # Wrap momentum as tensor for torch.compile. - batch_pre_ortho(local_grads, local_bufs, - torch.tensor(momentum), group["nesterov"]) - - # For non-nesterov, the result is the momentum buffer. - if not group["nesterov"]: - for p in active_params: - p.grad = self.state[p]["momentum_buffer"] - - # Identify batched experts for deferred NS. - # Detection is cheap (condition checks only); actual NS compute is - # deferred so it can overlap with the first chunk's A2A gather. - deferred_expert_work = [] - if self.expert_keys: - batched_expert_indices = [] - for i, (n, p) in enumerate(zip(names, params)): - if not (is_expert_param(n, self.expert_keys) - and p.grad is not None): - continue - # Eligible: plain tensor, or DTensor with no non-dim-0 shards. - if isinstance(p.data, DTensor): - has_tp = any( - _is_shard(pl) and pl.dim != 0 for pl in p.placements) - if has_tp: - continue - batched_expert_indices.append(i) - - if batched_expert_indices: - # Save refs for deferred NS; free grads from param list. - for i in batched_expert_indices: - p = params[i] - g = p.grad - local_g = (g._local_tensor - if isinstance(g, DTensor) else g) - local_data = (p.data._local_tensor if isinstance( - p.data, DTensor) else p.data) - deferred_expert_work.append((local_data, local_g)) - p.grad = None - - # Remove batched experts from lists before expansion. - keep = sorted( - set(range(len(params))) - set(batched_expert_indices)) - names = [names[i] for i in keep] - params = [params[i] for i in keep] - - def _run_deferred_expert_ns(): - """Execute deferred batched expert NS.""" - if not deferred_expert_work: - return - with record_function("muon::batched_expert_ns"): - ns_steps = group["ns_steps"] - for local_data, local_g in deferred_expert_work: - u = zeropower_via_newtonschulz5_batched( - local_g.to(COMM_DTYPE), steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:]) - local_data.mul_(1 - lr * weight_decay) - local_data.add_(u, alpha=-adjusted_lr) - - # Expand expert params by splitting on dim 0. - logger.debug("[_step_muon] before expand: %d params, expert_keys=%s", - len(params), self.expert_keys) - if self.expert_keys: - cache_key = tuple(id(p) for p in params) - cache = self._expert_expand_cache.get(cache_key) - - if cache is None: - # Cold path: full expansion + build cache metadata. - exp_names, exp_params = _expand_expert_params( - names, params, self.expert_keys) - - # Build per-expert-group info for hot-path grad updates. - grad_info = [] - exp_idx = 0 - for orig_idx, (n, p) in enumerate(zip(names, params)): - if not is_expert_param(n, self.expert_keys): - exp_idx += 1 - continue - - is_dt = isinstance(p.data, DTensor) - num_experts = (p.to_local() if is_dt else p.data).shape[0] - - # Detect TP mesh from the first expanded expert param. - tp_mesh = None - tp_pls = None - sample = exp_params[exp_idx] - if isinstance(sample.data, DTensor): - tp_mesh = sample.data.device_mesh - tp_pls = list(sample.data.placements) - - grad_info.append((orig_idx, num_experts, exp_idx, is_dt, - tp_mesh, tp_pls)) - exp_idx += num_experts - - self._expert_expand_cache[cache_key] = { - 'names': exp_names, - 'params': exp_params, - 'grad_info': grad_info, - } - names, params = exp_names, exp_params - else: - # Hot path: reuse cached params, only update expert grads. - for (orig_idx, num_experts, exp_start, is_dt, tp_mesh, - tp_pls) in cache['grad_info']: - p = params[orig_idx] - g = p.grad - local_grad = (g.to_local() - if is_dt and isinstance(g, DTensor) else g) - for i in range(num_experts): - expert_p = cache['params'][exp_start + i] - sg = local_grad[i] - if tp_mesh is not None: - expert_p.grad = DTensor.from_local( - sg, device_mesh=tp_mesh, placements=tp_pls) - else: - expert_p.grad = sg - p.grad = None - - names = cache['names'] - params = cache['params'] - else: - names, params = _expand_expert_params(names, params, - self.expert_keys) - logger.debug("[_step_muon] after expand: %d params", len(params)) - - param_dtensors = [] - name_dtensors = [] - - param_tensors = [] - name_tensors = [] - - # distributed_muon is a reference implementation for testing only. - # The parallel pipeline (all2all) path below is the production path. - if self.use_distributed_muon: - _run_deferred_expert_ns() - self.distributed_muon(names=names, - params=params, - group=group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits) - return - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - logger.debug( - "[route] %s → base (DTensor all-Replicate), " - "shape=%s, placements=%s", n, p.shape, p.placements) - param_tensors.append(p) - name_tensors.append(n) - else: - logger.debug( - "[route] %s → parallel (DTensor), shape=%s, " - "placements=%s, mesh=%s", n, p.shape, p.placements, - p.device_mesh.mesh_dim_names) - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - logger.debug("[route] %s → base (plain tensor), shape=%s", n, - p.data.shape) - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError(f"Unsupported parameter type: {type(p.data)}") - - logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, " - f"{len(param_tensors)} Tensors → base") - - def group_dtensors(dtensors, names): - # To support different placements, we group parameters by placements - # and run parallel Muon on each group. - - placement_to_params = defaultdict(lambda: ([], [])) - - assert len(dtensors) == len(names) - for p, n in zip(dtensors, names): - placement_to_params[tuple([p.placements, - p.device_mesh])][0].append(n) - placement_to_params[tuple([p.placements, - p.device_mesh])][1].append(p) - return placement_to_params - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - dtensor_group = group_dtensors(param_dtensors, name_dtensors) - - # Pre-launch the first chunk's A2A gather so that the NCCL - # communication overlaps with the (deferred) batched expert NS - # compute on the default CUDA stream. - prelaunch = None - if deferred_expert_work: - first_names, first_params = next(iter(dtensor_group.values())) - ordered, pts, rnk, csz = self._setup_parallel( - first_names, first_params, group, qk_logits) - first_chunk = ordered[:csz] - if first_chunk: - prelaunch = prelaunch_first_gather(first_chunk, pts, rnk, - group["none_grad"]) - - _run_deferred_expert_ns() - - first_group = True - for _, (names, params) in dtensor_group.items(): - pg = prelaunch if first_group else None - first_group = False - self.parallel( - names, - params, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - prelaunch_gather=pg, - ) - else: - _run_deferred_expert_ns() - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - ) - - def _register_states_for_offload(self): - """Register all optimizer state tensors with the CPU offload pool. - - Called once after the first step when states have been lazily created. - Offloads all param states (momentum buffers for Muon, moment1/moment2 - for AdamW) to free GPU memory between steps. - """ - pool = self._cpu_offload_pool - tracked = 0 - for group in self.param_groups: - for p in group["params"]: - if p not in self.state: - continue - state = self.state[p] - if group.get("use_muon", False): - if "momentum_buffer" in state: - pool.track(state["momentum_buffer"]) - tracked += 1 - else: - if "moment1" in state: - pool.track(state["moment1"]) - if "moment2" in state: - pool.track(state["moment2"]) - tracked += 1 - logger.info("[CPUOffload] Registered %d param states for offload", - tracked) - - @torch.no_grad - def step(self, closure=None, qk_logits=None): - """Perform a single optimization step. - - Args: - closure (Callable, optional): A closure that reevaluates the model - and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). - """ - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # H2D: reload optimizer states from CPU before computation. - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload() - - logger.debug("[Muon.step] expert_keys=%s, %d param groups", - self.expert_keys, len(self.param_groups)) - - for i, group in enumerate(self.param_groups): - if group["use_muon"]: - logger.debug("[Muon.step] group %d: use_muon=True, %d params", - i, len(group["params"])) - self._step_muon(group, qk_logits=qk_logits) - else: - logger.debug( - "[Muon.step] group %d: use_muon=False (AdamW), %d params", - i, len(group["params"])) - step_adamw(self.state, group) - - # D2H: offload optimizer states to CPU after computation. - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - return loss - - # ------------------------------------------------------------------ - # CPU offload public helpers - # ------------------------------------------------------------------ - - def turn_on_cpu_offload(self): - """Enable CPU offload for optimizer states.""" - if self.cpu_offload: - return - logger.info("[Muon] turn_on_cpu_offload") - self.cpu_offload = True - if not self.state: - return - self._cpu_offload_pool = CPUOffloadPool() - self._offload_initialized = False - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_off_cpu_offload(self): - """Disable CPU offload and keep optimizer states resident on GPU.""" - if not self.cpu_offload: - return - logger.info("[Muon] turn_off_cpu_offload") - if self._offload_initialized: - self._cpu_offload_pool.reload() - torch.cuda.current_stream().synchronize() - self._cpu_offload_pool = None - self._offload_initialized = False - self.cpu_offload = False - - # ------------------------------------------------------------------ - # Checkpoint support for cpu_offload - # ------------------------------------------------------------------ - - def state_dict(self) -> dict: - if self.cpu_offload: - raise RuntimeError( - "Muon.state_dict() requires turn_off_cpu_offload() before checkpoint save." - ) - return super().state_dict() - - def load_state_dict(self, state_dict: dict) -> None: - if self.cpu_offload: - raise RuntimeError( - "Muon.load_state_dict() requires turn_off_cpu_offload() before checkpoint load." - ) - super().load_state_dict(state_dict) - - # Invalidate adamw.py's module-level tensor caches so that - # the next step rebuilds them with the newly loaded state tensors. - _placement_cache.clear() - _tensor_cache.clear() diff --git a/build/torch210-cxx11-cu130-x86_64-linux/newton_schulz.py b/build/torch210-cxx11-cu130-x86_64-linux/newton_schulz.py deleted file mode 100644 index d939264b69a34e7a3fa78859f34dc265a1159d59..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-x86_64-linux/newton_schulz.py +++ /dev/null @@ -1,240 +0,0 @@ -from itertools import repeat -from math import inf, sqrt - -import numpy as np -import torch - -from .matmul_transpose_triton import matmul_transpose_assign - -COMM_DTYPE = torch.bfloat16 -DEFAULT_CHUNK_SIZE_RATIO = 4 - - -def _optimal_quintic(l, u, max_iter=1000): - """ - Use the simplified Remez algorithm to find the optimal odd quintic approximant - to the constant function x -> 1 over the interval [l, u]. - - Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum - approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the - two interior equioscillation nodes q, r until convergence. Returns the - closed-form equioscillating solution when l ≈ u. - - Raises ValueError if any intermediate value (a, b, c, E, q, r) is non-finite - (NaN or inf). Raises RuntimeError if convergence is not reached within - max_iter iterations. - """ - assert 0 <= l <= u - if 1 - 5e-6 <= l / u: - return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5) - q = (3 * l + u) / 4 - r = (l + 3 * u) / 4 - E = inf - for _ in range(max_iter): - old_E = E - LHS = np.array( - [ - [l, l**3, l**5, 1], - [q, q**3, q**5, -1], - [r, r**3, r**5, 1], - [u, u**3, u**5, -1], - ] - ) - a, b, c, E = np.linalg.solve(LHS, np.ones(4)) - if not np.all(np.isfinite([a, b, c, E])): - raise ValueError( - f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}" - ) - q, r = np.sqrt( - (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / (10 * c) - ) - if not np.all(np.isfinite([q, r])): - raise ValueError(f"_optimal_quintic: non-finite node update q={q}, r={r}") - if abs(old_E - E) <= 1e-15: - break - else: - raise RuntimeError( - f"_optimal_quintic: did not converge after {max_iter} iterations" - ) - return float(a), float(b), float(c) - - -def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0): - """ - Compute the Polar Express coefficient series for `num_iters` quintic iterations. - - Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that - compose to map singular values from [l, 1] toward 1. At each step: - 1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion` - prevents near-zero singular values from stalling by raising the effective - lower bound; if it is active (cushion*u > l), the coefficients are - rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u]. - 2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the - last iteration, providing numerical headroom at the cost of a slightly slower - final convergence step. - 3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1). - - Returns a list of (a, b, c) tuples, one per iteration. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - u = 1 - assert 0 <= l <= u - safety_factor = 1 + safety_factor_eps - coefficients = [] - for iter in range(num_iters): - a, b, c = _optimal_quintic(max(l, cushion * u), u) - if cushion * u > l: - pl = a * l + b * l**3 + c * l**5 - pu = a * u + b * u**3 + c * u**5 - rescaler = 2 / (pl + pu) - a *= rescaler - b *= rescaler - c *= rescaler - if iter < num_iters - 1: - a /= safety_factor - b /= safety_factor**3 - c /= safety_factor**5 - coefficients.append((a, b, c)) - l = a * l + b * l**3 + c * l**5 - u = 2 - l - return coefficients - - -# Precomputed Polar Express coefficients (a, b, c) for 10 quintic Newton-Schulz -# iterations. Each tuple is the minimax-optimal (Remez/equioscillation) odd quintic -# approximant to x->1 over the current singular-value interval, computed once at -# import time and reused across all optimizer steps. -# -# Contrast with the former hardcoded NS coefficients (5 fixed tuples): -# - Former: empirically tuned to maximize slope at zero; did not converge -# singular values to 1, yielding US'V^T with S' ~ Uniform(0.5, 1.5) instead -# of the true polar factor UV^T. -# - Polar Express: analytically optimal per step, adapting to the shrinking -# singular-value interval [l, u] as iterations progress; converges all -# singular values to 1, producing the exact polar factor UV^T. -_coeffs_list = _optimal_composition( - l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02 -) - - -# This code is adapted from: -# KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py) -# NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress) -# matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon) -@torch.no_grad() -def _zeropower_via_newtonschulz5(G, steps): - """ - Compute the polar factor of G via the Polar Express method. - - Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c) - are the Polar Express coefficients from `_coeffs_list`. Each step is the - optimal odd quintic approximant to x -> 1 over the current singular-value - interval, minimizing the maximum approximation error (Remez / minimax criterion). - The composition maps singular values from [l, 1] to near 1, producing the - polar factor (orthogonal factor in the polar decomposition G = UP). - - `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2, - cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE - X = G # no manual typecast - - if G.size(0) > G.size(1): - X = X.T - - X = X / (X.norm() + 1e-7) - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - # Perform the NS iterations - for a, b, c in hs: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(0) > G.size(1): - X = X.T - - return X - - -@torch.no_grad() -def _zeropower_via_newtonschulz5_batched(G, steps): - """Batched polar factor computation for 3D (E, out, in) tensors. - - Same algorithm as ``_zeropower_via_newtonschulz5`` but uses - ``torch.bmm`` / ``torch.baddbmm`` instead of the 2D Triton kernel, - processing all E expert matrices in a single batched call. - """ - assert len(G.shape) == 3 - assert G.dtype == COMM_DTYPE - X = G - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - # Per-expert Frobenius norm. - X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) - - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - for a, b, c in hs: - buf1 = torch.bmm(X, X.transpose(-2, -1)) - buf2 = torch.bmm(buf1, buf1.transpose(-2, -1)) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.baddbmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - return X - - -_ns_per_shape: dict[tuple[int, ...], callable] = {} -_use_compile = True - - -def set_ns_compile(enabled: bool): - """Toggle torch.compile for Newton-Schulz iteration.""" - global _use_compile - _use_compile = enabled - - -def zeropower_via_newtonschulz5(G, steps=5): - if not _use_compile: - return _zeropower_via_newtonschulz5(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() - - -def zeropower_via_newtonschulz5_batched(G, steps=5): - """Compile-cached batched Newton-Schulz for 3D expert tensors.""" - if not _use_compile: - return _zeropower_via_newtonschulz5_batched(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile( - _zeropower_via_newtonschulz5_batched, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() diff --git a/build/torch210-cxx11-cu130-x86_64-linux/optimizer/__init__.py b/build/torch210-cxx11-cu130-x86_64-linux/optimizer/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-x86_64-linux/optimizer/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch210-cxx11-cu130-x86_64-linux/pipeline.py b/build/torch210-cxx11-cu130-x86_64-linux/pipeline.py deleted file mode 100644 index c0c2d515856182d8d15ad27dd4e4e093b29397d6..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-x86_64-linux/pipeline.py +++ /dev/null @@ -1,468 +0,0 @@ -import logging -from typing import Generator - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -from .core import _muon_state, adjust_lr_for_muon -from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5 -from .qk_clip import compute_scales - -logger = logging.getLogger(__name__) - -# ====================================================================== -# Stage helpers -# ====================================================================== - - -def _launch_gather( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Allocate gather buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_gather``). - gathered_grads: ``{id(p): empty_tensor}`` for owned params, - ``None`` for non-owned. - recv_counts: Per-source-rank element counts. - """ - # Allocate gathered-grad buffers - gathered_grads: dict[int, torch.Tensor | None] = {} - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - gathered_grads[id(p)] = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - else: - gathered_grads[id(p)] = None - - # Build send buffer – batch grad copies via torch.cat - # (1-2 fused kernels vs N individual narrow().copy_() calls). - send_counts = [0] * num_ranks - for p in params: - state = param_to_state[id(p)] - send_counts[state.worker_rank] += state.rank_numels[rank] - - total_send = sum(send_counts) - if total_send > 0: - # Group grad slices by destination rank in a single pass. - dst_to_grads = [[] for _ in range(num_ranks)] - for p in params: - state = param_to_state[id(p)] - n = state.rank_numels[rank] - if n > 0: - g = p.grad.to_local() - dst_to_grads[state.worker_rank].append(g.reshape(-1)) - - # Flatten in dst order and cat once. - all_slices = [] - for dst in range(num_ranks): - all_slices.extend(dst_to_grads[dst]) - send_buf = torch.cat(all_slices) - if send_buf.dtype != COMM_DTYPE: - send_buf = send_buf.to(COMM_DTYPE) - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += state.rank_numels[src] - recv_counts[src] = total - - recv_buf = torch.empty(sum(recv_counts), dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - logger.debug(f"send_buf size: {send_buf.numel()}, " - f"recv_buf size: {recv_buf.numel()}, " - f"recv_counts: {recv_counts}, " - f"send_counts: {send_counts}, " - f"process_group: {str(process_group)}") - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, gathered_grads, recv_counts - - -def _complete_gather( - recv_buf: torch.Tensor, - recv_counts: list[int], - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - param_to_state: dict[int, _muon_state], - rank: int, -) -> None: - """Reconstruct gathered grads from the recv buffer (in-place).""" - off = 0 - for src in range(len(recv_counts)): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - - indices = state.rank_indices[src] - - shard_view = gathered_grads[id(p)][indices] - n = shard_view.numel() - if n == 0: - continue - - sg = recv_buf.narrow(0, off + inner_off, n) - sg = sg.reshape(shard_view.shape) - gathered_grads[id(p)][indices] = sg - - inner_off += n - assert inner_off == block - off += block - - -def _compute_ns( - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - ns_steps: int, -) -> dict[int, torch.Tensor | None]: - """Run Newton-Schulz orthogonalization on owned parameters. - - Returns: - computed_us: ``{id(p): orthogonalized_update}`` for owned params. - """ - computed_us: dict[int, torch.Tensor | None] = {} - for p in owned_params: - u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps) - gathered_grads[id(p)] = None # free gathered grad - computed_us[id(p)] = u - return computed_us - - -def _launch_scatter( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, - computed_us: dict[int, torch.Tensor | None], -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor], list[int]]: - """Allocate scatter buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_scatter``). - scattered_us: Empty dict, populated by ``_complete_scatter`` with - zero-copy views into ``recv_buf``. - recv_counts: Per-source-rank element counts. - """ - # scattered_us is populated by _complete_scatter with zero-copy views - # into recv_buf, avoiding N empty_like allocations + N copy_ calls. - # Pre-seed entries for params whose local shard is empty (rank_numels == 0) - # so _update_params can iterate all params without KeyError. - scattered_us: dict[int, torch.Tensor] = {} - for p in params: - if param_to_state[id(p)].rank_numels[rank] == 0: - scattered_us[id(p)] = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - # Build send buffer – batch via torch.cat - # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls). - send_counts = [0] * num_ranks - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - for dst_rank in range(num_ranks): - send_counts[dst_rank] += state.rank_numels[dst_rank] - - total_send = sum(send_counts) - if total_send > 0: - # Cache u_full conversions to avoid redundant .to() per dst_rank. - u_fulls = {} - for p in owned_params: - u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous() - - # Collect slices in dst order (matches all-to-all send layout). - all_slices = [] - for dst_rank in range(num_ranks): - for p in owned_params: - state = param_to_state[id(p)] - su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten() - if su.numel() > 0: - all_slices.append(su) - - send_buf = torch.cat(all_slices) if all_slices else torch.empty( - 0, dtype=COMM_DTYPE, device="cuda") - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += state.rank_numels[rank] - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, scattered_us, recv_counts - - -def _complete_scatter( - recv_buf: torch.Tensor, - recv_counts: list[int], - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], -) -> None: - """Populate scattered_us with zero-copy views into recv_buf. - - Instead of pre-allocating tensors and copying, we assign views directly - from ``recv_buf``. This eliminates N ``empty_like`` + N ``copy_`` calls. - The underlying storage of ``recv_buf`` is kept alive through the views - until ``scattered_us`` is cleared after ``_update_params``. - """ - off = 0 - for src in range(len(recv_counts)): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = state.rank_numels[rank] - if n == 0: - continue - - scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - - inner_off += n - - assert inner_off == block - off += block - - -def _update_params( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], - lr: float, - weight_decay: float, -) -> None: - """Apply weight decay, Muon update, and optional QK clipping. - - Uses batched ``_foreach_mul_`` for weight decay and batched - ``_foreach_add_`` for the Muon update, grouping parameters by - adjusted_lr to minimize kernel launches while preserving float32 - precision for the alpha scaling. - """ - if not params: - return - - # Batched weight decay: p *= (1 - lr * wd) — single fused kernel. - p_locals = [p._local_tensor for p in params] - torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay) - - # Group params by adjusted_lr so _foreach_add_ can use a single - # alpha per group (preserves float32 precision for alpha scaling). - lr_groups: dict[float, tuple[list, list]] = {} - for p in params: - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - if adjusted_lr not in lr_groups: - lr_groups[adjusted_lr] = ([], []) - lr_groups[adjusted_lr][0].append(p._local_tensor) - lr_groups[adjusted_lr][1].append(scattered_us[id(p)]) - - for adjusted_lr, (p_group, u_group) in lr_groups.items(): - torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr) - - # QK clipping – applied directly on the local tensor to - # avoid DTensor sharding-propagation issues with _StridedShard. - for p in params: - state = param_to_state[id(p)] - if state.qk_clip_state is None: - continue - scales_full = compute_scales(p, state.qk_clip_state) - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = state.rank_indices[rank][0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - -# ====================================================================== -# Pre-launch helper for overlapping first chunk's gather with other work. -# ====================================================================== - - -@torch.no_grad() -def prelaunch_first_gather( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - none_grad: bool, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Launch the first chunk's A2A gather early for overlap with other compute. - - Call this *before* expensive GPU work (e.g. batched expert NS) so that - the NCCL all-to-all runs concurrently on the NCCL stream while the - default stream executes compute. - - Returns the same 4-tuple that ``_launch_gather`` produces, which should - be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - with record_function("muon::prelaunch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - return work, recv_buf, gathered_grads, recv_counts - - -# ====================================================================== -# Main generator – thin orchestrator that wires stages together. -# ====================================================================== - - -@torch.no_grad() -def muon_chunk_pipeline( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - ns_steps: int, - lr: float, - weight_decay: float, - none_grad: bool, - prelaunch_gather: tuple | None = None, -) -> Generator[None, None, None]: - """Process one chunk of parameters through the full Muon pipeline. - - Stages: gather -> compute (Newton-Schulz) -> scatter -> update. - - Each ``yield`` lets :func:`run_pipeline` interleave other chunks so - that communication and computation overlap across chunks. Async - communication is launched via ``async_op=True`` and completed after - the yield with ``work.wait()``. - - Overlap happens because :func:`run_pipeline` admits one new chunk - per iteration (staggered admission). While chunk *N* does NS - compute on the default CUDA stream, chunk *N+1*'s async all-to-all - runs concurrently on the NCCL stream — no separate ``comm_stream`` - is required. - - If ``prelaunch_gather`` is provided, the gather was already launched - by :func:`prelaunch_first_gather` and we skip launching it again. - - Yields exactly **2** times: - - 1. After launching async all-to-all gather (or immediately if pre-launched). - 2. After launching async all-to-all scatter. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - if prelaunch_gather is not None: - # Gather was pre-launched; none_grad already handled by caller. - work, recv_buf, gathered_grads, recv_counts = prelaunch_gather - else: - # Normal path: launch async gather. - with record_function("muon::launch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - yield # --- YIELD 1: other chunks can launch their gather --- - - with record_function("muon::wait_gather"): - work.wait() - _complete_gather(recv_buf, recv_counts, owned_params, gathered_grads, - param_to_state, rank) - del recv_buf - - # Stage 3: Newton-Schulz orthogonalization. - with record_function("muon::newton_schulz"): - computed_us = _compute_ns(owned_params, gathered_grads, ns_steps) - gathered_grads.clear() - - # Stages 4-5: launch async scatter. - with record_function("muon::launch_scatter"): - work, recv_buf, scattered_us, recv_counts = _launch_scatter( - params, owned_params, param_to_state, rank, num_ranks, - process_group, computed_us) - computed_us.clear() - - yield # --- YIELD 2: other chunks can launch their scatter --- - - with record_function("muon::wait_scatter"): - work.wait() - _complete_scatter(recv_buf, recv_counts, params, param_to_state, rank, - scattered_us) - del recv_buf - - # Stage 6: apply parameter updates. - with record_function("muon::update_params"): - _update_params(params, param_to_state, rank, scattered_us, lr, - weight_decay) - scattered_us.clear() diff --git a/build/torch210-cxx11-cu130-x86_64-linux/qk_clip.py b/build/torch210-cxx11-cu130-x86_64-linux/qk_clip.py deleted file mode 100644 index 2aba711b3004b7f09e7141da7ef834bd61cc2430..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-x86_64-linux/qk_clip.py +++ /dev/null @@ -1,198 +0,0 @@ -import logging -import math -from dataclasses import dataclass - -import torch -from torch.distributed.tensor import DTensor - -from .core import normalize_fqn - -logger = logging.getLogger(__name__) - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - and return (kind, layer_index). - - Supported kinds: - MHA/GQA: 'wq', 'wk', 'q_proj', 'k_proj' - MLA: 'wq_b' (Q up-proj), 'wkv_b' (KV up-proj) - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.1.attn.wq_b.weight' -> ('wq_b', 1) - 'model.0.attn.wkv_b.weight' -> ('wkv_b', 0) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = normalize_fqn(name).split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: str | None # 'wq'/'q_proj'/'wq_b' or 'wk'/'k_proj'/'wkv_b' or None - indices: list[int] # which heads to consider for clipping - head_dim: int # from config (qk_head_dim for MLA wq_b) - threshold: float # from config - logit: torch.Tensor | None - - # MLA-specific fields - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 - - -def get_qk_clip_info(clip_config, n, qk_logits): - """Extract QK clipping info for a named parameter. - - Args: - clip_config: QK clipping configuration dict (or None). - MHA/GQA keys: head_dim, threshold, q_indices, k_indices - MLA extra keys: is_mla=True, qk_nope_head_dim, qk_rope_head_dim, v_head_dim - n: Parameter name string. - qk_logits: Dict mapping layer indices to logit tensors (or None). - - Returns: - QKClipInfo instance with clipping configuration for this parameter. - """ - if clip_config is None: - return None - - head_dim = clip_config.get('head_dim') - threshold = clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - # In TP settings, qk_logits may be DTensor - # We convert it to full tensor here for simplicity - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - -def compute_scales(p, qk_clip_state): - """Compute per-head scaling factors for QK clipping. - - Returns scales tensor (√γ per head) if any head exceeds threshold, else None. - For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim. - """ - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - # Check if any head exceeds threshold before allocating. - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - - if not head_scales: - return None - - # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows - if qk_clip_state.is_mla and kind == 'wkv_b': - effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim - else: - effective_head_dim = head_dim - - H_global = p.shape[0] // effective_head_dim - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale - return scales_full - - -def qk_clip(p, scales, info): - """Apply per-head scaling to a Q/K projection weight matrix. - - Args: - p: Parameter (nn.Parameter or raw tensor). - scales: [n_heads] tensor, each element = √γ_h. - info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions. - - MLA sub-region scaling per Algorithm 1 (MuonClip): - wq_b: q_nope rows → √γ, q_pe rows → γ - wkv_b: k_nope rows → √γ, v rows → unchanged - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not info.is_mla: - # MHA/GQA: uniform √γ applied to all rows in each head - W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: vectorized sub-region scaling within each head - if info.kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_head_dim = qk_nope + info.qk_rope_head_dim - W_3d = W.view(-1, qk_head_dim, W.shape[1]) # [H, qk_head_dim, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # q_nope → √γ - W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1, - 1)) # q_pe → γ - - elif info.kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - W_3d = W.view(-1, kv_stride, W.shape[1]) # [H, kv_stride, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # k_nope → √γ - # v rows: not touched (k_R shared rotary unchanged) diff --git a/build/torch210-cxx11-rocm70-x86_64-linux/_ops.py b/build/torch210-cxx11-rocm70-x86_64-linux/_ops.py deleted file mode 100644 index 034bff088659b5df6f6d401feb18c89dc5f33b29..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm70-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _optimizer_8d53b78_dirty -ops = torch.ops._optimizer_8d53b78_dirty - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_optimizer_8d53b78_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch210-cxx11-rocm70-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so b/build/torch210-cxx11-rocm70-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so deleted file mode 100755 index d24aaf499164890a854ed9b06cfb6439f0c392a9..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm70-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:055206c495ecade2fe4b5427db34f0a48152174e79808cbe1ce7d7ca86d32396 -size 1866400 diff --git a/build/torch210-cxx11-rocm70-x86_64-linux/adamw.py b/build/torch210-cxx11-rocm70-x86_64-linux/adamw.py deleted file mode 100644 index b5a95816a9f5b9e1889eaadae65373bfbced809a..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm70-x86_64-linux/adamw.py +++ /dev/null @@ -1,271 +0,0 @@ -import logging -from collections import defaultdict -from typing import cast - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -def fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: float | torch.Tensor, - weight_decay: float, - eps: float, - maximize: bool, -) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: dict | None = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) - - -def _to_local(t): - """Unwrap DTensor to local tensor for fused ops.""" - return t._local_tensor if isinstance(t, DTensor) else t - - -# --------------------------------------------------------------------------- -# Caches for eliminating per-step Python overhead. -# -# Placement grouping and tensor list assembly are identical every step -# (params don't change placement, moment/step tensors are the same objects -# after initialisation). We cache them keyed by id() of the param list -# stored in param_groups (stable across steps). -# -# Only gradients change each step and must be collected fresh. -# --------------------------------------------------------------------------- - -# id(group["params"]) → dict[placement_key, list[param]] -_placement_cache: dict[int, dict[tuple, list]] = {} - -# id(placement_group_list) → (params_local, moment1, moment2, state_steps) -_tensor_cache: dict[int, tuple[list, list, list, list]] = {} - - -def _step_adamw_params_slow(optimizer_state, params, group): - """Uncached fallback for the rare case where some params lack grads.""" - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - g = p.grad - if g is None: - continue - state = optimizer_state[p] - params_with_grads.append(_to_local(p)) - grads.append(_to_local(g)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - if not params_with_grads: - return - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw_params(optimizer_state, params, group): - """Run fused AdamW on a list of parameters sharing the same placement. - - After the first call, cached tensor lists (params_local, moment1, - moment2, state_steps) are reused — only gradients are collected fresh. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - params: List of parameters to update. - group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay. - """ - # Collect grads — the only thing that changes each step. - with record_function("adamw::collect_grads"): - grads = [] - for p in params: - g = p.grad - if g is None: - # Rare: fall back to slow path that filters per-param. - _step_adamw_params_slow(optimizer_state, params, group) - return - grads.append(_to_local(g)) - - tensor_key = id(params) - if tensor_key not in _tensor_cache: - with record_function("adamw::init_tensor_cache"): - params_local = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - state = optimizer_state[p] - params_local.append(_to_local(p)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(p.grad) - state["moment2"] = torch.zeros_like(p.grad) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - _tensor_cache[tensor_key] = (params_local, moment1, moment2, - state_steps) - - params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key] - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - with record_function("adamw::fused_adamw"): - fused_adamw( - params_local, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw(optimizer_state, group): - """Dispatch AdamW step, grouping parameters by type and placement. - - Placement grouping is cached after the first call since params never - change their placement between steps. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - group: Parameter group dict. - """ - params = group["params"] - placement_key = id(params) - - if placement_key not in _placement_cache: - with record_function("adamw::group_by_placement"): - placement_to_params: dict[tuple, - list[torch.Tensor]] = defaultdict(list) - for p in params: - match p: - case DTensor(): - logger.debug( - "[AdamW] DTensor param: shape=%s, placements=%s, " - "mesh=%s, grad=%s", p.shape, p.placements, - p.device_mesh.mesh_dim_names, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple( - [p.placements, p.device_mesh])].append(p) - case torch.Tensor(): - logger.debug( - "[AdamW] plain param: shape=%s, grad=%s", p.shape, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple([torch.Tensor, - None])].append(p) - - logger.debug("[AdamW] %d placement groups, %d total params", - len(placement_to_params), len(params)) - - _placement_cache[placement_key] = dict(placement_to_params) - - for group_params in _placement_cache[placement_key].values(): - step_adamw_params(optimizer_state, group_params, group) diff --git a/build/torch210-cxx11-rocm70-x86_64-linux/async_utils.py b/build/torch210-cxx11-rocm70-x86_64-linux/async_utils.py deleted file mode 100644 index a45c530ac9cad88e3555ec1047a6aa59f225347e..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm70-x86_64-linux/async_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import logging -from typing import Generator - -logger = logging.getLogger(__name__) - - -class _Task: - """Internal: wraps a generator, advances one yield at a time.""" - - def __init__(self, generator: Generator[None, None, None], index: int): - self._generator = generator - self._index = index - self._steps_completed = 0 - self.step() # run to first yield - - def step(self) -> bool: - try: - next(self._generator) - self._steps_completed += 1 - logger.debug("pipeline[%d] completed stage %d", self._index, - self._steps_completed) - return True - except StopIteration: - logger.debug("pipeline[%d] finished after %d stages", self._index, - self._steps_completed) - return False - - def close(self): - self._generator.close() - - -def run_pipeline( - pipelines: Generator[Generator[None, None, None], None, None], - max_concurrent: int, -) -> None: - """Run generator-based pipelines with bounded concurrency. - - Each pipeline is a generator that yields at stage boundaries. - The runtime interleaves pipelines so communication and computation - overlap across chunks. - """ - if max_concurrent <= 0: - raise ValueError(f"max_concurrent must be > 0, got {max_concurrent}") - - have_new = True - task_index = 0 - previous_tasks: list[_Task] = [] - - try: - while have_new or previous_tasks: - running_tasks: list[_Task] = [] - - # Admit one new pipeline per iteration (staggered admission). - # Admitting one at a time ensures that while chunk N does NS - # compute on the default stream, chunk N+1's NCCL all-to-all - # runs concurrently on the NCCL stream — creating real - # communication/computation overlap on the GPU. - if have_new and len(previous_tasks) < max_concurrent: - try: - gen = next(pipelines) - task = _Task(gen, task_index) - task_index += 1 - running_tasks.append(task) - except StopIteration: - have_new = False - - # Advance every previously-yielded task by one step. - for task in previous_tasks: - if task.step(): - running_tasks.append(task) - - previous_tasks = running_tasks - except BaseException: - # Clean up all in-flight generators to release GPU resources. - for task in previous_tasks: - task.close() - raise diff --git a/build/torch210-cxx11-rocm70-x86_64-linux/core.py b/build/torch210-cxx11-rocm70-x86_64-linux/core.py deleted file mode 100644 index c69d515afef305ad0ed66374095fa2d2468d99cc..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm70-x86_64-linux/core.py +++ /dev/null @@ -1,219 +0,0 @@ -import logging -import math -from dataclasses import dataclass -from typing import List - -import torch -from torch.distributed import ProcessGroup -from torch.distributed.tensor import DTensor - -# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into -# parameter FQNs. Activation checkpointing similarly inserts -# "_checkpoint_wrapped_module". Strip these so name-based matching (skip_keys, -# expert_keys, QK layer parsing) works regardless of wrapper nesting. -_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) - -logger = logging.getLogger(__name__) - - -def normalize_fqn(name: str) -> str: - """Strip torch.compile / checkpoint wrapper components from a parameter FQN.""" - return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS) - - -@dataclass -class _muon_state: - worker_rank: int - process_group: ProcessGroup - rank_indices: dict[int, tuple] # local_rank -> per-dim indices - rank_numels: dict[int, int] # local_rank -> numel - name: str - qk_clip_state: torch.Tensor | None = None - - -def _batch_momentum( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update (no nesterov).""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - - -def _batch_momentum_nesterov( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update with nesterov correction.""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - nesterov_terms = torch._foreach_mul(momentum_bufs, momentum) - torch._foreach_add_(grads, nesterov_terms) - - -_compiled_momentum: dict[bool, callable] = {} -_use_momentum_compile = True - - -def set_momentum_compile(enabled: bool): - """Toggle torch.compile for batched momentum.""" - global _use_momentum_compile - _use_momentum_compile = enabled - - -def batch_pre_ortho( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, - nesterov: bool, -) -> None: - """Batched momentum update on lists of plain tensors. - - Mirrors dion's ``muon_update_pre_orthogonalize``. - Inputs must be plain CUDA tensors (not DTensor). - Modifies ``momentum_bufs`` and (for nesterov) ``grads`` in-place. - - When compile is enabled, uses separately compiled functions for - nesterov=True/False to avoid graph breaks from the branch. - """ - fn = _batch_momentum_nesterov if nesterov else _batch_momentum - if _use_momentum_compile: - if nesterov not in _compiled_momentum: - _compiled_momentum[nesterov] = torch.compile(fn) - fn = _compiled_momentum[nesterov] - fn(grads, momentum_bufs, momentum) - - -def _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay): - """Weight-decay + update on plain tensors. - - Not compiled: per-param @torch.compile caused ~0.25ms TorchDynamo cache - lookup per call × 256+ params = massive overhead. The pipeline path uses - batched _foreach_* ops instead; this function remains for base() and - distributed_muon(). - """ - p_data.mul_(1 - lr * weight_decay) - p_data.add_(u_data, alpha=-adjusted_lr) - - -def update_p(p, u, lr, adjusted_lr, weight_decay): - """Apply weight decay and orthogonalized update to parameter. - - Args: - p: Parameter (torch.nn.Parameter or DTensor). - u: Orthogonalized update tensor. - lr: Base learning rate. - adjusted_lr: Size-adjusted learning rate. - weight_decay: Weight decay coefficient. - """ - # Unwrap Parameter -> underlying data tensor. - p_data = p.data if isinstance(p, torch.nn.Parameter) else p - # Unwrap DTensor -> local CUDA tensor for compiled kernel. - if isinstance(p_data, DTensor): - p_data = p_data._local_tensor - u_data = u._local_tensor if isinstance(u, DTensor) else u - _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay) - - -def adjust_lr_for_muon(lr, param_shape): - """Scale learning rate based on parameter matrix dimensions. - - Args: - lr: Base learning rate. - param_shape: Shape of the parameter tensor. - - Returns: - Adjusted learning rate. - """ - A, B = param_shape[:2] - # We adjust the learning rate and weight decay based on the size of the parameter matrix - # as described in the paper - adjusted_ratio = 0.2 * math.sqrt(max(A, B)) - adjusted_lr = lr * adjusted_ratio - return adjusted_lr - - -def _match_key(parts, key): - """Check if key matches as contiguous components in parts. - - Single-component keys (e.g. "experts") match any single component. - Multi-component keys (e.g. "experts.w1") match as a contiguous subsequence. - """ - key_parts = key.split(".") - key_len = len(key_parts) - if key_len == 1: - return key in parts - return any(parts[i:i + key_len] == key_parts - for i in range(len(parts) - key_len + 1)) - - -def is_expert_param(name, expert_keys): - """Check if a parameter name matches any expert key (component-level).""" - if not expert_keys: - return False - parts = normalize_fqn(name).split(".") - return any(_match_key(parts, key) for key in expert_keys) - - -def default_is_muon(name, x, expert_keys=None): - normalized = normalize_fqn(name) - parts = normalized.split(".") - skip_keys = [ - "embed_tokens", - "lm_head", - "tok_embeddings", - "output", - "mhc_attn", - "mhc_ffn", - "lambda_proj", - ] - if any(key in parts for key in skip_keys): - logger.info( - "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d", - normalized, name, x.ndim) - return False - effective_ndim = x.ndim - is_expert = is_expert_param(name, expert_keys) - if is_expert: - effective_ndim -= 1 - result = effective_ndim >= 2 - logger.info( - "[is_muon] %s (orig: %s): ndim=%d, expert=%s, effective_ndim=%d → %s", - normalized, name, x.ndim, is_expert, effective_ndim, - "Muon" if result else "AdamW") - return result - - -def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None): - if is_muon_func is None: - is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys) - - muon_params, muon_names = [], [] - non_muon_params, non_muon_names = [], [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - non_muon_names.append(n) - - logger.info("[param_groups] expert_keys=%s, Muon=%d, AdamW=%d", - expert_keys, len(muon_names), len(non_muon_names)) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] diff --git a/build/torch210-cxx11-rocm70-x86_64-linux/cpu_offload.py b/build/torch210-cxx11-rocm70-x86_64-linux/cpu_offload.py deleted file mode 100644 index fb5e69154a1d4a6c884491413a37a9acf0f66c80..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm70-x86_64-linux/cpu_offload.py +++ /dev/null @@ -1,206 +0,0 @@ -"""CPU offloading for optimizer states. - -Manages a pinned CPU memory pool and async CUDA streams to offload -optimizer state tensors (momentum buffers, Adam moments) to CPU between -optimizer steps, freeing GPU memory. - -All tracked tensors are packed into a single flat pinned CPU buffer -(per dtype). D2H and H2D copies are performed per-tensor directly -between individual GPU tensors and their slice of the CPU flat buffer -— no GPU staging buffer is allocated, so there is **no temporary GPU -memory spike** during offload or reload. - -Individual tensor storages are freed after offload via -``untyped_storage().resize_(0)``, preserving tensor identity so -downstream caches remain valid. -""" - -import logging -from collections import defaultdict - -import torch -from torch.distributed.tensor import DTensor - -logger = logging.getLogger(__name__) - - -class CPUOffloadPool: - """Pinned CPU memory pool for async optimizer state offloading. - - Tracked tensors are grouped by dtype. Each group gets a single flat - pinned CPU buffer. D2H / H2D copies are per-tensor (into slices of - the flat buffer) to avoid allocating a GPU staging buffer. - """ - - def __init__(self): - self._managed: list[torch.Tensor] = [] - self._storage_nbytes: dict[int, int] = {} # id(t) → bytes - - # Per-dtype group: populated on first offload. - # dtype → dict with keys: - # "indices" : list[int] managed-list indices - # "offsets" : list[tuple[int,int]] (start, numel) in flat buf - # "total" : int total numel - # "cpu_flat" : Tensor pinned CPU buffer - self._groups: dict[torch.dtype, dict] = {} - - self._offload_stream: torch.cuda.Stream | None = None - self._device: torch.device | None = None - self._initialized: bool = False - self._logged: bool = False - - # ------------------------------------------------------------------ - @staticmethod - def _local(t: torch.Tensor) -> torch.Tensor: - """Unwrap DTensor to its local CUDA tensor.""" - return t._local_tensor if isinstance(t, DTensor) else t - - def _ensure_stream(self): - if self._offload_stream is None: - self._offload_stream = torch.cuda.Stream(device=self._device) - - # ------------------------------------------------------------------ - def track(self, tensor: torch.Tensor): - """Register a GPU tensor for CPU offloading. Idempotent.""" - tid = id(tensor) - if tid in self._storage_nbytes: - return - local = self._local(tensor) - if self._device is None: - self._device = local.device - storage = local.untyped_storage() - # Skip tensors with empty storage (e.g. empty FSDP shards) - if storage.size() == 0: - return - self._storage_nbytes[tid] = storage.size() - self._managed.append(tensor) - - # ------------------------------------------------------------------ - def _init_buffers(self): - """Build per-dtype flat buffers on first offload.""" - # Group managed tensors by dtype. - dtype_map: dict[torch.dtype, list[tuple[int, int]]] = defaultdict(list) - for idx, t in enumerate(self._managed): - local = self._local(t) - dtype_map[local.dtype].append((idx, local.numel())) - - total_cpu_bytes = 0 - for dtype, entries in dtype_map.items(): - offsets: list[tuple[int, int]] = [] - indices: list[int] = [] - off = 0 - for idx, n in entries: - indices.append(idx) - offsets.append((off, n)) - off += n - cpu_flat = torch.empty(off, dtype=dtype, device="cpu", pin_memory=True) - self._groups[dtype] = { - "indices": indices, - "offsets": offsets, - "total": off, - "cpu_flat": cpu_flat, - } - total_cpu_bytes += off * cpu_flat.element_size() - - self._initialized = True - logger.info( - "[CPUOffload] Pool initialized: %d tensors, %d dtype group(s), " - "%.2f MB pinned CPU memory", - len(self._managed), - len(self._groups), - total_cpu_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def offload(self): - """Per-tensor async D2H into CPU flat buffer, then free GPU storage.""" - if not self._managed: - return - if not self._initialized: - self._init_buffers() - self._ensure_stream() - - # Offload stream waits for compute to finish. - compute_event = torch.cuda.current_stream(self._device).record_event() - self._offload_stream.wait_event(compute_event) - - offloaded_bytes = 0 - - # Per-tensor D2H copies directly into CPU flat buffer slices. - # No GPU staging buffer → no temporary GPU memory spike. - with torch.cuda.stream(self._offload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - cpu_flat[off : off + n].copy_(local.reshape(-1), non_blocking=True) - - offloaded_bytes += grp["total"] * cpu_flat.element_size() - - # Wait for all D2H copies to land, then free GPU storage. - self._offload_stream.synchronize() - for t in self._managed: - storage = self._local(t).untyped_storage() - if storage.size() != 0: - storage.resize_(0) - else: - raise RuntimeError( - f"Tensor storage is already freed (size=0) before offload. " - f"This indicates a double-free or external interference. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - - if not self._logged: - logger.info( - "[CPUOffload] Offloaded %.2f MB (GPU → CPU)", - offloaded_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def reload(self): - """Per-tensor H2D from CPU flat buffer on the default stream. - - Runs on the current (default) CUDA stream to avoid stream - interaction issues with the parallel Muon pipeline. Since - pinned CPU memory is the source, the copies overlap with - GPU idle time between steps. - """ - if not self._managed or not self._initialized: - return - - reloaded_bytes = 0 - - # Re-allocate all GPU storages first. - for t in self._managed: - local = self._local(t) - storage = local.untyped_storage() - if storage.size() != 0: - raise RuntimeError( - f"Storage should have been freed (size=0) before reload, " - f"but got size={storage.size()}. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - storage.resize_(self._storage_nbytes[id(t)]) - - # Per-tensor H2D copies from CPU flat buffer slices. - # non_blocking=True with pinned source allows DMA overlap. - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off : off + n], non_blocking=True) - - reloaded_bytes += grp["total"] * cpu_flat.element_size() - - if not self._logged: - logger.info( - "[CPUOffload] Reloaded %.2f MB (CPU → GPU)", reloaded_bytes / (1024**2) - ) diff --git a/build/torch210-cxx11-rocm70-x86_64-linux/distributed/utils.py b/build/torch210-cxx11-rocm70-x86_64-linux/distributed/utils.py deleted file mode 100644 index 683096bcd50489e7a09959753fdbc29ba0f3b6a6..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm70-x86_64-linux/distributed/utils.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import (Placement, Shard, - _StridedShard) - - -def _is_shard(placement: Placement) -> bool: - """Check if a placement is a shard type (Shard or _StridedShard). - - In PyTorch 2.10+, _StridedShard no longer inherits from Shard, so - ``placement.is_shard()`` returns False for _StridedShard. This helper - handles both old and new hierarchies. - """ - return isinstance(placement, (Shard, _StridedShard)) - - -def get_slices_of_dtensor( - target: DTensor | torch.Tensor, - local_rank: int, - shard_mesh: DeviceMesh, - shard_placements: tuple[Placement], -) -> tuple[slice | torch.Tensor, ...]: - """ - Get per-dimension indices for a given rank's shard of the target tensor. - - Uses ``Shard.local_shard_size_and_offset`` and - ``_StridedShard.local_shard_size_and_offset`` for correct handling of - both contiguous and strided (non-contiguous) sharding. - - Args: - target (DTensor | torch.Tensor): The target tensor (for its shape). - local_rank (int): The local rank within the shard group. - shard_mesh (DeviceMesh): The shard mesh (only shard dimensions). - shard_placements (tuple[Placement]): The shard placements. - - Returns: - A tuple of indices (one per tensor dim). Each element is either: - - A ``slice`` (for contiguous or unsharded dims) - - A 1-D ``torch.LongTensor`` of indices (for strided sharding) - """ - - # find the global rank of the local rank in the shard mesh - rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank] - - rank_coords = (shard_mesh.mesh == rank).nonzero() - - assert len(rank_coords) == 1 - rank_coords = tuple(rank_coords[0].tolist()) - - assert len(rank_coords) == len(shard_placements) - - # Track per-shard-dim indices. - # None means "not yet sharded on this dim". - dim_indices: dict[int, torch.Tensor] = {} - - # Caution: Assuming replicate-to-shard of the shard mesh goes with - # left-to-right sharding. This is ensured by the sorting logic of - # construct_shard_mesh function. - for mesh_dim_idx, (rank_coord, placement) in enumerate( - zip(rank_coords, shard_placements)): - assert _is_shard(placement) - - num_chunks = shard_mesh.mesh.shape[mesh_dim_idx] - shard_dim = placement.dim - - # Current effective size on this dim (may already be sub-sharded) - if shard_dim in dim_indices: - curr_size = len(dim_indices[shard_dim]) - else: - curr_size = target.size()[shard_dim] - - # Compute indices for this level of sharding - if isinstance(placement, _StridedShard): - _shard_size, offsets = _StridedShard.local_shard_size_and_offset( - placement, - curr_size, - num_chunks, - rank_coord, - return_first_offset=False) - new_indices = torch.tensor(offsets, dtype=torch.long) - else: - shard_size, offset = Shard.local_shard_size_and_offset( - curr_size, num_chunks, rank_coord) - new_indices = torch.arange(offset, - offset + shard_size, - dtype=torch.long) - - # Compose with previous indices on this dim - if shard_dim in dim_indices: - dim_indices[shard_dim] = dim_indices[shard_dim][new_indices] - else: - dim_indices[shard_dim] = new_indices - - # Build result tuple - result: list[slice | torch.Tensor] = [] - for d in range(len(target.size())): - if d not in dim_indices: - result.append(slice(None)) - else: - indices = dim_indices[d] - # Convert contiguous indices to slice for efficiency - if len(indices) > 0: - start = indices[0].item() - expected = torch.arange(start, - start + len(indices), - dtype=torch.long) - if torch.equal(indices, expected): - result.append(slice(start, start + len(indices))) - else: - result.append(indices) - else: - result.append(slice(0, 0)) - - return tuple(result) - - -_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, - ProcessGroup]] = dict() - - -def construct_shard_mesh( - placements: tuple[Placement], - mesh: DeviceMesh, -) -> tuple[DeviceMesh, ProcessGroup, tuple[Placement, ...]]: - """Construct shard sub-mesh and ProcessGroup for all-to-all communication. - - Given a DTensor's placements and device mesh, extracts the "shard group" - — the set of ranks that together hold all shards of the same replica — - and creates a ProcessGroup for all-to-all among them. - - Steps: - 1. Sort placements: Replicate first, then Shard by (dim, granularity). - 2. Permute the mesh tensor to match the sorted order. - 3. Collapse Replicate dims → list of shard sub-meshes (one per replica). - 4. Create/retrieve a cached ProcessGroup for the current rank's sub-mesh. - - Example — 8 GPUs, mesh shape (2, 2, 2), - placements ``[Shard(0), Replicate, _StridedShard(0)]``:: - - Step 1 — Sort: [Replicate, _StridedShard(0), Shard(0)] - Permutation: [1, 2, 0] - - Step 2 — Permute mesh dims by [1, 2, 0]: - Original: Permuted: - [[[0,1],[2,3]], [[[0,2],[1,3]], - [[4,5],[6,7]]] [[4,6],[5,7]]] - - Step 3 — Unbind replicate dim (dim 0), giving 2 shard sub-meshes: - sub-mesh 0 = [[0,2],[1,3]] (replica group 0) - sub-mesh 1 = [[4,6],[5,7]] (replica group 1) - shard_placements = (_StridedShard(0), Shard(0)) - - Step 4 — Rank 0 → ProcessGroup([0,1,4,5]) - Rank 2 → ProcessGroup([2,3,6,7]) - - Returns: - ``(shard_mesh, process_group, shard_placements)`` - """ - my_rank = dist.get_rank() - assert mesh.mesh.device.type == 'cpu' - - # -- Fast path: 1D all-shard mesh → reuse existing PG. ---------------- - # Reuses the mesh's existing ProcessGroup directly, avoiding the - # overhead of dist.new_group(). The standard path below also handles - # subset calls safely via use_local_synchronization=True, but this - # fast path is still beneficial for the common 1D shard case. - if mesh.ndim == 1 and len(placements) == 1 and _is_shard(placements[0]): - key = (*mesh.mesh.shape, *mesh.mesh.flatten().tolist()) - if key not in _ranks_to_dist_cache: - _ranks_to_dist_cache[key] = (mesh, mesh.get_group()) - return (*_ranks_to_dist_cache[key], tuple(placements)) - - mesh_tensor = mesh.mesh.clone() - - # -- Step 1: Sort placements (Replicate first, then Shard by dim). ------ - # _StridedShard comes BEFORE regular Shard on the same dim so that - # get_slices_of_dtensor applies the outer sharding first, matching - # DTensor's left-to-right (outer-to-inner) composition order. - def _sort_key(item): - index, placement = item - assert not placement.is_partial(), "Partial placement not supported" - if placement.is_replicate(): - return (-1, 0, index) - assert _is_shard(placement), f"Unsupported: {type(placement)}" - split = (-1 / placement.split_factor if isinstance( - placement, _StridedShard) else 0) - return (placement.dim, split, index) - - indexed = sorted(enumerate(placements), key=_sort_key) - perm, sorted_placements = zip(*indexed) - - # -- Step 2: Permute mesh to match sorted placement order. -------------- - sorted_mesh = mesh_tensor.permute(perm) - - # -- Step 3: Collapse replicate dims → list of shard sub-meshes. -------- - # E.g. mesh (2, 3, 4, 4) with [R, R, S(0), S(1)] → 6 sub-meshes of (4, 4) - num_rep = sum(1 for p in sorted_placements if p.is_replicate()) - if num_rep > 0: - if num_rep > 1: - sorted_mesh = sorted_mesh.flatten(0, num_rep - 1) - shard_meshes = list(torch.unbind(sorted_mesh, dim=0)) - else: - shard_meshes = [sorted_mesh] - shard_placements = sorted_placements[num_rep:] - assert len(shard_placements) == len(set(shard_placements)) - - # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. -- - # Each rank only creates the group it belongs to, using - # use_local_synchronization=True so that only group members need to - # coordinate. This avoids deadlocks when different PP stages call - # construct_shard_mesh for different parameters. - def _cache_key(t: torch.Tensor) -> tuple: - return (*t.shape, *t.flatten().tolist()) - - my_key = None - for sm in shard_meshes: - if (my_rank == sm).any().item(): - key = _cache_key(sm) - assert my_key is None, "Rank appears in multiple shard groups" - my_key = key - if key not in _ranks_to_dist_cache: - pg = dist.new_group(sm.flatten().tolist(), - use_local_synchronization=True) - _ranks_to_dist_cache[key] = ( - DeviceMesh(device_type="cuda", mesh=sm), - pg, - ) - - return (*_ranks_to_dist_cache[my_key], shard_placements) diff --git a/build/torch210-cxx11-rocm70-x86_64-linux/matmul_transpose_triton.py b/build/torch210-cxx11-rocm70-x86_64-linux/matmul_transpose_triton.py deleted file mode 100644 index 792de23d82c3fb45fe33d397ab9b76a0787259d0..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm70-x86_64-linux/matmul_transpose_triton.py +++ /dev/null @@ -1,122 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], - restore_value=['y'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -@torch.library.custom_op("muon::matmul_transpose_assign", - mutates_args=("d_out", )) -def matmul_transpose_assign(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """Compute d_out = d_in @ d_in.T using an optimized Triton kernel.""" - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -@matmul_transpose_assign.register_fake -def _(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """FakeTensor impl: d_out is already allocated, mutation is declared.""" - pass diff --git a/build/torch210-cxx11-rocm70-x86_64-linux/metadata.json b/build/torch210-cxx11-rocm70-x86_64-linux/metadata.json deleted file mode 100644 index c55a35717622f1dd5c8ba376ea3a814cbcc10d78..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm70-x86_64-linux/metadata.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch210-cxx11-rocm70-x86_64-linux/muon.py b/build/torch210-cxx11-rocm70-x86_64-linux/muon.py deleted file mode 100644 index 14c0e22471fa6d47a51ed95e0e0c341dc18d5194..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm70-x86_64-linux/muon.py +++ /dev/null @@ -1,1068 +0,0 @@ -import logging -import types -from collections import defaultdict -from typing import Any - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Replicate, Shard -from torch.profiler import record_function - -from .adamw import _placement_cache, _tensor_cache, step_adamw -from .async_utils import run_pipeline -from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho, - get_default_muon_param_groups, is_expert_param, update_p) -from .cpu_offload import CPUOffloadPool -from .distributed.utils import (_is_shard, construct_shard_mesh, - get_slices_of_dtensor) -from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO, - _zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5_batched) -from .pipeline import muon_chunk_pipeline, prelaunch_first_gather -from .qk_clip import compute_scales, get_qk_clip_info, qk_clip - -logger = logging.getLogger(__name__) - - -def _expand_expert_params(names, params, expert_keys): - """Expand expert params by splitting on dim 0 (expert dimension). - - Params whose name matches any key in ``expert_keys`` are treated as - expert-parallel tensors. Their outermost dimension is the expert - dimension: an ``(E, out, in)`` tensor becomes ``E`` separate 2D - ``nn.Parameter`` views so that in-place updates propagate back to - the original storage. - - Non-expert params with ``ndim > 2`` trigger an ``AssertionError`` — - if they are expert params, their key must be added to ``expert_keys``. - - The grad must already be set on each expert param (e.g. after momentum). - - For DTensor expert params, placements that shard on dim 0 (expert dim) - are consumed by the split. Non-dim-0 shard placements (e.g. TP) are - preserved: each 2D slice is wrapped as a DTensor on the corresponding - submesh so the parallel pipeline handles the TP communication. - """ - expanded_names = [] - expanded_params = [] - - for n, p in zip(names, params): - is_expert = is_expert_param(n, expert_keys) - is_dtensor = isinstance(p.data, DTensor) - - if is_expert: - if is_dtensor: - logger.debug( - "[expand_expert] %s: expert DTensor, shape=%s, " - "placements=%s, mesh=%s, local_shape=%s", n, p.shape, - p.placements, p.device_mesh.mesh_dim_names, - p.to_local().shape) - else: - logger.debug( - "[expand_expert] %s: expert plain tensor, shape=%s", n, - p.data.shape) - - if not is_expert: - assert p.data.ndim <= 2, ( - f"Param {n} has ndim={p.data.ndim} but does not match " - f"expert_keys={expert_keys}. If this is an expert param, " - f"add its key to expert_keys.") - expanded_names.append(n) - expanded_params.append(p) - continue - - g = p.grad - assert g is not None, ( - f"Expert param {n} must have grad set before expansion") - - tp_mesh = None - tp_placements_2d = None - - if is_dtensor: - local_data = p.to_local() - local_grad = g.to_local() if isinstance(g, DTensor) else g - - # Find non-dim-0 shard placements (e.g. TP sharding). - # After splitting on dim 0, Shard(k) becomes Shard(k-1). - tp_dim_indices = [] - tp_placements_2d = [] - for i, pl in enumerate(p.placements): - if _is_shard(pl) and pl.dim != 0: - tp_dim_indices.append(i) - tp_placements_2d.append(Shard(pl.dim - 1)) - - if tp_dim_indices: - tp_dim_names = tuple(p.device_mesh.mesh_dim_names[i] - for i in tp_dim_indices) - if len(tp_dim_names) == 1: - tp_mesh = p.device_mesh[tp_dim_names[0]] - else: - tp_mesh = p.device_mesh[tp_dim_names] - else: - local_data = p.data - local_grad = g - - # Expand: split dim 0, reshape each slice to 2D. - num_local_experts = local_data.shape[0] - for i in range(num_local_experts): - slice_data = local_data[i] - slice_grad = local_grad[i] - - if tp_mesh is not None: - # Wrap as DTensor on TP submesh so the pipeline handles - # TP communication (gather/scatter across TP ranks). - dt_data = DTensor.from_local(slice_data, - device_mesh=tp_mesh, - placements=tp_placements_2d) - dt_grad = DTensor.from_local(slice_grad, - device_mesh=tp_mesh, - placements=tp_placements_2d) - expert_param = torch.nn.Parameter(dt_data, requires_grad=False) - expert_param.grad = dt_grad - else: - expert_param = torch.nn.Parameter(slice_data, - requires_grad=False) - expert_param.grad = slice_grad - - expanded_names.append(f"{n}[{i}]") - expanded_params.append(expert_param) - - p.grad = None # allow expert grad storage to be freed after pipeline - - return expanded_names, expanded_params - - -class Muon(torch.optim.Optimizer): - """ - Muon - MomentUm Orthogonalized by Newton-schulz - - Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- - processing step, in which each 2D parameter's update is replaced with the nearest orthogonal - matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has - the advantage that it can be stably run in bfloat16 on the GPU. - - Some warnings: - - We believe this optimizer is unlikely to work well for training with small batch size. - - We believe it may not work well for finetuning pretrained models, but we haven't tested this. - - Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. - lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) - momentum: The momentum used by the internal SGD. (0.95 is a good default) - nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) - ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. - Parameters that are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW instead. - adamw_lr: The learning rate for the internal AdamW. - adamw_betas: The betas for the internal AdamW. - adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - warmup_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher warmup_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. - chunk_size : Batch size of parameters to process in each - all2all gather/compute/scatter step. - Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified. - use_distributed_muon: Use distributed muon by Liu et al. (2024). - For testing purpose only. - expert_keys: List of strings to identify expert-parallel parameters. - If any key appears in a parameter's name, its outermost - dimension is treated as the expert dimension and expanded - into per-expert 2D params for Muon. For example, - ``expert_keys=["experts"]`` matches any param whose name - contains "experts". 3D+ params not matched by any key - will raise an error. - """ - - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config=None, - warmup_step=5, - chunk_size=-1, - use_distributed_muon=False, - expert_keys=None): - defaults = dict( - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - nesterov=nesterov, - ns_steps=ns_steps, - adamw_betas=adamw_betas, - adamw_eps=adamw_eps, - none_grad=none_grad, - use_muon=True, - ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) - super().__init__(params, defaults) - - self.debug = debug - self.clip_config = clip_config if clip_config is not None else { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100, - } - self.warmup_step = warmup_step - self.chunk_size = chunk_size - self.use_distributed_muon = use_distributed_muon - self.expert_keys = expert_keys - self.cpu_offload = False - self._cpu_offload_pool: CPUOffloadPool | None = None - self._offload_initialized = False - self._parallel_cache: dict[tuple[str, ...], dict] = {} - self._expert_expand_cache: dict[tuple[int, ...], dict] = {} - - def _calc_flops(self, G, steps): - assert len(G.shape) == 2 - M, N = G.shape - if M > N: - M, N = N, M - - return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) - - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - p.placements, p.device_mesh) - - return shard_mesh, shard_pg, shard_placements - - def init_state_and_assign_params(self, names, params, group, qk_logits): - param_to_state = {} - param_to_flops = {} - - total_flops = 0 - for p in params: - g = p.grad - if g is None: - continue - assert g.ndim == 2, "Muon only supports 2D parameters." - - flops = self._calc_flops(g, group["ns_steps"]) - param_to_flops[id(p)] = flops - total_flops += flops - - if self.debug: - logger.debug("Total TFLOPs for Muon: %.2f TFLOPs", - total_flops / 1e12) - - paired = list(zip(names, params)) - - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) - - round_robin = 0 - mesh = ordered_params[0].device_mesh - placements = ordered_params[0].placements - - shard_mesh, shard_pg, shard_placements = self.get_shard_mesh( - ordered_params[0]) - shard_mesh_flattened = shard_mesh.mesh.flatten() - num_ranks = dist.get_world_size(group=shard_pg) - - for n, p in zip(ordered_names, ordered_params): - if mesh != p.device_mesh: - raise ValueError("All parameters must be on the same mesh.") - if placements != p.placements: - raise ValueError("All parameters must have same placements.") - - worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks - round_robin = (round_robin + 1) % len(shard_mesh_flattened) - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - # Precompute per-rank indices and numels for all-to-all. - rank_indices: dict[int, tuple] = {} - rank_numels: dict[int, int] = {} - for r in range(num_ranks): - indices = get_slices_of_dtensor(p, r, shard_mesh, - shard_placements) - rank_indices[r] = indices - numel = 1 - for idx, dim_size in zip(indices, p.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) - rank_numels[r] = numel - - param_to_state[id(p)] = _muon_state( - worker_rank=worker_rank, - process_group=shard_pg, - rank_indices=rank_indices, - rank_numels=rank_numels, - name=n, - qk_clip_state=qk_clip_state, - ) - - return param_to_state, ordered_params - - def base(self, names, params, group, lr, weight_decay, qk_logits): - # Momentum is already applied by _step_muon before this method. - for n, p in zip(names, params): - g = p.grad - if g is None: - continue - - u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) - - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - def distributed_muon( - self, - names: list[str], - params: list[torch.nn.Parameter], - group: dict[str, Any], - lr: float, - weight_decay: float, - qk_logits: list[torch.Tensor | DTensor] | None, - ): - """Batched Distributed Muon — for testing/correctness verification only. - - Uses all-gather to reconstruct full tensors, computes Newton-Schulz on - the full grad, then slices back to local shards. This is simpler but - slower than the parallel pipeline (all2all) path, so it serves as a - reference implementation for verifying correctness. - """ - with record_function("distributed_muon"): - # Momentum is already applied by _step_muon before this method. - ns_steps = group["ns_steps"] - - # Separate plain tensors (no communication) from DTensors. - plain_names, plain_params = [], [] - dtensor_names, dtensor_params = [], [] - for n, p in zip(names, params): - if p.grad is None: - continue - if isinstance(p.data, DTensor): - dtensor_names.append(n) - dtensor_params.append(p) - else: - plain_names.append(n) - plain_params.append(p) - - # Process plain tensors per-param (no communication). - for n, p in zip(plain_names, plain_params): - u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE), - steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, - qk_logits) - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - if not dtensor_params: - return - - # Group DTensors by (placements, mesh) for batched all-gather. - placement_groups: dict[tuple, - tuple[list, - list]] = defaultdict(lambda: ([], [])) - for n, p in zip(dtensor_names, dtensor_params): - key = (p.placements, p.device_mesh) - placement_groups[key][0].append(n) - placement_groups[key][1].append(p) - - logger.info( - "distributed_muon: %d placement groups, %d total dtensors", - len(placement_groups), len(dtensor_params)) - - for (placements, mesh), (grp_names, - grp_params) in placement_groups.items(): - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - placements, mesh) - rank = dist.get_rank(shard_pg) - world_size = dist.get_world_size(shard_pg) - - logger.info(" group: %d params, placements=%s, world_size=%d", - len(grp_params), placements, world_size) - - # Separate params that can be batched (all shard dims evenly - # divisible) from those needing per-param full_tensor - # (e.g. MoE gate weights with fewer rows than shard ranks). - # all_gather_into_tensor requires equal buffer sizes across - # ranks, so uneven splits must use DTensor full_tensor(). - batch_names, batch_params = [], [] - single_names, single_params = [], [] - for n, p in zip(grp_names, grp_params): - even = all(p.shape[pl.dim] % - shard_mesh.mesh.shape[dim_idx] == 0 - for dim_idx, pl in enumerate(shard_placements)) - if even: - batch_names.append(n) - batch_params.append(p) - else: - single_names.append(n) - single_params.append(p) - - # Process uneven-split params per-param via full_tensor(). - for n, p in zip(single_names, single_params): - with record_function("distributed_muon::newton_schulz"): - g_full = p.grad.full_tensor().to(COMM_DTYPE) - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - if not batch_params: - continue - - logger.info(" batched=%d, single=%d", len(batch_params), - len(single_params)) - - # Concat all local grad shards into a single flat buffer. - with record_function("distributed_muon::gather"): - grad_locals = [ - p.grad.to_local().to(COMM_DTYPE).flatten() - for p in batch_params - ] - numels = [g.numel() for g in grad_locals] - grad_concat = torch.cat(grad_locals) - del grad_locals - - # Single all-gather (replaces N separate full_tensor). - grad_gathered = torch.empty( - grad_concat.numel() * world_size, - dtype=COMM_DTYPE, - device="cuda", - ) - dist.all_gather_into_tensor(grad_gathered, - grad_concat, - group=shard_pg) - - total_numel = grad_concat.numel() - del grad_concat - - # Precompute per-param offsets within the concat buffer. - offsets = [] - off = 0 - for ne in numels: - offsets.append(off) - off += ne - - # Per-param: reconstruct full grad → NS → local update. - for i, (n, p) in enumerate(zip(batch_names, batch_params)): - with record_function("distributed_muon::newton_schulz"): - g_full = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - for r in range(world_size): - r_start = r * total_numel + offsets[i] - shard = grad_gathered[r_start:r_start + numels[i]] - indices = get_slices_of_dtensor( - p, r, shard_mesh, shard_placements) - g_full[indices] = shard.reshape( - g_full[indices].shape) - - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - def _setup_parallel(self, names, params, group, qk_logits): - """Compute (or retrieve cached) parallel pipeline metadata. - - Returns: - (ordered_params, param_to_state, rank, chunk_size) - """ - cache_key = tuple(names) - - if cache_key not in self._parallel_cache: - # First call: compute metadata and populate cache. - param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - shard_pg = param_to_state[id(ordered_params[0])].process_group - rank = dist.get_rank(group=shard_pg) - - if self.chunk_size == -1: - shard_ranks = dist.get_world_size(shard_pg) - chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO - elif self.chunk_size > 0: - chunk_size = self.chunk_size - else: - raise ValueError( - "chunk_size must be -1 or a positive integer.") - - ordered_names = [ - param_to_state[id(p)].name for p in ordered_params - ] - name_to_state = { - param_to_state[id(p)].name: param_to_state[id(p)] - for p in ordered_params - } - self._parallel_cache[cache_key] = { - 'ordered_names': ordered_names, - 'name_to_state': name_to_state, - 'rank': rank, - 'chunk_size': chunk_size, - } - else: - # Cached path: rebuild param_to_state with current id(p) keys. - cache = self._parallel_cache[cache_key] - rank = cache['rank'] - chunk_size = cache['chunk_size'] - - name_to_param = dict(zip(names, params)) - ordered_params = [name_to_param[n] for n in cache['ordered_names']] - - param_to_state = {} - for p, n in zip(ordered_params, cache['ordered_names']): - cached_state = cache['name_to_state'][n] - param_to_state[id(p)] = _muon_state( - worker_rank=cached_state.worker_rank, - process_group=cached_state.process_group, - rank_indices=cached_state.rank_indices, - rank_numels=cached_state.rank_numels, - name=n, - qk_clip_state=get_qk_clip_info(self.clip_config, n, - qk_logits), - ) - - return ordered_params, param_to_state, rank, chunk_size - - def parallel(self, - names, - params, - group, - lr, - weight_decay, - qk_logits, - prelaunch_gather=None): - """ - Perform a parallel optimization step using Muon. - - Parameters are chunked and each chunk is processed by a - :func:`muon_chunk_pipeline` generator. :func:`run_pipeline` - interleaves multiple chunks so that communication and computation - overlap across chunks (the same overlap previously achieved by the - warmup + main-loop index scheduling). - - If ``prelaunch_gather`` is provided, it is passed to the first - chunk's generator to skip re-launching the already in-flight - A2A gather. - """ - - # Momentum is already applied by _step_muon before this method. - - ordered_params, param_to_state, rank, chunk_size = ( - self._setup_parallel(names, params, group, qk_logits)) - - def pipelines(): - first = True - for start in range(0, len(ordered_params), chunk_size): - chunk = ordered_params[start:start + chunk_size] - if chunk: - kwargs = dict( - params=chunk, - param_to_state=param_to_state, - rank=rank, - ns_steps=group["ns_steps"], - lr=lr, - weight_decay=weight_decay, - none_grad=group["none_grad"], - ) - if first and prelaunch_gather is not None: - kwargs['prelaunch_gather'] = prelaunch_gather - first = False - yield muon_chunk_pipeline(**kwargs) - - with record_function("muon::pipeline"): - run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1) - - def _step_muon(self, group, qk_logits=None): - params = group["params"] - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - # Apply momentum to all params before routing/expansion. - # Batched using _foreach_* ops (compiled, fullgraph=True). - with record_function("muon::momentum"): - active_params = [p for p in params if p.grad is not None] - if active_params: - # Ensure momentum buffers exist (avoid zeros_like when already present). - for p in active_params: - if "momentum_buffer" not in self.state[p]: - self.state[p]["momentum_buffer"] = torch.zeros_like( - p.grad) - - # Extract local tensors for compiled batch function. - local_grads = [ - p.grad._local_tensor - if isinstance(p.grad, DTensor) else p.grad - for p in active_params - ] - local_bufs = [ - self.state[p]["momentum_buffer"]._local_tensor - if isinstance(self.state[p]["momentum_buffer"], DTensor) - else self.state[p]["momentum_buffer"] - for p in active_params - ] - - # Wrap momentum as tensor for torch.compile. - batch_pre_ortho(local_grads, local_bufs, - torch.tensor(momentum), group["nesterov"]) - - # For non-nesterov, the result is the momentum buffer. - if not group["nesterov"]: - for p in active_params: - p.grad = self.state[p]["momentum_buffer"] - - # Identify batched experts for deferred NS. - # Detection is cheap (condition checks only); actual NS compute is - # deferred so it can overlap with the first chunk's A2A gather. - deferred_expert_work = [] - if self.expert_keys: - batched_expert_indices = [] - for i, (n, p) in enumerate(zip(names, params)): - if not (is_expert_param(n, self.expert_keys) - and p.grad is not None): - continue - # Eligible: plain tensor, or DTensor with no non-dim-0 shards. - if isinstance(p.data, DTensor): - has_tp = any( - _is_shard(pl) and pl.dim != 0 for pl in p.placements) - if has_tp: - continue - batched_expert_indices.append(i) - - if batched_expert_indices: - # Save refs for deferred NS; free grads from param list. - for i in batched_expert_indices: - p = params[i] - g = p.grad - local_g = (g._local_tensor - if isinstance(g, DTensor) else g) - local_data = (p.data._local_tensor if isinstance( - p.data, DTensor) else p.data) - deferred_expert_work.append((local_data, local_g)) - p.grad = None - - # Remove batched experts from lists before expansion. - keep = sorted( - set(range(len(params))) - set(batched_expert_indices)) - names = [names[i] for i in keep] - params = [params[i] for i in keep] - - def _run_deferred_expert_ns(): - """Execute deferred batched expert NS.""" - if not deferred_expert_work: - return - with record_function("muon::batched_expert_ns"): - ns_steps = group["ns_steps"] - for local_data, local_g in deferred_expert_work: - u = zeropower_via_newtonschulz5_batched( - local_g.to(COMM_DTYPE), steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:]) - local_data.mul_(1 - lr * weight_decay) - local_data.add_(u, alpha=-adjusted_lr) - - # Expand expert params by splitting on dim 0. - logger.debug("[_step_muon] before expand: %d params, expert_keys=%s", - len(params), self.expert_keys) - if self.expert_keys: - cache_key = tuple(id(p) for p in params) - cache = self._expert_expand_cache.get(cache_key) - - if cache is None: - # Cold path: full expansion + build cache metadata. - exp_names, exp_params = _expand_expert_params( - names, params, self.expert_keys) - - # Build per-expert-group info for hot-path grad updates. - grad_info = [] - exp_idx = 0 - for orig_idx, (n, p) in enumerate(zip(names, params)): - if not is_expert_param(n, self.expert_keys): - exp_idx += 1 - continue - - is_dt = isinstance(p.data, DTensor) - num_experts = (p.to_local() if is_dt else p.data).shape[0] - - # Detect TP mesh from the first expanded expert param. - tp_mesh = None - tp_pls = None - sample = exp_params[exp_idx] - if isinstance(sample.data, DTensor): - tp_mesh = sample.data.device_mesh - tp_pls = list(sample.data.placements) - - grad_info.append((orig_idx, num_experts, exp_idx, is_dt, - tp_mesh, tp_pls)) - exp_idx += num_experts - - self._expert_expand_cache[cache_key] = { - 'names': exp_names, - 'params': exp_params, - 'grad_info': grad_info, - } - names, params = exp_names, exp_params - else: - # Hot path: reuse cached params, only update expert grads. - for (orig_idx, num_experts, exp_start, is_dt, tp_mesh, - tp_pls) in cache['grad_info']: - p = params[orig_idx] - g = p.grad - local_grad = (g.to_local() - if is_dt and isinstance(g, DTensor) else g) - for i in range(num_experts): - expert_p = cache['params'][exp_start + i] - sg = local_grad[i] - if tp_mesh is not None: - expert_p.grad = DTensor.from_local( - sg, device_mesh=tp_mesh, placements=tp_pls) - else: - expert_p.grad = sg - p.grad = None - - names = cache['names'] - params = cache['params'] - else: - names, params = _expand_expert_params(names, params, - self.expert_keys) - logger.debug("[_step_muon] after expand: %d params", len(params)) - - param_dtensors = [] - name_dtensors = [] - - param_tensors = [] - name_tensors = [] - - # distributed_muon is a reference implementation for testing only. - # The parallel pipeline (all2all) path below is the production path. - if self.use_distributed_muon: - _run_deferred_expert_ns() - self.distributed_muon(names=names, - params=params, - group=group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits) - return - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - logger.debug( - "[route] %s → base (DTensor all-Replicate), " - "shape=%s, placements=%s", n, p.shape, p.placements) - param_tensors.append(p) - name_tensors.append(n) - else: - logger.debug( - "[route] %s → parallel (DTensor), shape=%s, " - "placements=%s, mesh=%s", n, p.shape, p.placements, - p.device_mesh.mesh_dim_names) - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - logger.debug("[route] %s → base (plain tensor), shape=%s", n, - p.data.shape) - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError(f"Unsupported parameter type: {type(p.data)}") - - logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, " - f"{len(param_tensors)} Tensors → base") - - def group_dtensors(dtensors, names): - # To support different placements, we group parameters by placements - # and run parallel Muon on each group. - - placement_to_params = defaultdict(lambda: ([], [])) - - assert len(dtensors) == len(names) - for p, n in zip(dtensors, names): - placement_to_params[tuple([p.placements, - p.device_mesh])][0].append(n) - placement_to_params[tuple([p.placements, - p.device_mesh])][1].append(p) - return placement_to_params - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - dtensor_group = group_dtensors(param_dtensors, name_dtensors) - - # Pre-launch the first chunk's A2A gather so that the NCCL - # communication overlaps with the (deferred) batched expert NS - # compute on the default CUDA stream. - prelaunch = None - if deferred_expert_work: - first_names, first_params = next(iter(dtensor_group.values())) - ordered, pts, rnk, csz = self._setup_parallel( - first_names, first_params, group, qk_logits) - first_chunk = ordered[:csz] - if first_chunk: - prelaunch = prelaunch_first_gather(first_chunk, pts, rnk, - group["none_grad"]) - - _run_deferred_expert_ns() - - first_group = True - for _, (names, params) in dtensor_group.items(): - pg = prelaunch if first_group else None - first_group = False - self.parallel( - names, - params, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - prelaunch_gather=pg, - ) - else: - _run_deferred_expert_ns() - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - ) - - def _register_states_for_offload(self): - """Register all optimizer state tensors with the CPU offload pool. - - Called once after the first step when states have been lazily created. - Offloads all param states (momentum buffers for Muon, moment1/moment2 - for AdamW) to free GPU memory between steps. - """ - pool = self._cpu_offload_pool - tracked = 0 - for group in self.param_groups: - for p in group["params"]: - if p not in self.state: - continue - state = self.state[p] - if group.get("use_muon", False): - if "momentum_buffer" in state: - pool.track(state["momentum_buffer"]) - tracked += 1 - else: - if "moment1" in state: - pool.track(state["moment1"]) - if "moment2" in state: - pool.track(state["moment2"]) - tracked += 1 - logger.info("[CPUOffload] Registered %d param states for offload", - tracked) - - @torch.no_grad - def step(self, closure=None, qk_logits=None): - """Perform a single optimization step. - - Args: - closure (Callable, optional): A closure that reevaluates the model - and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). - """ - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # H2D: reload optimizer states from CPU before computation. - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload() - - logger.debug("[Muon.step] expert_keys=%s, %d param groups", - self.expert_keys, len(self.param_groups)) - - for i, group in enumerate(self.param_groups): - if group["use_muon"]: - logger.debug("[Muon.step] group %d: use_muon=True, %d params", - i, len(group["params"])) - self._step_muon(group, qk_logits=qk_logits) - else: - logger.debug( - "[Muon.step] group %d: use_muon=False (AdamW), %d params", - i, len(group["params"])) - step_adamw(self.state, group) - - # D2H: offload optimizer states to CPU after computation. - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - return loss - - # ------------------------------------------------------------------ - # CPU offload public helpers - # ------------------------------------------------------------------ - - def turn_on_cpu_offload(self): - """Enable CPU offload for optimizer states.""" - if self.cpu_offload: - return - logger.info("[Muon] turn_on_cpu_offload") - self.cpu_offload = True - if not self.state: - return - self._cpu_offload_pool = CPUOffloadPool() - self._offload_initialized = False - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_off_cpu_offload(self): - """Disable CPU offload and keep optimizer states resident on GPU.""" - if not self.cpu_offload: - return - logger.info("[Muon] turn_off_cpu_offload") - if self._offload_initialized: - self._cpu_offload_pool.reload() - torch.cuda.current_stream().synchronize() - self._cpu_offload_pool = None - self._offload_initialized = False - self.cpu_offload = False - - # ------------------------------------------------------------------ - # Checkpoint support for cpu_offload - # ------------------------------------------------------------------ - - def state_dict(self) -> dict: - if self.cpu_offload: - raise RuntimeError( - "Muon.state_dict() requires turn_off_cpu_offload() before checkpoint save." - ) - return super().state_dict() - - def load_state_dict(self, state_dict: dict) -> None: - if self.cpu_offload: - raise RuntimeError( - "Muon.load_state_dict() requires turn_off_cpu_offload() before checkpoint load." - ) - super().load_state_dict(state_dict) - - # Invalidate adamw.py's module-level tensor caches so that - # the next step rebuilds them with the newly loaded state tensors. - _placement_cache.clear() - _tensor_cache.clear() diff --git a/build/torch210-cxx11-rocm70-x86_64-linux/newton_schulz.py b/build/torch210-cxx11-rocm70-x86_64-linux/newton_schulz.py deleted file mode 100644 index d939264b69a34e7a3fa78859f34dc265a1159d59..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm70-x86_64-linux/newton_schulz.py +++ /dev/null @@ -1,240 +0,0 @@ -from itertools import repeat -from math import inf, sqrt - -import numpy as np -import torch - -from .matmul_transpose_triton import matmul_transpose_assign - -COMM_DTYPE = torch.bfloat16 -DEFAULT_CHUNK_SIZE_RATIO = 4 - - -def _optimal_quintic(l, u, max_iter=1000): - """ - Use the simplified Remez algorithm to find the optimal odd quintic approximant - to the constant function x -> 1 over the interval [l, u]. - - Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum - approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the - two interior equioscillation nodes q, r until convergence. Returns the - closed-form equioscillating solution when l ≈ u. - - Raises ValueError if any intermediate value (a, b, c, E, q, r) is non-finite - (NaN or inf). Raises RuntimeError if convergence is not reached within - max_iter iterations. - """ - assert 0 <= l <= u - if 1 - 5e-6 <= l / u: - return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5) - q = (3 * l + u) / 4 - r = (l + 3 * u) / 4 - E = inf - for _ in range(max_iter): - old_E = E - LHS = np.array( - [ - [l, l**3, l**5, 1], - [q, q**3, q**5, -1], - [r, r**3, r**5, 1], - [u, u**3, u**5, -1], - ] - ) - a, b, c, E = np.linalg.solve(LHS, np.ones(4)) - if not np.all(np.isfinite([a, b, c, E])): - raise ValueError( - f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}" - ) - q, r = np.sqrt( - (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / (10 * c) - ) - if not np.all(np.isfinite([q, r])): - raise ValueError(f"_optimal_quintic: non-finite node update q={q}, r={r}") - if abs(old_E - E) <= 1e-15: - break - else: - raise RuntimeError( - f"_optimal_quintic: did not converge after {max_iter} iterations" - ) - return float(a), float(b), float(c) - - -def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0): - """ - Compute the Polar Express coefficient series for `num_iters` quintic iterations. - - Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that - compose to map singular values from [l, 1] toward 1. At each step: - 1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion` - prevents near-zero singular values from stalling by raising the effective - lower bound; if it is active (cushion*u > l), the coefficients are - rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u]. - 2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the - last iteration, providing numerical headroom at the cost of a slightly slower - final convergence step. - 3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1). - - Returns a list of (a, b, c) tuples, one per iteration. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - u = 1 - assert 0 <= l <= u - safety_factor = 1 + safety_factor_eps - coefficients = [] - for iter in range(num_iters): - a, b, c = _optimal_quintic(max(l, cushion * u), u) - if cushion * u > l: - pl = a * l + b * l**3 + c * l**5 - pu = a * u + b * u**3 + c * u**5 - rescaler = 2 / (pl + pu) - a *= rescaler - b *= rescaler - c *= rescaler - if iter < num_iters - 1: - a /= safety_factor - b /= safety_factor**3 - c /= safety_factor**5 - coefficients.append((a, b, c)) - l = a * l + b * l**3 + c * l**5 - u = 2 - l - return coefficients - - -# Precomputed Polar Express coefficients (a, b, c) for 10 quintic Newton-Schulz -# iterations. Each tuple is the minimax-optimal (Remez/equioscillation) odd quintic -# approximant to x->1 over the current singular-value interval, computed once at -# import time and reused across all optimizer steps. -# -# Contrast with the former hardcoded NS coefficients (5 fixed tuples): -# - Former: empirically tuned to maximize slope at zero; did not converge -# singular values to 1, yielding US'V^T with S' ~ Uniform(0.5, 1.5) instead -# of the true polar factor UV^T. -# - Polar Express: analytically optimal per step, adapting to the shrinking -# singular-value interval [l, u] as iterations progress; converges all -# singular values to 1, producing the exact polar factor UV^T. -_coeffs_list = _optimal_composition( - l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02 -) - - -# This code is adapted from: -# KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py) -# NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress) -# matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon) -@torch.no_grad() -def _zeropower_via_newtonschulz5(G, steps): - """ - Compute the polar factor of G via the Polar Express method. - - Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c) - are the Polar Express coefficients from `_coeffs_list`. Each step is the - optimal odd quintic approximant to x -> 1 over the current singular-value - interval, minimizing the maximum approximation error (Remez / minimax criterion). - The composition maps singular values from [l, 1] to near 1, producing the - polar factor (orthogonal factor in the polar decomposition G = UP). - - `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2, - cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE - X = G # no manual typecast - - if G.size(0) > G.size(1): - X = X.T - - X = X / (X.norm() + 1e-7) - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - # Perform the NS iterations - for a, b, c in hs: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(0) > G.size(1): - X = X.T - - return X - - -@torch.no_grad() -def _zeropower_via_newtonschulz5_batched(G, steps): - """Batched polar factor computation for 3D (E, out, in) tensors. - - Same algorithm as ``_zeropower_via_newtonschulz5`` but uses - ``torch.bmm`` / ``torch.baddbmm`` instead of the 2D Triton kernel, - processing all E expert matrices in a single batched call. - """ - assert len(G.shape) == 3 - assert G.dtype == COMM_DTYPE - X = G - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - # Per-expert Frobenius norm. - X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) - - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - for a, b, c in hs: - buf1 = torch.bmm(X, X.transpose(-2, -1)) - buf2 = torch.bmm(buf1, buf1.transpose(-2, -1)) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.baddbmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - return X - - -_ns_per_shape: dict[tuple[int, ...], callable] = {} -_use_compile = True - - -def set_ns_compile(enabled: bool): - """Toggle torch.compile for Newton-Schulz iteration.""" - global _use_compile - _use_compile = enabled - - -def zeropower_via_newtonschulz5(G, steps=5): - if not _use_compile: - return _zeropower_via_newtonschulz5(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() - - -def zeropower_via_newtonschulz5_batched(G, steps=5): - """Compile-cached batched Newton-Schulz for 3D expert tensors.""" - if not _use_compile: - return _zeropower_via_newtonschulz5_batched(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile( - _zeropower_via_newtonschulz5_batched, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() diff --git a/build/torch210-cxx11-rocm70-x86_64-linux/optimizer/__init__.py b/build/torch210-cxx11-rocm70-x86_64-linux/optimizer/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm70-x86_64-linux/optimizer/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch210-cxx11-rocm70-x86_64-linux/pipeline.py b/build/torch210-cxx11-rocm70-x86_64-linux/pipeline.py deleted file mode 100644 index c0c2d515856182d8d15ad27dd4e4e093b29397d6..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm70-x86_64-linux/pipeline.py +++ /dev/null @@ -1,468 +0,0 @@ -import logging -from typing import Generator - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -from .core import _muon_state, adjust_lr_for_muon -from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5 -from .qk_clip import compute_scales - -logger = logging.getLogger(__name__) - -# ====================================================================== -# Stage helpers -# ====================================================================== - - -def _launch_gather( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Allocate gather buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_gather``). - gathered_grads: ``{id(p): empty_tensor}`` for owned params, - ``None`` for non-owned. - recv_counts: Per-source-rank element counts. - """ - # Allocate gathered-grad buffers - gathered_grads: dict[int, torch.Tensor | None] = {} - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - gathered_grads[id(p)] = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - else: - gathered_grads[id(p)] = None - - # Build send buffer – batch grad copies via torch.cat - # (1-2 fused kernels vs N individual narrow().copy_() calls). - send_counts = [0] * num_ranks - for p in params: - state = param_to_state[id(p)] - send_counts[state.worker_rank] += state.rank_numels[rank] - - total_send = sum(send_counts) - if total_send > 0: - # Group grad slices by destination rank in a single pass. - dst_to_grads = [[] for _ in range(num_ranks)] - for p in params: - state = param_to_state[id(p)] - n = state.rank_numels[rank] - if n > 0: - g = p.grad.to_local() - dst_to_grads[state.worker_rank].append(g.reshape(-1)) - - # Flatten in dst order and cat once. - all_slices = [] - for dst in range(num_ranks): - all_slices.extend(dst_to_grads[dst]) - send_buf = torch.cat(all_slices) - if send_buf.dtype != COMM_DTYPE: - send_buf = send_buf.to(COMM_DTYPE) - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += state.rank_numels[src] - recv_counts[src] = total - - recv_buf = torch.empty(sum(recv_counts), dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - logger.debug(f"send_buf size: {send_buf.numel()}, " - f"recv_buf size: {recv_buf.numel()}, " - f"recv_counts: {recv_counts}, " - f"send_counts: {send_counts}, " - f"process_group: {str(process_group)}") - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, gathered_grads, recv_counts - - -def _complete_gather( - recv_buf: torch.Tensor, - recv_counts: list[int], - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - param_to_state: dict[int, _muon_state], - rank: int, -) -> None: - """Reconstruct gathered grads from the recv buffer (in-place).""" - off = 0 - for src in range(len(recv_counts)): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - - indices = state.rank_indices[src] - - shard_view = gathered_grads[id(p)][indices] - n = shard_view.numel() - if n == 0: - continue - - sg = recv_buf.narrow(0, off + inner_off, n) - sg = sg.reshape(shard_view.shape) - gathered_grads[id(p)][indices] = sg - - inner_off += n - assert inner_off == block - off += block - - -def _compute_ns( - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - ns_steps: int, -) -> dict[int, torch.Tensor | None]: - """Run Newton-Schulz orthogonalization on owned parameters. - - Returns: - computed_us: ``{id(p): orthogonalized_update}`` for owned params. - """ - computed_us: dict[int, torch.Tensor | None] = {} - for p in owned_params: - u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps) - gathered_grads[id(p)] = None # free gathered grad - computed_us[id(p)] = u - return computed_us - - -def _launch_scatter( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, - computed_us: dict[int, torch.Tensor | None], -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor], list[int]]: - """Allocate scatter buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_scatter``). - scattered_us: Empty dict, populated by ``_complete_scatter`` with - zero-copy views into ``recv_buf``. - recv_counts: Per-source-rank element counts. - """ - # scattered_us is populated by _complete_scatter with zero-copy views - # into recv_buf, avoiding N empty_like allocations + N copy_ calls. - # Pre-seed entries for params whose local shard is empty (rank_numels == 0) - # so _update_params can iterate all params without KeyError. - scattered_us: dict[int, torch.Tensor] = {} - for p in params: - if param_to_state[id(p)].rank_numels[rank] == 0: - scattered_us[id(p)] = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - # Build send buffer – batch via torch.cat - # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls). - send_counts = [0] * num_ranks - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - for dst_rank in range(num_ranks): - send_counts[dst_rank] += state.rank_numels[dst_rank] - - total_send = sum(send_counts) - if total_send > 0: - # Cache u_full conversions to avoid redundant .to() per dst_rank. - u_fulls = {} - for p in owned_params: - u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous() - - # Collect slices in dst order (matches all-to-all send layout). - all_slices = [] - for dst_rank in range(num_ranks): - for p in owned_params: - state = param_to_state[id(p)] - su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten() - if su.numel() > 0: - all_slices.append(su) - - send_buf = torch.cat(all_slices) if all_slices else torch.empty( - 0, dtype=COMM_DTYPE, device="cuda") - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += state.rank_numels[rank] - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, scattered_us, recv_counts - - -def _complete_scatter( - recv_buf: torch.Tensor, - recv_counts: list[int], - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], -) -> None: - """Populate scattered_us with zero-copy views into recv_buf. - - Instead of pre-allocating tensors and copying, we assign views directly - from ``recv_buf``. This eliminates N ``empty_like`` + N ``copy_`` calls. - The underlying storage of ``recv_buf`` is kept alive through the views - until ``scattered_us`` is cleared after ``_update_params``. - """ - off = 0 - for src in range(len(recv_counts)): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = state.rank_numels[rank] - if n == 0: - continue - - scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - - inner_off += n - - assert inner_off == block - off += block - - -def _update_params( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], - lr: float, - weight_decay: float, -) -> None: - """Apply weight decay, Muon update, and optional QK clipping. - - Uses batched ``_foreach_mul_`` for weight decay and batched - ``_foreach_add_`` for the Muon update, grouping parameters by - adjusted_lr to minimize kernel launches while preserving float32 - precision for the alpha scaling. - """ - if not params: - return - - # Batched weight decay: p *= (1 - lr * wd) — single fused kernel. - p_locals = [p._local_tensor for p in params] - torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay) - - # Group params by adjusted_lr so _foreach_add_ can use a single - # alpha per group (preserves float32 precision for alpha scaling). - lr_groups: dict[float, tuple[list, list]] = {} - for p in params: - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - if adjusted_lr not in lr_groups: - lr_groups[adjusted_lr] = ([], []) - lr_groups[adjusted_lr][0].append(p._local_tensor) - lr_groups[adjusted_lr][1].append(scattered_us[id(p)]) - - for adjusted_lr, (p_group, u_group) in lr_groups.items(): - torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr) - - # QK clipping – applied directly on the local tensor to - # avoid DTensor sharding-propagation issues with _StridedShard. - for p in params: - state = param_to_state[id(p)] - if state.qk_clip_state is None: - continue - scales_full = compute_scales(p, state.qk_clip_state) - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = state.rank_indices[rank][0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - -# ====================================================================== -# Pre-launch helper for overlapping first chunk's gather with other work. -# ====================================================================== - - -@torch.no_grad() -def prelaunch_first_gather( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - none_grad: bool, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Launch the first chunk's A2A gather early for overlap with other compute. - - Call this *before* expensive GPU work (e.g. batched expert NS) so that - the NCCL all-to-all runs concurrently on the NCCL stream while the - default stream executes compute. - - Returns the same 4-tuple that ``_launch_gather`` produces, which should - be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - with record_function("muon::prelaunch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - return work, recv_buf, gathered_grads, recv_counts - - -# ====================================================================== -# Main generator – thin orchestrator that wires stages together. -# ====================================================================== - - -@torch.no_grad() -def muon_chunk_pipeline( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - ns_steps: int, - lr: float, - weight_decay: float, - none_grad: bool, - prelaunch_gather: tuple | None = None, -) -> Generator[None, None, None]: - """Process one chunk of parameters through the full Muon pipeline. - - Stages: gather -> compute (Newton-Schulz) -> scatter -> update. - - Each ``yield`` lets :func:`run_pipeline` interleave other chunks so - that communication and computation overlap across chunks. Async - communication is launched via ``async_op=True`` and completed after - the yield with ``work.wait()``. - - Overlap happens because :func:`run_pipeline` admits one new chunk - per iteration (staggered admission). While chunk *N* does NS - compute on the default CUDA stream, chunk *N+1*'s async all-to-all - runs concurrently on the NCCL stream — no separate ``comm_stream`` - is required. - - If ``prelaunch_gather`` is provided, the gather was already launched - by :func:`prelaunch_first_gather` and we skip launching it again. - - Yields exactly **2** times: - - 1. After launching async all-to-all gather (or immediately if pre-launched). - 2. After launching async all-to-all scatter. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - if prelaunch_gather is not None: - # Gather was pre-launched; none_grad already handled by caller. - work, recv_buf, gathered_grads, recv_counts = prelaunch_gather - else: - # Normal path: launch async gather. - with record_function("muon::launch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - yield # --- YIELD 1: other chunks can launch their gather --- - - with record_function("muon::wait_gather"): - work.wait() - _complete_gather(recv_buf, recv_counts, owned_params, gathered_grads, - param_to_state, rank) - del recv_buf - - # Stage 3: Newton-Schulz orthogonalization. - with record_function("muon::newton_schulz"): - computed_us = _compute_ns(owned_params, gathered_grads, ns_steps) - gathered_grads.clear() - - # Stages 4-5: launch async scatter. - with record_function("muon::launch_scatter"): - work, recv_buf, scattered_us, recv_counts = _launch_scatter( - params, owned_params, param_to_state, rank, num_ranks, - process_group, computed_us) - computed_us.clear() - - yield # --- YIELD 2: other chunks can launch their scatter --- - - with record_function("muon::wait_scatter"): - work.wait() - _complete_scatter(recv_buf, recv_counts, params, param_to_state, rank, - scattered_us) - del recv_buf - - # Stage 6: apply parameter updates. - with record_function("muon::update_params"): - _update_params(params, param_to_state, rank, scattered_us, lr, - weight_decay) - scattered_us.clear() diff --git a/build/torch210-cxx11-rocm70-x86_64-linux/qk_clip.py b/build/torch210-cxx11-rocm70-x86_64-linux/qk_clip.py deleted file mode 100644 index 2aba711b3004b7f09e7141da7ef834bd61cc2430..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm70-x86_64-linux/qk_clip.py +++ /dev/null @@ -1,198 +0,0 @@ -import logging -import math -from dataclasses import dataclass - -import torch -from torch.distributed.tensor import DTensor - -from .core import normalize_fqn - -logger = logging.getLogger(__name__) - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - and return (kind, layer_index). - - Supported kinds: - MHA/GQA: 'wq', 'wk', 'q_proj', 'k_proj' - MLA: 'wq_b' (Q up-proj), 'wkv_b' (KV up-proj) - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.1.attn.wq_b.weight' -> ('wq_b', 1) - 'model.0.attn.wkv_b.weight' -> ('wkv_b', 0) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = normalize_fqn(name).split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: str | None # 'wq'/'q_proj'/'wq_b' or 'wk'/'k_proj'/'wkv_b' or None - indices: list[int] # which heads to consider for clipping - head_dim: int # from config (qk_head_dim for MLA wq_b) - threshold: float # from config - logit: torch.Tensor | None - - # MLA-specific fields - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 - - -def get_qk_clip_info(clip_config, n, qk_logits): - """Extract QK clipping info for a named parameter. - - Args: - clip_config: QK clipping configuration dict (or None). - MHA/GQA keys: head_dim, threshold, q_indices, k_indices - MLA extra keys: is_mla=True, qk_nope_head_dim, qk_rope_head_dim, v_head_dim - n: Parameter name string. - qk_logits: Dict mapping layer indices to logit tensors (or None). - - Returns: - QKClipInfo instance with clipping configuration for this parameter. - """ - if clip_config is None: - return None - - head_dim = clip_config.get('head_dim') - threshold = clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - # In TP settings, qk_logits may be DTensor - # We convert it to full tensor here for simplicity - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - -def compute_scales(p, qk_clip_state): - """Compute per-head scaling factors for QK clipping. - - Returns scales tensor (√γ per head) if any head exceeds threshold, else None. - For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim. - """ - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - # Check if any head exceeds threshold before allocating. - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - - if not head_scales: - return None - - # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows - if qk_clip_state.is_mla and kind == 'wkv_b': - effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim - else: - effective_head_dim = head_dim - - H_global = p.shape[0] // effective_head_dim - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale - return scales_full - - -def qk_clip(p, scales, info): - """Apply per-head scaling to a Q/K projection weight matrix. - - Args: - p: Parameter (nn.Parameter or raw tensor). - scales: [n_heads] tensor, each element = √γ_h. - info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions. - - MLA sub-region scaling per Algorithm 1 (MuonClip): - wq_b: q_nope rows → √γ, q_pe rows → γ - wkv_b: k_nope rows → √γ, v rows → unchanged - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not info.is_mla: - # MHA/GQA: uniform √γ applied to all rows in each head - W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: vectorized sub-region scaling within each head - if info.kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_head_dim = qk_nope + info.qk_rope_head_dim - W_3d = W.view(-1, qk_head_dim, W.shape[1]) # [H, qk_head_dim, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # q_nope → √γ - W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1, - 1)) # q_pe → γ - - elif info.kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - W_3d = W.view(-1, kv_stride, W.shape[1]) # [H, kv_stride, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # k_nope → √γ - # v rows: not touched (k_R shared rotary unchanged) diff --git a/build/torch210-cxx11-rocm71-x86_64-linux/_ops.py b/build/torch210-cxx11-rocm71-x86_64-linux/_ops.py deleted file mode 100644 index 034bff088659b5df6f6d401feb18c89dc5f33b29..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm71-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _optimizer_8d53b78_dirty -ops = torch.ops._optimizer_8d53b78_dirty - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_optimizer_8d53b78_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch210-cxx11-rocm71-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so b/build/torch210-cxx11-rocm71-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so deleted file mode 100755 index 482f777c94bf8d03311ce7c0be278ece8264857f..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm71-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:315ff09ffa88ec806cb8abe49edb2ca6951e9ac34be3d3e10f159093f9576ee0 -size 1866112 diff --git a/build/torch210-cxx11-rocm71-x86_64-linux/adamw.py b/build/torch210-cxx11-rocm71-x86_64-linux/adamw.py deleted file mode 100644 index b5a95816a9f5b9e1889eaadae65373bfbced809a..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm71-x86_64-linux/adamw.py +++ /dev/null @@ -1,271 +0,0 @@ -import logging -from collections import defaultdict -from typing import cast - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -def fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: float | torch.Tensor, - weight_decay: float, - eps: float, - maximize: bool, -) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: dict | None = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) - - -def _to_local(t): - """Unwrap DTensor to local tensor for fused ops.""" - return t._local_tensor if isinstance(t, DTensor) else t - - -# --------------------------------------------------------------------------- -# Caches for eliminating per-step Python overhead. -# -# Placement grouping and tensor list assembly are identical every step -# (params don't change placement, moment/step tensors are the same objects -# after initialisation). We cache them keyed by id() of the param list -# stored in param_groups (stable across steps). -# -# Only gradients change each step and must be collected fresh. -# --------------------------------------------------------------------------- - -# id(group["params"]) → dict[placement_key, list[param]] -_placement_cache: dict[int, dict[tuple, list]] = {} - -# id(placement_group_list) → (params_local, moment1, moment2, state_steps) -_tensor_cache: dict[int, tuple[list, list, list, list]] = {} - - -def _step_adamw_params_slow(optimizer_state, params, group): - """Uncached fallback for the rare case where some params lack grads.""" - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - g = p.grad - if g is None: - continue - state = optimizer_state[p] - params_with_grads.append(_to_local(p)) - grads.append(_to_local(g)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - if not params_with_grads: - return - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw_params(optimizer_state, params, group): - """Run fused AdamW on a list of parameters sharing the same placement. - - After the first call, cached tensor lists (params_local, moment1, - moment2, state_steps) are reused — only gradients are collected fresh. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - params: List of parameters to update. - group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay. - """ - # Collect grads — the only thing that changes each step. - with record_function("adamw::collect_grads"): - grads = [] - for p in params: - g = p.grad - if g is None: - # Rare: fall back to slow path that filters per-param. - _step_adamw_params_slow(optimizer_state, params, group) - return - grads.append(_to_local(g)) - - tensor_key = id(params) - if tensor_key not in _tensor_cache: - with record_function("adamw::init_tensor_cache"): - params_local = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - state = optimizer_state[p] - params_local.append(_to_local(p)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(p.grad) - state["moment2"] = torch.zeros_like(p.grad) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - _tensor_cache[tensor_key] = (params_local, moment1, moment2, - state_steps) - - params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key] - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - with record_function("adamw::fused_adamw"): - fused_adamw( - params_local, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw(optimizer_state, group): - """Dispatch AdamW step, grouping parameters by type and placement. - - Placement grouping is cached after the first call since params never - change their placement between steps. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - group: Parameter group dict. - """ - params = group["params"] - placement_key = id(params) - - if placement_key not in _placement_cache: - with record_function("adamw::group_by_placement"): - placement_to_params: dict[tuple, - list[torch.Tensor]] = defaultdict(list) - for p in params: - match p: - case DTensor(): - logger.debug( - "[AdamW] DTensor param: shape=%s, placements=%s, " - "mesh=%s, grad=%s", p.shape, p.placements, - p.device_mesh.mesh_dim_names, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple( - [p.placements, p.device_mesh])].append(p) - case torch.Tensor(): - logger.debug( - "[AdamW] plain param: shape=%s, grad=%s", p.shape, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple([torch.Tensor, - None])].append(p) - - logger.debug("[AdamW] %d placement groups, %d total params", - len(placement_to_params), len(params)) - - _placement_cache[placement_key] = dict(placement_to_params) - - for group_params in _placement_cache[placement_key].values(): - step_adamw_params(optimizer_state, group_params, group) diff --git a/build/torch210-cxx11-rocm71-x86_64-linux/async_utils.py b/build/torch210-cxx11-rocm71-x86_64-linux/async_utils.py deleted file mode 100644 index a45c530ac9cad88e3555ec1047a6aa59f225347e..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm71-x86_64-linux/async_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import logging -from typing import Generator - -logger = logging.getLogger(__name__) - - -class _Task: - """Internal: wraps a generator, advances one yield at a time.""" - - def __init__(self, generator: Generator[None, None, None], index: int): - self._generator = generator - self._index = index - self._steps_completed = 0 - self.step() # run to first yield - - def step(self) -> bool: - try: - next(self._generator) - self._steps_completed += 1 - logger.debug("pipeline[%d] completed stage %d", self._index, - self._steps_completed) - return True - except StopIteration: - logger.debug("pipeline[%d] finished after %d stages", self._index, - self._steps_completed) - return False - - def close(self): - self._generator.close() - - -def run_pipeline( - pipelines: Generator[Generator[None, None, None], None, None], - max_concurrent: int, -) -> None: - """Run generator-based pipelines with bounded concurrency. - - Each pipeline is a generator that yields at stage boundaries. - The runtime interleaves pipelines so communication and computation - overlap across chunks. - """ - if max_concurrent <= 0: - raise ValueError(f"max_concurrent must be > 0, got {max_concurrent}") - - have_new = True - task_index = 0 - previous_tasks: list[_Task] = [] - - try: - while have_new or previous_tasks: - running_tasks: list[_Task] = [] - - # Admit one new pipeline per iteration (staggered admission). - # Admitting one at a time ensures that while chunk N does NS - # compute on the default stream, chunk N+1's NCCL all-to-all - # runs concurrently on the NCCL stream — creating real - # communication/computation overlap on the GPU. - if have_new and len(previous_tasks) < max_concurrent: - try: - gen = next(pipelines) - task = _Task(gen, task_index) - task_index += 1 - running_tasks.append(task) - except StopIteration: - have_new = False - - # Advance every previously-yielded task by one step. - for task in previous_tasks: - if task.step(): - running_tasks.append(task) - - previous_tasks = running_tasks - except BaseException: - # Clean up all in-flight generators to release GPU resources. - for task in previous_tasks: - task.close() - raise diff --git a/build/torch210-cxx11-rocm71-x86_64-linux/core.py b/build/torch210-cxx11-rocm71-x86_64-linux/core.py deleted file mode 100644 index c69d515afef305ad0ed66374095fa2d2468d99cc..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm71-x86_64-linux/core.py +++ /dev/null @@ -1,219 +0,0 @@ -import logging -import math -from dataclasses import dataclass -from typing import List - -import torch -from torch.distributed import ProcessGroup -from torch.distributed.tensor import DTensor - -# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into -# parameter FQNs. Activation checkpointing similarly inserts -# "_checkpoint_wrapped_module". Strip these so name-based matching (skip_keys, -# expert_keys, QK layer parsing) works regardless of wrapper nesting. -_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) - -logger = logging.getLogger(__name__) - - -def normalize_fqn(name: str) -> str: - """Strip torch.compile / checkpoint wrapper components from a parameter FQN.""" - return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS) - - -@dataclass -class _muon_state: - worker_rank: int - process_group: ProcessGroup - rank_indices: dict[int, tuple] # local_rank -> per-dim indices - rank_numels: dict[int, int] # local_rank -> numel - name: str - qk_clip_state: torch.Tensor | None = None - - -def _batch_momentum( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update (no nesterov).""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - - -def _batch_momentum_nesterov( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update with nesterov correction.""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - nesterov_terms = torch._foreach_mul(momentum_bufs, momentum) - torch._foreach_add_(grads, nesterov_terms) - - -_compiled_momentum: dict[bool, callable] = {} -_use_momentum_compile = True - - -def set_momentum_compile(enabled: bool): - """Toggle torch.compile for batched momentum.""" - global _use_momentum_compile - _use_momentum_compile = enabled - - -def batch_pre_ortho( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, - nesterov: bool, -) -> None: - """Batched momentum update on lists of plain tensors. - - Mirrors dion's ``muon_update_pre_orthogonalize``. - Inputs must be plain CUDA tensors (not DTensor). - Modifies ``momentum_bufs`` and (for nesterov) ``grads`` in-place. - - When compile is enabled, uses separately compiled functions for - nesterov=True/False to avoid graph breaks from the branch. - """ - fn = _batch_momentum_nesterov if nesterov else _batch_momentum - if _use_momentum_compile: - if nesterov not in _compiled_momentum: - _compiled_momentum[nesterov] = torch.compile(fn) - fn = _compiled_momentum[nesterov] - fn(grads, momentum_bufs, momentum) - - -def _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay): - """Weight-decay + update on plain tensors. - - Not compiled: per-param @torch.compile caused ~0.25ms TorchDynamo cache - lookup per call × 256+ params = massive overhead. The pipeline path uses - batched _foreach_* ops instead; this function remains for base() and - distributed_muon(). - """ - p_data.mul_(1 - lr * weight_decay) - p_data.add_(u_data, alpha=-adjusted_lr) - - -def update_p(p, u, lr, adjusted_lr, weight_decay): - """Apply weight decay and orthogonalized update to parameter. - - Args: - p: Parameter (torch.nn.Parameter or DTensor). - u: Orthogonalized update tensor. - lr: Base learning rate. - adjusted_lr: Size-adjusted learning rate. - weight_decay: Weight decay coefficient. - """ - # Unwrap Parameter -> underlying data tensor. - p_data = p.data if isinstance(p, torch.nn.Parameter) else p - # Unwrap DTensor -> local CUDA tensor for compiled kernel. - if isinstance(p_data, DTensor): - p_data = p_data._local_tensor - u_data = u._local_tensor if isinstance(u, DTensor) else u - _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay) - - -def adjust_lr_for_muon(lr, param_shape): - """Scale learning rate based on parameter matrix dimensions. - - Args: - lr: Base learning rate. - param_shape: Shape of the parameter tensor. - - Returns: - Adjusted learning rate. - """ - A, B = param_shape[:2] - # We adjust the learning rate and weight decay based on the size of the parameter matrix - # as described in the paper - adjusted_ratio = 0.2 * math.sqrt(max(A, B)) - adjusted_lr = lr * adjusted_ratio - return adjusted_lr - - -def _match_key(parts, key): - """Check if key matches as contiguous components in parts. - - Single-component keys (e.g. "experts") match any single component. - Multi-component keys (e.g. "experts.w1") match as a contiguous subsequence. - """ - key_parts = key.split(".") - key_len = len(key_parts) - if key_len == 1: - return key in parts - return any(parts[i:i + key_len] == key_parts - for i in range(len(parts) - key_len + 1)) - - -def is_expert_param(name, expert_keys): - """Check if a parameter name matches any expert key (component-level).""" - if not expert_keys: - return False - parts = normalize_fqn(name).split(".") - return any(_match_key(parts, key) for key in expert_keys) - - -def default_is_muon(name, x, expert_keys=None): - normalized = normalize_fqn(name) - parts = normalized.split(".") - skip_keys = [ - "embed_tokens", - "lm_head", - "tok_embeddings", - "output", - "mhc_attn", - "mhc_ffn", - "lambda_proj", - ] - if any(key in parts for key in skip_keys): - logger.info( - "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d", - normalized, name, x.ndim) - return False - effective_ndim = x.ndim - is_expert = is_expert_param(name, expert_keys) - if is_expert: - effective_ndim -= 1 - result = effective_ndim >= 2 - logger.info( - "[is_muon] %s (orig: %s): ndim=%d, expert=%s, effective_ndim=%d → %s", - normalized, name, x.ndim, is_expert, effective_ndim, - "Muon" if result else "AdamW") - return result - - -def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None): - if is_muon_func is None: - is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys) - - muon_params, muon_names = [], [] - non_muon_params, non_muon_names = [], [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - non_muon_names.append(n) - - logger.info("[param_groups] expert_keys=%s, Muon=%d, AdamW=%d", - expert_keys, len(muon_names), len(non_muon_names)) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] diff --git a/build/torch210-cxx11-rocm71-x86_64-linux/cpu_offload.py b/build/torch210-cxx11-rocm71-x86_64-linux/cpu_offload.py deleted file mode 100644 index fb5e69154a1d4a6c884491413a37a9acf0f66c80..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm71-x86_64-linux/cpu_offload.py +++ /dev/null @@ -1,206 +0,0 @@ -"""CPU offloading for optimizer states. - -Manages a pinned CPU memory pool and async CUDA streams to offload -optimizer state tensors (momentum buffers, Adam moments) to CPU between -optimizer steps, freeing GPU memory. - -All tracked tensors are packed into a single flat pinned CPU buffer -(per dtype). D2H and H2D copies are performed per-tensor directly -between individual GPU tensors and their slice of the CPU flat buffer -— no GPU staging buffer is allocated, so there is **no temporary GPU -memory spike** during offload or reload. - -Individual tensor storages are freed after offload via -``untyped_storage().resize_(0)``, preserving tensor identity so -downstream caches remain valid. -""" - -import logging -from collections import defaultdict - -import torch -from torch.distributed.tensor import DTensor - -logger = logging.getLogger(__name__) - - -class CPUOffloadPool: - """Pinned CPU memory pool for async optimizer state offloading. - - Tracked tensors are grouped by dtype. Each group gets a single flat - pinned CPU buffer. D2H / H2D copies are per-tensor (into slices of - the flat buffer) to avoid allocating a GPU staging buffer. - """ - - def __init__(self): - self._managed: list[torch.Tensor] = [] - self._storage_nbytes: dict[int, int] = {} # id(t) → bytes - - # Per-dtype group: populated on first offload. - # dtype → dict with keys: - # "indices" : list[int] managed-list indices - # "offsets" : list[tuple[int,int]] (start, numel) in flat buf - # "total" : int total numel - # "cpu_flat" : Tensor pinned CPU buffer - self._groups: dict[torch.dtype, dict] = {} - - self._offload_stream: torch.cuda.Stream | None = None - self._device: torch.device | None = None - self._initialized: bool = False - self._logged: bool = False - - # ------------------------------------------------------------------ - @staticmethod - def _local(t: torch.Tensor) -> torch.Tensor: - """Unwrap DTensor to its local CUDA tensor.""" - return t._local_tensor if isinstance(t, DTensor) else t - - def _ensure_stream(self): - if self._offload_stream is None: - self._offload_stream = torch.cuda.Stream(device=self._device) - - # ------------------------------------------------------------------ - def track(self, tensor: torch.Tensor): - """Register a GPU tensor for CPU offloading. Idempotent.""" - tid = id(tensor) - if tid in self._storage_nbytes: - return - local = self._local(tensor) - if self._device is None: - self._device = local.device - storage = local.untyped_storage() - # Skip tensors with empty storage (e.g. empty FSDP shards) - if storage.size() == 0: - return - self._storage_nbytes[tid] = storage.size() - self._managed.append(tensor) - - # ------------------------------------------------------------------ - def _init_buffers(self): - """Build per-dtype flat buffers on first offload.""" - # Group managed tensors by dtype. - dtype_map: dict[torch.dtype, list[tuple[int, int]]] = defaultdict(list) - for idx, t in enumerate(self._managed): - local = self._local(t) - dtype_map[local.dtype].append((idx, local.numel())) - - total_cpu_bytes = 0 - for dtype, entries in dtype_map.items(): - offsets: list[tuple[int, int]] = [] - indices: list[int] = [] - off = 0 - for idx, n in entries: - indices.append(idx) - offsets.append((off, n)) - off += n - cpu_flat = torch.empty(off, dtype=dtype, device="cpu", pin_memory=True) - self._groups[dtype] = { - "indices": indices, - "offsets": offsets, - "total": off, - "cpu_flat": cpu_flat, - } - total_cpu_bytes += off * cpu_flat.element_size() - - self._initialized = True - logger.info( - "[CPUOffload] Pool initialized: %d tensors, %d dtype group(s), " - "%.2f MB pinned CPU memory", - len(self._managed), - len(self._groups), - total_cpu_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def offload(self): - """Per-tensor async D2H into CPU flat buffer, then free GPU storage.""" - if not self._managed: - return - if not self._initialized: - self._init_buffers() - self._ensure_stream() - - # Offload stream waits for compute to finish. - compute_event = torch.cuda.current_stream(self._device).record_event() - self._offload_stream.wait_event(compute_event) - - offloaded_bytes = 0 - - # Per-tensor D2H copies directly into CPU flat buffer slices. - # No GPU staging buffer → no temporary GPU memory spike. - with torch.cuda.stream(self._offload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - cpu_flat[off : off + n].copy_(local.reshape(-1), non_blocking=True) - - offloaded_bytes += grp["total"] * cpu_flat.element_size() - - # Wait for all D2H copies to land, then free GPU storage. - self._offload_stream.synchronize() - for t in self._managed: - storage = self._local(t).untyped_storage() - if storage.size() != 0: - storage.resize_(0) - else: - raise RuntimeError( - f"Tensor storage is already freed (size=0) before offload. " - f"This indicates a double-free or external interference. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - - if not self._logged: - logger.info( - "[CPUOffload] Offloaded %.2f MB (GPU → CPU)", - offloaded_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def reload(self): - """Per-tensor H2D from CPU flat buffer on the default stream. - - Runs on the current (default) CUDA stream to avoid stream - interaction issues with the parallel Muon pipeline. Since - pinned CPU memory is the source, the copies overlap with - GPU idle time between steps. - """ - if not self._managed or not self._initialized: - return - - reloaded_bytes = 0 - - # Re-allocate all GPU storages first. - for t in self._managed: - local = self._local(t) - storage = local.untyped_storage() - if storage.size() != 0: - raise RuntimeError( - f"Storage should have been freed (size=0) before reload, " - f"but got size={storage.size()}. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - storage.resize_(self._storage_nbytes[id(t)]) - - # Per-tensor H2D copies from CPU flat buffer slices. - # non_blocking=True with pinned source allows DMA overlap. - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off : off + n], non_blocking=True) - - reloaded_bytes += grp["total"] * cpu_flat.element_size() - - if not self._logged: - logger.info( - "[CPUOffload] Reloaded %.2f MB (CPU → GPU)", reloaded_bytes / (1024**2) - ) diff --git a/build/torch210-cxx11-rocm71-x86_64-linux/distributed/utils.py b/build/torch210-cxx11-rocm71-x86_64-linux/distributed/utils.py deleted file mode 100644 index 683096bcd50489e7a09959753fdbc29ba0f3b6a6..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm71-x86_64-linux/distributed/utils.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import (Placement, Shard, - _StridedShard) - - -def _is_shard(placement: Placement) -> bool: - """Check if a placement is a shard type (Shard or _StridedShard). - - In PyTorch 2.10+, _StridedShard no longer inherits from Shard, so - ``placement.is_shard()`` returns False for _StridedShard. This helper - handles both old and new hierarchies. - """ - return isinstance(placement, (Shard, _StridedShard)) - - -def get_slices_of_dtensor( - target: DTensor | torch.Tensor, - local_rank: int, - shard_mesh: DeviceMesh, - shard_placements: tuple[Placement], -) -> tuple[slice | torch.Tensor, ...]: - """ - Get per-dimension indices for a given rank's shard of the target tensor. - - Uses ``Shard.local_shard_size_and_offset`` and - ``_StridedShard.local_shard_size_and_offset`` for correct handling of - both contiguous and strided (non-contiguous) sharding. - - Args: - target (DTensor | torch.Tensor): The target tensor (for its shape). - local_rank (int): The local rank within the shard group. - shard_mesh (DeviceMesh): The shard mesh (only shard dimensions). - shard_placements (tuple[Placement]): The shard placements. - - Returns: - A tuple of indices (one per tensor dim). Each element is either: - - A ``slice`` (for contiguous or unsharded dims) - - A 1-D ``torch.LongTensor`` of indices (for strided sharding) - """ - - # find the global rank of the local rank in the shard mesh - rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank] - - rank_coords = (shard_mesh.mesh == rank).nonzero() - - assert len(rank_coords) == 1 - rank_coords = tuple(rank_coords[0].tolist()) - - assert len(rank_coords) == len(shard_placements) - - # Track per-shard-dim indices. - # None means "not yet sharded on this dim". - dim_indices: dict[int, torch.Tensor] = {} - - # Caution: Assuming replicate-to-shard of the shard mesh goes with - # left-to-right sharding. This is ensured by the sorting logic of - # construct_shard_mesh function. - for mesh_dim_idx, (rank_coord, placement) in enumerate( - zip(rank_coords, shard_placements)): - assert _is_shard(placement) - - num_chunks = shard_mesh.mesh.shape[mesh_dim_idx] - shard_dim = placement.dim - - # Current effective size on this dim (may already be sub-sharded) - if shard_dim in dim_indices: - curr_size = len(dim_indices[shard_dim]) - else: - curr_size = target.size()[shard_dim] - - # Compute indices for this level of sharding - if isinstance(placement, _StridedShard): - _shard_size, offsets = _StridedShard.local_shard_size_and_offset( - placement, - curr_size, - num_chunks, - rank_coord, - return_first_offset=False) - new_indices = torch.tensor(offsets, dtype=torch.long) - else: - shard_size, offset = Shard.local_shard_size_and_offset( - curr_size, num_chunks, rank_coord) - new_indices = torch.arange(offset, - offset + shard_size, - dtype=torch.long) - - # Compose with previous indices on this dim - if shard_dim in dim_indices: - dim_indices[shard_dim] = dim_indices[shard_dim][new_indices] - else: - dim_indices[shard_dim] = new_indices - - # Build result tuple - result: list[slice | torch.Tensor] = [] - for d in range(len(target.size())): - if d not in dim_indices: - result.append(slice(None)) - else: - indices = dim_indices[d] - # Convert contiguous indices to slice for efficiency - if len(indices) > 0: - start = indices[0].item() - expected = torch.arange(start, - start + len(indices), - dtype=torch.long) - if torch.equal(indices, expected): - result.append(slice(start, start + len(indices))) - else: - result.append(indices) - else: - result.append(slice(0, 0)) - - return tuple(result) - - -_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, - ProcessGroup]] = dict() - - -def construct_shard_mesh( - placements: tuple[Placement], - mesh: DeviceMesh, -) -> tuple[DeviceMesh, ProcessGroup, tuple[Placement, ...]]: - """Construct shard sub-mesh and ProcessGroup for all-to-all communication. - - Given a DTensor's placements and device mesh, extracts the "shard group" - — the set of ranks that together hold all shards of the same replica — - and creates a ProcessGroup for all-to-all among them. - - Steps: - 1. Sort placements: Replicate first, then Shard by (dim, granularity). - 2. Permute the mesh tensor to match the sorted order. - 3. Collapse Replicate dims → list of shard sub-meshes (one per replica). - 4. Create/retrieve a cached ProcessGroup for the current rank's sub-mesh. - - Example — 8 GPUs, mesh shape (2, 2, 2), - placements ``[Shard(0), Replicate, _StridedShard(0)]``:: - - Step 1 — Sort: [Replicate, _StridedShard(0), Shard(0)] - Permutation: [1, 2, 0] - - Step 2 — Permute mesh dims by [1, 2, 0]: - Original: Permuted: - [[[0,1],[2,3]], [[[0,2],[1,3]], - [[4,5],[6,7]]] [[4,6],[5,7]]] - - Step 3 — Unbind replicate dim (dim 0), giving 2 shard sub-meshes: - sub-mesh 0 = [[0,2],[1,3]] (replica group 0) - sub-mesh 1 = [[4,6],[5,7]] (replica group 1) - shard_placements = (_StridedShard(0), Shard(0)) - - Step 4 — Rank 0 → ProcessGroup([0,1,4,5]) - Rank 2 → ProcessGroup([2,3,6,7]) - - Returns: - ``(shard_mesh, process_group, shard_placements)`` - """ - my_rank = dist.get_rank() - assert mesh.mesh.device.type == 'cpu' - - # -- Fast path: 1D all-shard mesh → reuse existing PG. ---------------- - # Reuses the mesh's existing ProcessGroup directly, avoiding the - # overhead of dist.new_group(). The standard path below also handles - # subset calls safely via use_local_synchronization=True, but this - # fast path is still beneficial for the common 1D shard case. - if mesh.ndim == 1 and len(placements) == 1 and _is_shard(placements[0]): - key = (*mesh.mesh.shape, *mesh.mesh.flatten().tolist()) - if key not in _ranks_to_dist_cache: - _ranks_to_dist_cache[key] = (mesh, mesh.get_group()) - return (*_ranks_to_dist_cache[key], tuple(placements)) - - mesh_tensor = mesh.mesh.clone() - - # -- Step 1: Sort placements (Replicate first, then Shard by dim). ------ - # _StridedShard comes BEFORE regular Shard on the same dim so that - # get_slices_of_dtensor applies the outer sharding first, matching - # DTensor's left-to-right (outer-to-inner) composition order. - def _sort_key(item): - index, placement = item - assert not placement.is_partial(), "Partial placement not supported" - if placement.is_replicate(): - return (-1, 0, index) - assert _is_shard(placement), f"Unsupported: {type(placement)}" - split = (-1 / placement.split_factor if isinstance( - placement, _StridedShard) else 0) - return (placement.dim, split, index) - - indexed = sorted(enumerate(placements), key=_sort_key) - perm, sorted_placements = zip(*indexed) - - # -- Step 2: Permute mesh to match sorted placement order. -------------- - sorted_mesh = mesh_tensor.permute(perm) - - # -- Step 3: Collapse replicate dims → list of shard sub-meshes. -------- - # E.g. mesh (2, 3, 4, 4) with [R, R, S(0), S(1)] → 6 sub-meshes of (4, 4) - num_rep = sum(1 for p in sorted_placements if p.is_replicate()) - if num_rep > 0: - if num_rep > 1: - sorted_mesh = sorted_mesh.flatten(0, num_rep - 1) - shard_meshes = list(torch.unbind(sorted_mesh, dim=0)) - else: - shard_meshes = [sorted_mesh] - shard_placements = sorted_placements[num_rep:] - assert len(shard_placements) == len(set(shard_placements)) - - # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. -- - # Each rank only creates the group it belongs to, using - # use_local_synchronization=True so that only group members need to - # coordinate. This avoids deadlocks when different PP stages call - # construct_shard_mesh for different parameters. - def _cache_key(t: torch.Tensor) -> tuple: - return (*t.shape, *t.flatten().tolist()) - - my_key = None - for sm in shard_meshes: - if (my_rank == sm).any().item(): - key = _cache_key(sm) - assert my_key is None, "Rank appears in multiple shard groups" - my_key = key - if key not in _ranks_to_dist_cache: - pg = dist.new_group(sm.flatten().tolist(), - use_local_synchronization=True) - _ranks_to_dist_cache[key] = ( - DeviceMesh(device_type="cuda", mesh=sm), - pg, - ) - - return (*_ranks_to_dist_cache[my_key], shard_placements) diff --git a/build/torch210-cxx11-rocm71-x86_64-linux/matmul_transpose_triton.py b/build/torch210-cxx11-rocm71-x86_64-linux/matmul_transpose_triton.py deleted file mode 100644 index 792de23d82c3fb45fe33d397ab9b76a0787259d0..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm71-x86_64-linux/matmul_transpose_triton.py +++ /dev/null @@ -1,122 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], - restore_value=['y'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -@torch.library.custom_op("muon::matmul_transpose_assign", - mutates_args=("d_out", )) -def matmul_transpose_assign(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """Compute d_out = d_in @ d_in.T using an optimized Triton kernel.""" - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -@matmul_transpose_assign.register_fake -def _(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """FakeTensor impl: d_out is already allocated, mutation is declared.""" - pass diff --git a/build/torch210-cxx11-rocm71-x86_64-linux/metadata.json b/build/torch210-cxx11-rocm71-x86_64-linux/metadata.json deleted file mode 100644 index c55a35717622f1dd5c8ba376ea3a814cbcc10d78..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm71-x86_64-linux/metadata.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch210-cxx11-rocm71-x86_64-linux/muon.py b/build/torch210-cxx11-rocm71-x86_64-linux/muon.py deleted file mode 100644 index 14c0e22471fa6d47a51ed95e0e0c341dc18d5194..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm71-x86_64-linux/muon.py +++ /dev/null @@ -1,1068 +0,0 @@ -import logging -import types -from collections import defaultdict -from typing import Any - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Replicate, Shard -from torch.profiler import record_function - -from .adamw import _placement_cache, _tensor_cache, step_adamw -from .async_utils import run_pipeline -from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho, - get_default_muon_param_groups, is_expert_param, update_p) -from .cpu_offload import CPUOffloadPool -from .distributed.utils import (_is_shard, construct_shard_mesh, - get_slices_of_dtensor) -from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO, - _zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5_batched) -from .pipeline import muon_chunk_pipeline, prelaunch_first_gather -from .qk_clip import compute_scales, get_qk_clip_info, qk_clip - -logger = logging.getLogger(__name__) - - -def _expand_expert_params(names, params, expert_keys): - """Expand expert params by splitting on dim 0 (expert dimension). - - Params whose name matches any key in ``expert_keys`` are treated as - expert-parallel tensors. Their outermost dimension is the expert - dimension: an ``(E, out, in)`` tensor becomes ``E`` separate 2D - ``nn.Parameter`` views so that in-place updates propagate back to - the original storage. - - Non-expert params with ``ndim > 2`` trigger an ``AssertionError`` — - if they are expert params, their key must be added to ``expert_keys``. - - The grad must already be set on each expert param (e.g. after momentum). - - For DTensor expert params, placements that shard on dim 0 (expert dim) - are consumed by the split. Non-dim-0 shard placements (e.g. TP) are - preserved: each 2D slice is wrapped as a DTensor on the corresponding - submesh so the parallel pipeline handles the TP communication. - """ - expanded_names = [] - expanded_params = [] - - for n, p in zip(names, params): - is_expert = is_expert_param(n, expert_keys) - is_dtensor = isinstance(p.data, DTensor) - - if is_expert: - if is_dtensor: - logger.debug( - "[expand_expert] %s: expert DTensor, shape=%s, " - "placements=%s, mesh=%s, local_shape=%s", n, p.shape, - p.placements, p.device_mesh.mesh_dim_names, - p.to_local().shape) - else: - logger.debug( - "[expand_expert] %s: expert plain tensor, shape=%s", n, - p.data.shape) - - if not is_expert: - assert p.data.ndim <= 2, ( - f"Param {n} has ndim={p.data.ndim} but does not match " - f"expert_keys={expert_keys}. If this is an expert param, " - f"add its key to expert_keys.") - expanded_names.append(n) - expanded_params.append(p) - continue - - g = p.grad - assert g is not None, ( - f"Expert param {n} must have grad set before expansion") - - tp_mesh = None - tp_placements_2d = None - - if is_dtensor: - local_data = p.to_local() - local_grad = g.to_local() if isinstance(g, DTensor) else g - - # Find non-dim-0 shard placements (e.g. TP sharding). - # After splitting on dim 0, Shard(k) becomes Shard(k-1). - tp_dim_indices = [] - tp_placements_2d = [] - for i, pl in enumerate(p.placements): - if _is_shard(pl) and pl.dim != 0: - tp_dim_indices.append(i) - tp_placements_2d.append(Shard(pl.dim - 1)) - - if tp_dim_indices: - tp_dim_names = tuple(p.device_mesh.mesh_dim_names[i] - for i in tp_dim_indices) - if len(tp_dim_names) == 1: - tp_mesh = p.device_mesh[tp_dim_names[0]] - else: - tp_mesh = p.device_mesh[tp_dim_names] - else: - local_data = p.data - local_grad = g - - # Expand: split dim 0, reshape each slice to 2D. - num_local_experts = local_data.shape[0] - for i in range(num_local_experts): - slice_data = local_data[i] - slice_grad = local_grad[i] - - if tp_mesh is not None: - # Wrap as DTensor on TP submesh so the pipeline handles - # TP communication (gather/scatter across TP ranks). - dt_data = DTensor.from_local(slice_data, - device_mesh=tp_mesh, - placements=tp_placements_2d) - dt_grad = DTensor.from_local(slice_grad, - device_mesh=tp_mesh, - placements=tp_placements_2d) - expert_param = torch.nn.Parameter(dt_data, requires_grad=False) - expert_param.grad = dt_grad - else: - expert_param = torch.nn.Parameter(slice_data, - requires_grad=False) - expert_param.grad = slice_grad - - expanded_names.append(f"{n}[{i}]") - expanded_params.append(expert_param) - - p.grad = None # allow expert grad storage to be freed after pipeline - - return expanded_names, expanded_params - - -class Muon(torch.optim.Optimizer): - """ - Muon - MomentUm Orthogonalized by Newton-schulz - - Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- - processing step, in which each 2D parameter's update is replaced with the nearest orthogonal - matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has - the advantage that it can be stably run in bfloat16 on the GPU. - - Some warnings: - - We believe this optimizer is unlikely to work well for training with small batch size. - - We believe it may not work well for finetuning pretrained models, but we haven't tested this. - - Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. - lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) - momentum: The momentum used by the internal SGD. (0.95 is a good default) - nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) - ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. - Parameters that are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW instead. - adamw_lr: The learning rate for the internal AdamW. - adamw_betas: The betas for the internal AdamW. - adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - warmup_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher warmup_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. - chunk_size : Batch size of parameters to process in each - all2all gather/compute/scatter step. - Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified. - use_distributed_muon: Use distributed muon by Liu et al. (2024). - For testing purpose only. - expert_keys: List of strings to identify expert-parallel parameters. - If any key appears in a parameter's name, its outermost - dimension is treated as the expert dimension and expanded - into per-expert 2D params for Muon. For example, - ``expert_keys=["experts"]`` matches any param whose name - contains "experts". 3D+ params not matched by any key - will raise an error. - """ - - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config=None, - warmup_step=5, - chunk_size=-1, - use_distributed_muon=False, - expert_keys=None): - defaults = dict( - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - nesterov=nesterov, - ns_steps=ns_steps, - adamw_betas=adamw_betas, - adamw_eps=adamw_eps, - none_grad=none_grad, - use_muon=True, - ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) - super().__init__(params, defaults) - - self.debug = debug - self.clip_config = clip_config if clip_config is not None else { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100, - } - self.warmup_step = warmup_step - self.chunk_size = chunk_size - self.use_distributed_muon = use_distributed_muon - self.expert_keys = expert_keys - self.cpu_offload = False - self._cpu_offload_pool: CPUOffloadPool | None = None - self._offload_initialized = False - self._parallel_cache: dict[tuple[str, ...], dict] = {} - self._expert_expand_cache: dict[tuple[int, ...], dict] = {} - - def _calc_flops(self, G, steps): - assert len(G.shape) == 2 - M, N = G.shape - if M > N: - M, N = N, M - - return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) - - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - p.placements, p.device_mesh) - - return shard_mesh, shard_pg, shard_placements - - def init_state_and_assign_params(self, names, params, group, qk_logits): - param_to_state = {} - param_to_flops = {} - - total_flops = 0 - for p in params: - g = p.grad - if g is None: - continue - assert g.ndim == 2, "Muon only supports 2D parameters." - - flops = self._calc_flops(g, group["ns_steps"]) - param_to_flops[id(p)] = flops - total_flops += flops - - if self.debug: - logger.debug("Total TFLOPs for Muon: %.2f TFLOPs", - total_flops / 1e12) - - paired = list(zip(names, params)) - - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) - - round_robin = 0 - mesh = ordered_params[0].device_mesh - placements = ordered_params[0].placements - - shard_mesh, shard_pg, shard_placements = self.get_shard_mesh( - ordered_params[0]) - shard_mesh_flattened = shard_mesh.mesh.flatten() - num_ranks = dist.get_world_size(group=shard_pg) - - for n, p in zip(ordered_names, ordered_params): - if mesh != p.device_mesh: - raise ValueError("All parameters must be on the same mesh.") - if placements != p.placements: - raise ValueError("All parameters must have same placements.") - - worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks - round_robin = (round_robin + 1) % len(shard_mesh_flattened) - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - # Precompute per-rank indices and numels for all-to-all. - rank_indices: dict[int, tuple] = {} - rank_numels: dict[int, int] = {} - for r in range(num_ranks): - indices = get_slices_of_dtensor(p, r, shard_mesh, - shard_placements) - rank_indices[r] = indices - numel = 1 - for idx, dim_size in zip(indices, p.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) - rank_numels[r] = numel - - param_to_state[id(p)] = _muon_state( - worker_rank=worker_rank, - process_group=shard_pg, - rank_indices=rank_indices, - rank_numels=rank_numels, - name=n, - qk_clip_state=qk_clip_state, - ) - - return param_to_state, ordered_params - - def base(self, names, params, group, lr, weight_decay, qk_logits): - # Momentum is already applied by _step_muon before this method. - for n, p in zip(names, params): - g = p.grad - if g is None: - continue - - u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) - - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - def distributed_muon( - self, - names: list[str], - params: list[torch.nn.Parameter], - group: dict[str, Any], - lr: float, - weight_decay: float, - qk_logits: list[torch.Tensor | DTensor] | None, - ): - """Batched Distributed Muon — for testing/correctness verification only. - - Uses all-gather to reconstruct full tensors, computes Newton-Schulz on - the full grad, then slices back to local shards. This is simpler but - slower than the parallel pipeline (all2all) path, so it serves as a - reference implementation for verifying correctness. - """ - with record_function("distributed_muon"): - # Momentum is already applied by _step_muon before this method. - ns_steps = group["ns_steps"] - - # Separate plain tensors (no communication) from DTensors. - plain_names, plain_params = [], [] - dtensor_names, dtensor_params = [], [] - for n, p in zip(names, params): - if p.grad is None: - continue - if isinstance(p.data, DTensor): - dtensor_names.append(n) - dtensor_params.append(p) - else: - plain_names.append(n) - plain_params.append(p) - - # Process plain tensors per-param (no communication). - for n, p in zip(plain_names, plain_params): - u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE), - steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, - qk_logits) - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - if not dtensor_params: - return - - # Group DTensors by (placements, mesh) for batched all-gather. - placement_groups: dict[tuple, - tuple[list, - list]] = defaultdict(lambda: ([], [])) - for n, p in zip(dtensor_names, dtensor_params): - key = (p.placements, p.device_mesh) - placement_groups[key][0].append(n) - placement_groups[key][1].append(p) - - logger.info( - "distributed_muon: %d placement groups, %d total dtensors", - len(placement_groups), len(dtensor_params)) - - for (placements, mesh), (grp_names, - grp_params) in placement_groups.items(): - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - placements, mesh) - rank = dist.get_rank(shard_pg) - world_size = dist.get_world_size(shard_pg) - - logger.info(" group: %d params, placements=%s, world_size=%d", - len(grp_params), placements, world_size) - - # Separate params that can be batched (all shard dims evenly - # divisible) from those needing per-param full_tensor - # (e.g. MoE gate weights with fewer rows than shard ranks). - # all_gather_into_tensor requires equal buffer sizes across - # ranks, so uneven splits must use DTensor full_tensor(). - batch_names, batch_params = [], [] - single_names, single_params = [], [] - for n, p in zip(grp_names, grp_params): - even = all(p.shape[pl.dim] % - shard_mesh.mesh.shape[dim_idx] == 0 - for dim_idx, pl in enumerate(shard_placements)) - if even: - batch_names.append(n) - batch_params.append(p) - else: - single_names.append(n) - single_params.append(p) - - # Process uneven-split params per-param via full_tensor(). - for n, p in zip(single_names, single_params): - with record_function("distributed_muon::newton_schulz"): - g_full = p.grad.full_tensor().to(COMM_DTYPE) - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - if not batch_params: - continue - - logger.info(" batched=%d, single=%d", len(batch_params), - len(single_params)) - - # Concat all local grad shards into a single flat buffer. - with record_function("distributed_muon::gather"): - grad_locals = [ - p.grad.to_local().to(COMM_DTYPE).flatten() - for p in batch_params - ] - numels = [g.numel() for g in grad_locals] - grad_concat = torch.cat(grad_locals) - del grad_locals - - # Single all-gather (replaces N separate full_tensor). - grad_gathered = torch.empty( - grad_concat.numel() * world_size, - dtype=COMM_DTYPE, - device="cuda", - ) - dist.all_gather_into_tensor(grad_gathered, - grad_concat, - group=shard_pg) - - total_numel = grad_concat.numel() - del grad_concat - - # Precompute per-param offsets within the concat buffer. - offsets = [] - off = 0 - for ne in numels: - offsets.append(off) - off += ne - - # Per-param: reconstruct full grad → NS → local update. - for i, (n, p) in enumerate(zip(batch_names, batch_params)): - with record_function("distributed_muon::newton_schulz"): - g_full = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - for r in range(world_size): - r_start = r * total_numel + offsets[i] - shard = grad_gathered[r_start:r_start + numels[i]] - indices = get_slices_of_dtensor( - p, r, shard_mesh, shard_placements) - g_full[indices] = shard.reshape( - g_full[indices].shape) - - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - def _setup_parallel(self, names, params, group, qk_logits): - """Compute (or retrieve cached) parallel pipeline metadata. - - Returns: - (ordered_params, param_to_state, rank, chunk_size) - """ - cache_key = tuple(names) - - if cache_key not in self._parallel_cache: - # First call: compute metadata and populate cache. - param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - shard_pg = param_to_state[id(ordered_params[0])].process_group - rank = dist.get_rank(group=shard_pg) - - if self.chunk_size == -1: - shard_ranks = dist.get_world_size(shard_pg) - chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO - elif self.chunk_size > 0: - chunk_size = self.chunk_size - else: - raise ValueError( - "chunk_size must be -1 or a positive integer.") - - ordered_names = [ - param_to_state[id(p)].name for p in ordered_params - ] - name_to_state = { - param_to_state[id(p)].name: param_to_state[id(p)] - for p in ordered_params - } - self._parallel_cache[cache_key] = { - 'ordered_names': ordered_names, - 'name_to_state': name_to_state, - 'rank': rank, - 'chunk_size': chunk_size, - } - else: - # Cached path: rebuild param_to_state with current id(p) keys. - cache = self._parallel_cache[cache_key] - rank = cache['rank'] - chunk_size = cache['chunk_size'] - - name_to_param = dict(zip(names, params)) - ordered_params = [name_to_param[n] for n in cache['ordered_names']] - - param_to_state = {} - for p, n in zip(ordered_params, cache['ordered_names']): - cached_state = cache['name_to_state'][n] - param_to_state[id(p)] = _muon_state( - worker_rank=cached_state.worker_rank, - process_group=cached_state.process_group, - rank_indices=cached_state.rank_indices, - rank_numels=cached_state.rank_numels, - name=n, - qk_clip_state=get_qk_clip_info(self.clip_config, n, - qk_logits), - ) - - return ordered_params, param_to_state, rank, chunk_size - - def parallel(self, - names, - params, - group, - lr, - weight_decay, - qk_logits, - prelaunch_gather=None): - """ - Perform a parallel optimization step using Muon. - - Parameters are chunked and each chunk is processed by a - :func:`muon_chunk_pipeline` generator. :func:`run_pipeline` - interleaves multiple chunks so that communication and computation - overlap across chunks (the same overlap previously achieved by the - warmup + main-loop index scheduling). - - If ``prelaunch_gather`` is provided, it is passed to the first - chunk's generator to skip re-launching the already in-flight - A2A gather. - """ - - # Momentum is already applied by _step_muon before this method. - - ordered_params, param_to_state, rank, chunk_size = ( - self._setup_parallel(names, params, group, qk_logits)) - - def pipelines(): - first = True - for start in range(0, len(ordered_params), chunk_size): - chunk = ordered_params[start:start + chunk_size] - if chunk: - kwargs = dict( - params=chunk, - param_to_state=param_to_state, - rank=rank, - ns_steps=group["ns_steps"], - lr=lr, - weight_decay=weight_decay, - none_grad=group["none_grad"], - ) - if first and prelaunch_gather is not None: - kwargs['prelaunch_gather'] = prelaunch_gather - first = False - yield muon_chunk_pipeline(**kwargs) - - with record_function("muon::pipeline"): - run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1) - - def _step_muon(self, group, qk_logits=None): - params = group["params"] - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - # Apply momentum to all params before routing/expansion. - # Batched using _foreach_* ops (compiled, fullgraph=True). - with record_function("muon::momentum"): - active_params = [p for p in params if p.grad is not None] - if active_params: - # Ensure momentum buffers exist (avoid zeros_like when already present). - for p in active_params: - if "momentum_buffer" not in self.state[p]: - self.state[p]["momentum_buffer"] = torch.zeros_like( - p.grad) - - # Extract local tensors for compiled batch function. - local_grads = [ - p.grad._local_tensor - if isinstance(p.grad, DTensor) else p.grad - for p in active_params - ] - local_bufs = [ - self.state[p]["momentum_buffer"]._local_tensor - if isinstance(self.state[p]["momentum_buffer"], DTensor) - else self.state[p]["momentum_buffer"] - for p in active_params - ] - - # Wrap momentum as tensor for torch.compile. - batch_pre_ortho(local_grads, local_bufs, - torch.tensor(momentum), group["nesterov"]) - - # For non-nesterov, the result is the momentum buffer. - if not group["nesterov"]: - for p in active_params: - p.grad = self.state[p]["momentum_buffer"] - - # Identify batched experts for deferred NS. - # Detection is cheap (condition checks only); actual NS compute is - # deferred so it can overlap with the first chunk's A2A gather. - deferred_expert_work = [] - if self.expert_keys: - batched_expert_indices = [] - for i, (n, p) in enumerate(zip(names, params)): - if not (is_expert_param(n, self.expert_keys) - and p.grad is not None): - continue - # Eligible: plain tensor, or DTensor with no non-dim-0 shards. - if isinstance(p.data, DTensor): - has_tp = any( - _is_shard(pl) and pl.dim != 0 for pl in p.placements) - if has_tp: - continue - batched_expert_indices.append(i) - - if batched_expert_indices: - # Save refs for deferred NS; free grads from param list. - for i in batched_expert_indices: - p = params[i] - g = p.grad - local_g = (g._local_tensor - if isinstance(g, DTensor) else g) - local_data = (p.data._local_tensor if isinstance( - p.data, DTensor) else p.data) - deferred_expert_work.append((local_data, local_g)) - p.grad = None - - # Remove batched experts from lists before expansion. - keep = sorted( - set(range(len(params))) - set(batched_expert_indices)) - names = [names[i] for i in keep] - params = [params[i] for i in keep] - - def _run_deferred_expert_ns(): - """Execute deferred batched expert NS.""" - if not deferred_expert_work: - return - with record_function("muon::batched_expert_ns"): - ns_steps = group["ns_steps"] - for local_data, local_g in deferred_expert_work: - u = zeropower_via_newtonschulz5_batched( - local_g.to(COMM_DTYPE), steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:]) - local_data.mul_(1 - lr * weight_decay) - local_data.add_(u, alpha=-adjusted_lr) - - # Expand expert params by splitting on dim 0. - logger.debug("[_step_muon] before expand: %d params, expert_keys=%s", - len(params), self.expert_keys) - if self.expert_keys: - cache_key = tuple(id(p) for p in params) - cache = self._expert_expand_cache.get(cache_key) - - if cache is None: - # Cold path: full expansion + build cache metadata. - exp_names, exp_params = _expand_expert_params( - names, params, self.expert_keys) - - # Build per-expert-group info for hot-path grad updates. - grad_info = [] - exp_idx = 0 - for orig_idx, (n, p) in enumerate(zip(names, params)): - if not is_expert_param(n, self.expert_keys): - exp_idx += 1 - continue - - is_dt = isinstance(p.data, DTensor) - num_experts = (p.to_local() if is_dt else p.data).shape[0] - - # Detect TP mesh from the first expanded expert param. - tp_mesh = None - tp_pls = None - sample = exp_params[exp_idx] - if isinstance(sample.data, DTensor): - tp_mesh = sample.data.device_mesh - tp_pls = list(sample.data.placements) - - grad_info.append((orig_idx, num_experts, exp_idx, is_dt, - tp_mesh, tp_pls)) - exp_idx += num_experts - - self._expert_expand_cache[cache_key] = { - 'names': exp_names, - 'params': exp_params, - 'grad_info': grad_info, - } - names, params = exp_names, exp_params - else: - # Hot path: reuse cached params, only update expert grads. - for (orig_idx, num_experts, exp_start, is_dt, tp_mesh, - tp_pls) in cache['grad_info']: - p = params[orig_idx] - g = p.grad - local_grad = (g.to_local() - if is_dt and isinstance(g, DTensor) else g) - for i in range(num_experts): - expert_p = cache['params'][exp_start + i] - sg = local_grad[i] - if tp_mesh is not None: - expert_p.grad = DTensor.from_local( - sg, device_mesh=tp_mesh, placements=tp_pls) - else: - expert_p.grad = sg - p.grad = None - - names = cache['names'] - params = cache['params'] - else: - names, params = _expand_expert_params(names, params, - self.expert_keys) - logger.debug("[_step_muon] after expand: %d params", len(params)) - - param_dtensors = [] - name_dtensors = [] - - param_tensors = [] - name_tensors = [] - - # distributed_muon is a reference implementation for testing only. - # The parallel pipeline (all2all) path below is the production path. - if self.use_distributed_muon: - _run_deferred_expert_ns() - self.distributed_muon(names=names, - params=params, - group=group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits) - return - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - logger.debug( - "[route] %s → base (DTensor all-Replicate), " - "shape=%s, placements=%s", n, p.shape, p.placements) - param_tensors.append(p) - name_tensors.append(n) - else: - logger.debug( - "[route] %s → parallel (DTensor), shape=%s, " - "placements=%s, mesh=%s", n, p.shape, p.placements, - p.device_mesh.mesh_dim_names) - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - logger.debug("[route] %s → base (plain tensor), shape=%s", n, - p.data.shape) - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError(f"Unsupported parameter type: {type(p.data)}") - - logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, " - f"{len(param_tensors)} Tensors → base") - - def group_dtensors(dtensors, names): - # To support different placements, we group parameters by placements - # and run parallel Muon on each group. - - placement_to_params = defaultdict(lambda: ([], [])) - - assert len(dtensors) == len(names) - for p, n in zip(dtensors, names): - placement_to_params[tuple([p.placements, - p.device_mesh])][0].append(n) - placement_to_params[tuple([p.placements, - p.device_mesh])][1].append(p) - return placement_to_params - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - dtensor_group = group_dtensors(param_dtensors, name_dtensors) - - # Pre-launch the first chunk's A2A gather so that the NCCL - # communication overlaps with the (deferred) batched expert NS - # compute on the default CUDA stream. - prelaunch = None - if deferred_expert_work: - first_names, first_params = next(iter(dtensor_group.values())) - ordered, pts, rnk, csz = self._setup_parallel( - first_names, first_params, group, qk_logits) - first_chunk = ordered[:csz] - if first_chunk: - prelaunch = prelaunch_first_gather(first_chunk, pts, rnk, - group["none_grad"]) - - _run_deferred_expert_ns() - - first_group = True - for _, (names, params) in dtensor_group.items(): - pg = prelaunch if first_group else None - first_group = False - self.parallel( - names, - params, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - prelaunch_gather=pg, - ) - else: - _run_deferred_expert_ns() - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - ) - - def _register_states_for_offload(self): - """Register all optimizer state tensors with the CPU offload pool. - - Called once after the first step when states have been lazily created. - Offloads all param states (momentum buffers for Muon, moment1/moment2 - for AdamW) to free GPU memory between steps. - """ - pool = self._cpu_offload_pool - tracked = 0 - for group in self.param_groups: - for p in group["params"]: - if p not in self.state: - continue - state = self.state[p] - if group.get("use_muon", False): - if "momentum_buffer" in state: - pool.track(state["momentum_buffer"]) - tracked += 1 - else: - if "moment1" in state: - pool.track(state["moment1"]) - if "moment2" in state: - pool.track(state["moment2"]) - tracked += 1 - logger.info("[CPUOffload] Registered %d param states for offload", - tracked) - - @torch.no_grad - def step(self, closure=None, qk_logits=None): - """Perform a single optimization step. - - Args: - closure (Callable, optional): A closure that reevaluates the model - and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). - """ - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # H2D: reload optimizer states from CPU before computation. - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload() - - logger.debug("[Muon.step] expert_keys=%s, %d param groups", - self.expert_keys, len(self.param_groups)) - - for i, group in enumerate(self.param_groups): - if group["use_muon"]: - logger.debug("[Muon.step] group %d: use_muon=True, %d params", - i, len(group["params"])) - self._step_muon(group, qk_logits=qk_logits) - else: - logger.debug( - "[Muon.step] group %d: use_muon=False (AdamW), %d params", - i, len(group["params"])) - step_adamw(self.state, group) - - # D2H: offload optimizer states to CPU after computation. - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - return loss - - # ------------------------------------------------------------------ - # CPU offload public helpers - # ------------------------------------------------------------------ - - def turn_on_cpu_offload(self): - """Enable CPU offload for optimizer states.""" - if self.cpu_offload: - return - logger.info("[Muon] turn_on_cpu_offload") - self.cpu_offload = True - if not self.state: - return - self._cpu_offload_pool = CPUOffloadPool() - self._offload_initialized = False - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_off_cpu_offload(self): - """Disable CPU offload and keep optimizer states resident on GPU.""" - if not self.cpu_offload: - return - logger.info("[Muon] turn_off_cpu_offload") - if self._offload_initialized: - self._cpu_offload_pool.reload() - torch.cuda.current_stream().synchronize() - self._cpu_offload_pool = None - self._offload_initialized = False - self.cpu_offload = False - - # ------------------------------------------------------------------ - # Checkpoint support for cpu_offload - # ------------------------------------------------------------------ - - def state_dict(self) -> dict: - if self.cpu_offload: - raise RuntimeError( - "Muon.state_dict() requires turn_off_cpu_offload() before checkpoint save." - ) - return super().state_dict() - - def load_state_dict(self, state_dict: dict) -> None: - if self.cpu_offload: - raise RuntimeError( - "Muon.load_state_dict() requires turn_off_cpu_offload() before checkpoint load." - ) - super().load_state_dict(state_dict) - - # Invalidate adamw.py's module-level tensor caches so that - # the next step rebuilds them with the newly loaded state tensors. - _placement_cache.clear() - _tensor_cache.clear() diff --git a/build/torch210-cxx11-rocm71-x86_64-linux/newton_schulz.py b/build/torch210-cxx11-rocm71-x86_64-linux/newton_schulz.py deleted file mode 100644 index d939264b69a34e7a3fa78859f34dc265a1159d59..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm71-x86_64-linux/newton_schulz.py +++ /dev/null @@ -1,240 +0,0 @@ -from itertools import repeat -from math import inf, sqrt - -import numpy as np -import torch - -from .matmul_transpose_triton import matmul_transpose_assign - -COMM_DTYPE = torch.bfloat16 -DEFAULT_CHUNK_SIZE_RATIO = 4 - - -def _optimal_quintic(l, u, max_iter=1000): - """ - Use the simplified Remez algorithm to find the optimal odd quintic approximant - to the constant function x -> 1 over the interval [l, u]. - - Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum - approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the - two interior equioscillation nodes q, r until convergence. Returns the - closed-form equioscillating solution when l ≈ u. - - Raises ValueError if any intermediate value (a, b, c, E, q, r) is non-finite - (NaN or inf). Raises RuntimeError if convergence is not reached within - max_iter iterations. - """ - assert 0 <= l <= u - if 1 - 5e-6 <= l / u: - return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5) - q = (3 * l + u) / 4 - r = (l + 3 * u) / 4 - E = inf - for _ in range(max_iter): - old_E = E - LHS = np.array( - [ - [l, l**3, l**5, 1], - [q, q**3, q**5, -1], - [r, r**3, r**5, 1], - [u, u**3, u**5, -1], - ] - ) - a, b, c, E = np.linalg.solve(LHS, np.ones(4)) - if not np.all(np.isfinite([a, b, c, E])): - raise ValueError( - f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}" - ) - q, r = np.sqrt( - (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / (10 * c) - ) - if not np.all(np.isfinite([q, r])): - raise ValueError(f"_optimal_quintic: non-finite node update q={q}, r={r}") - if abs(old_E - E) <= 1e-15: - break - else: - raise RuntimeError( - f"_optimal_quintic: did not converge after {max_iter} iterations" - ) - return float(a), float(b), float(c) - - -def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0): - """ - Compute the Polar Express coefficient series for `num_iters` quintic iterations. - - Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that - compose to map singular values from [l, 1] toward 1. At each step: - 1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion` - prevents near-zero singular values from stalling by raising the effective - lower bound; if it is active (cushion*u > l), the coefficients are - rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u]. - 2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the - last iteration, providing numerical headroom at the cost of a slightly slower - final convergence step. - 3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1). - - Returns a list of (a, b, c) tuples, one per iteration. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - u = 1 - assert 0 <= l <= u - safety_factor = 1 + safety_factor_eps - coefficients = [] - for iter in range(num_iters): - a, b, c = _optimal_quintic(max(l, cushion * u), u) - if cushion * u > l: - pl = a * l + b * l**3 + c * l**5 - pu = a * u + b * u**3 + c * u**5 - rescaler = 2 / (pl + pu) - a *= rescaler - b *= rescaler - c *= rescaler - if iter < num_iters - 1: - a /= safety_factor - b /= safety_factor**3 - c /= safety_factor**5 - coefficients.append((a, b, c)) - l = a * l + b * l**3 + c * l**5 - u = 2 - l - return coefficients - - -# Precomputed Polar Express coefficients (a, b, c) for 10 quintic Newton-Schulz -# iterations. Each tuple is the minimax-optimal (Remez/equioscillation) odd quintic -# approximant to x->1 over the current singular-value interval, computed once at -# import time and reused across all optimizer steps. -# -# Contrast with the former hardcoded NS coefficients (5 fixed tuples): -# - Former: empirically tuned to maximize slope at zero; did not converge -# singular values to 1, yielding US'V^T with S' ~ Uniform(0.5, 1.5) instead -# of the true polar factor UV^T. -# - Polar Express: analytically optimal per step, adapting to the shrinking -# singular-value interval [l, u] as iterations progress; converges all -# singular values to 1, producing the exact polar factor UV^T. -_coeffs_list = _optimal_composition( - l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02 -) - - -# This code is adapted from: -# KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py) -# NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress) -# matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon) -@torch.no_grad() -def _zeropower_via_newtonschulz5(G, steps): - """ - Compute the polar factor of G via the Polar Express method. - - Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c) - are the Polar Express coefficients from `_coeffs_list`. Each step is the - optimal odd quintic approximant to x -> 1 over the current singular-value - interval, minimizing the maximum approximation error (Remez / minimax criterion). - The composition maps singular values from [l, 1] to near 1, producing the - polar factor (orthogonal factor in the polar decomposition G = UP). - - `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2, - cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE - X = G # no manual typecast - - if G.size(0) > G.size(1): - X = X.T - - X = X / (X.norm() + 1e-7) - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - # Perform the NS iterations - for a, b, c in hs: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(0) > G.size(1): - X = X.T - - return X - - -@torch.no_grad() -def _zeropower_via_newtonschulz5_batched(G, steps): - """Batched polar factor computation for 3D (E, out, in) tensors. - - Same algorithm as ``_zeropower_via_newtonschulz5`` but uses - ``torch.bmm`` / ``torch.baddbmm`` instead of the 2D Triton kernel, - processing all E expert matrices in a single batched call. - """ - assert len(G.shape) == 3 - assert G.dtype == COMM_DTYPE - X = G - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - # Per-expert Frobenius norm. - X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) - - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - for a, b, c in hs: - buf1 = torch.bmm(X, X.transpose(-2, -1)) - buf2 = torch.bmm(buf1, buf1.transpose(-2, -1)) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.baddbmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - return X - - -_ns_per_shape: dict[tuple[int, ...], callable] = {} -_use_compile = True - - -def set_ns_compile(enabled: bool): - """Toggle torch.compile for Newton-Schulz iteration.""" - global _use_compile - _use_compile = enabled - - -def zeropower_via_newtonschulz5(G, steps=5): - if not _use_compile: - return _zeropower_via_newtonschulz5(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() - - -def zeropower_via_newtonschulz5_batched(G, steps=5): - """Compile-cached batched Newton-Schulz for 3D expert tensors.""" - if not _use_compile: - return _zeropower_via_newtonschulz5_batched(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile( - _zeropower_via_newtonschulz5_batched, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() diff --git a/build/torch210-cxx11-rocm71-x86_64-linux/optimizer/__init__.py b/build/torch210-cxx11-rocm71-x86_64-linux/optimizer/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm71-x86_64-linux/optimizer/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch210-cxx11-rocm71-x86_64-linux/pipeline.py b/build/torch210-cxx11-rocm71-x86_64-linux/pipeline.py deleted file mode 100644 index c0c2d515856182d8d15ad27dd4e4e093b29397d6..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm71-x86_64-linux/pipeline.py +++ /dev/null @@ -1,468 +0,0 @@ -import logging -from typing import Generator - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -from .core import _muon_state, adjust_lr_for_muon -from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5 -from .qk_clip import compute_scales - -logger = logging.getLogger(__name__) - -# ====================================================================== -# Stage helpers -# ====================================================================== - - -def _launch_gather( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Allocate gather buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_gather``). - gathered_grads: ``{id(p): empty_tensor}`` for owned params, - ``None`` for non-owned. - recv_counts: Per-source-rank element counts. - """ - # Allocate gathered-grad buffers - gathered_grads: dict[int, torch.Tensor | None] = {} - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - gathered_grads[id(p)] = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - else: - gathered_grads[id(p)] = None - - # Build send buffer – batch grad copies via torch.cat - # (1-2 fused kernels vs N individual narrow().copy_() calls). - send_counts = [0] * num_ranks - for p in params: - state = param_to_state[id(p)] - send_counts[state.worker_rank] += state.rank_numels[rank] - - total_send = sum(send_counts) - if total_send > 0: - # Group grad slices by destination rank in a single pass. - dst_to_grads = [[] for _ in range(num_ranks)] - for p in params: - state = param_to_state[id(p)] - n = state.rank_numels[rank] - if n > 0: - g = p.grad.to_local() - dst_to_grads[state.worker_rank].append(g.reshape(-1)) - - # Flatten in dst order and cat once. - all_slices = [] - for dst in range(num_ranks): - all_slices.extend(dst_to_grads[dst]) - send_buf = torch.cat(all_slices) - if send_buf.dtype != COMM_DTYPE: - send_buf = send_buf.to(COMM_DTYPE) - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += state.rank_numels[src] - recv_counts[src] = total - - recv_buf = torch.empty(sum(recv_counts), dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - logger.debug(f"send_buf size: {send_buf.numel()}, " - f"recv_buf size: {recv_buf.numel()}, " - f"recv_counts: {recv_counts}, " - f"send_counts: {send_counts}, " - f"process_group: {str(process_group)}") - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, gathered_grads, recv_counts - - -def _complete_gather( - recv_buf: torch.Tensor, - recv_counts: list[int], - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - param_to_state: dict[int, _muon_state], - rank: int, -) -> None: - """Reconstruct gathered grads from the recv buffer (in-place).""" - off = 0 - for src in range(len(recv_counts)): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - - indices = state.rank_indices[src] - - shard_view = gathered_grads[id(p)][indices] - n = shard_view.numel() - if n == 0: - continue - - sg = recv_buf.narrow(0, off + inner_off, n) - sg = sg.reshape(shard_view.shape) - gathered_grads[id(p)][indices] = sg - - inner_off += n - assert inner_off == block - off += block - - -def _compute_ns( - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - ns_steps: int, -) -> dict[int, torch.Tensor | None]: - """Run Newton-Schulz orthogonalization on owned parameters. - - Returns: - computed_us: ``{id(p): orthogonalized_update}`` for owned params. - """ - computed_us: dict[int, torch.Tensor | None] = {} - for p in owned_params: - u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps) - gathered_grads[id(p)] = None # free gathered grad - computed_us[id(p)] = u - return computed_us - - -def _launch_scatter( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, - computed_us: dict[int, torch.Tensor | None], -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor], list[int]]: - """Allocate scatter buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_scatter``). - scattered_us: Empty dict, populated by ``_complete_scatter`` with - zero-copy views into ``recv_buf``. - recv_counts: Per-source-rank element counts. - """ - # scattered_us is populated by _complete_scatter with zero-copy views - # into recv_buf, avoiding N empty_like allocations + N copy_ calls. - # Pre-seed entries for params whose local shard is empty (rank_numels == 0) - # so _update_params can iterate all params without KeyError. - scattered_us: dict[int, torch.Tensor] = {} - for p in params: - if param_to_state[id(p)].rank_numels[rank] == 0: - scattered_us[id(p)] = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - # Build send buffer – batch via torch.cat - # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls). - send_counts = [0] * num_ranks - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - for dst_rank in range(num_ranks): - send_counts[dst_rank] += state.rank_numels[dst_rank] - - total_send = sum(send_counts) - if total_send > 0: - # Cache u_full conversions to avoid redundant .to() per dst_rank. - u_fulls = {} - for p in owned_params: - u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous() - - # Collect slices in dst order (matches all-to-all send layout). - all_slices = [] - for dst_rank in range(num_ranks): - for p in owned_params: - state = param_to_state[id(p)] - su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten() - if su.numel() > 0: - all_slices.append(su) - - send_buf = torch.cat(all_slices) if all_slices else torch.empty( - 0, dtype=COMM_DTYPE, device="cuda") - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += state.rank_numels[rank] - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, scattered_us, recv_counts - - -def _complete_scatter( - recv_buf: torch.Tensor, - recv_counts: list[int], - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], -) -> None: - """Populate scattered_us with zero-copy views into recv_buf. - - Instead of pre-allocating tensors and copying, we assign views directly - from ``recv_buf``. This eliminates N ``empty_like`` + N ``copy_`` calls. - The underlying storage of ``recv_buf`` is kept alive through the views - until ``scattered_us`` is cleared after ``_update_params``. - """ - off = 0 - for src in range(len(recv_counts)): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = state.rank_numels[rank] - if n == 0: - continue - - scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - - inner_off += n - - assert inner_off == block - off += block - - -def _update_params( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], - lr: float, - weight_decay: float, -) -> None: - """Apply weight decay, Muon update, and optional QK clipping. - - Uses batched ``_foreach_mul_`` for weight decay and batched - ``_foreach_add_`` for the Muon update, grouping parameters by - adjusted_lr to minimize kernel launches while preserving float32 - precision for the alpha scaling. - """ - if not params: - return - - # Batched weight decay: p *= (1 - lr * wd) — single fused kernel. - p_locals = [p._local_tensor for p in params] - torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay) - - # Group params by adjusted_lr so _foreach_add_ can use a single - # alpha per group (preserves float32 precision for alpha scaling). - lr_groups: dict[float, tuple[list, list]] = {} - for p in params: - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - if adjusted_lr not in lr_groups: - lr_groups[adjusted_lr] = ([], []) - lr_groups[adjusted_lr][0].append(p._local_tensor) - lr_groups[adjusted_lr][1].append(scattered_us[id(p)]) - - for adjusted_lr, (p_group, u_group) in lr_groups.items(): - torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr) - - # QK clipping – applied directly on the local tensor to - # avoid DTensor sharding-propagation issues with _StridedShard. - for p in params: - state = param_to_state[id(p)] - if state.qk_clip_state is None: - continue - scales_full = compute_scales(p, state.qk_clip_state) - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = state.rank_indices[rank][0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - -# ====================================================================== -# Pre-launch helper for overlapping first chunk's gather with other work. -# ====================================================================== - - -@torch.no_grad() -def prelaunch_first_gather( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - none_grad: bool, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Launch the first chunk's A2A gather early for overlap with other compute. - - Call this *before* expensive GPU work (e.g. batched expert NS) so that - the NCCL all-to-all runs concurrently on the NCCL stream while the - default stream executes compute. - - Returns the same 4-tuple that ``_launch_gather`` produces, which should - be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - with record_function("muon::prelaunch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - return work, recv_buf, gathered_grads, recv_counts - - -# ====================================================================== -# Main generator – thin orchestrator that wires stages together. -# ====================================================================== - - -@torch.no_grad() -def muon_chunk_pipeline( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - ns_steps: int, - lr: float, - weight_decay: float, - none_grad: bool, - prelaunch_gather: tuple | None = None, -) -> Generator[None, None, None]: - """Process one chunk of parameters through the full Muon pipeline. - - Stages: gather -> compute (Newton-Schulz) -> scatter -> update. - - Each ``yield`` lets :func:`run_pipeline` interleave other chunks so - that communication and computation overlap across chunks. Async - communication is launched via ``async_op=True`` and completed after - the yield with ``work.wait()``. - - Overlap happens because :func:`run_pipeline` admits one new chunk - per iteration (staggered admission). While chunk *N* does NS - compute on the default CUDA stream, chunk *N+1*'s async all-to-all - runs concurrently on the NCCL stream — no separate ``comm_stream`` - is required. - - If ``prelaunch_gather`` is provided, the gather was already launched - by :func:`prelaunch_first_gather` and we skip launching it again. - - Yields exactly **2** times: - - 1. After launching async all-to-all gather (or immediately if pre-launched). - 2. After launching async all-to-all scatter. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - if prelaunch_gather is not None: - # Gather was pre-launched; none_grad already handled by caller. - work, recv_buf, gathered_grads, recv_counts = prelaunch_gather - else: - # Normal path: launch async gather. - with record_function("muon::launch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - yield # --- YIELD 1: other chunks can launch their gather --- - - with record_function("muon::wait_gather"): - work.wait() - _complete_gather(recv_buf, recv_counts, owned_params, gathered_grads, - param_to_state, rank) - del recv_buf - - # Stage 3: Newton-Schulz orthogonalization. - with record_function("muon::newton_schulz"): - computed_us = _compute_ns(owned_params, gathered_grads, ns_steps) - gathered_grads.clear() - - # Stages 4-5: launch async scatter. - with record_function("muon::launch_scatter"): - work, recv_buf, scattered_us, recv_counts = _launch_scatter( - params, owned_params, param_to_state, rank, num_ranks, - process_group, computed_us) - computed_us.clear() - - yield # --- YIELD 2: other chunks can launch their scatter --- - - with record_function("muon::wait_scatter"): - work.wait() - _complete_scatter(recv_buf, recv_counts, params, param_to_state, rank, - scattered_us) - del recv_buf - - # Stage 6: apply parameter updates. - with record_function("muon::update_params"): - _update_params(params, param_to_state, rank, scattered_us, lr, - weight_decay) - scattered_us.clear() diff --git a/build/torch210-cxx11-rocm71-x86_64-linux/qk_clip.py b/build/torch210-cxx11-rocm71-x86_64-linux/qk_clip.py deleted file mode 100644 index 2aba711b3004b7f09e7141da7ef834bd61cc2430..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-rocm71-x86_64-linux/qk_clip.py +++ /dev/null @@ -1,198 +0,0 @@ -import logging -import math -from dataclasses import dataclass - -import torch -from torch.distributed.tensor import DTensor - -from .core import normalize_fqn - -logger = logging.getLogger(__name__) - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - and return (kind, layer_index). - - Supported kinds: - MHA/GQA: 'wq', 'wk', 'q_proj', 'k_proj' - MLA: 'wq_b' (Q up-proj), 'wkv_b' (KV up-proj) - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.1.attn.wq_b.weight' -> ('wq_b', 1) - 'model.0.attn.wkv_b.weight' -> ('wkv_b', 0) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = normalize_fqn(name).split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: str | None # 'wq'/'q_proj'/'wq_b' or 'wk'/'k_proj'/'wkv_b' or None - indices: list[int] # which heads to consider for clipping - head_dim: int # from config (qk_head_dim for MLA wq_b) - threshold: float # from config - logit: torch.Tensor | None - - # MLA-specific fields - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 - - -def get_qk_clip_info(clip_config, n, qk_logits): - """Extract QK clipping info for a named parameter. - - Args: - clip_config: QK clipping configuration dict (or None). - MHA/GQA keys: head_dim, threshold, q_indices, k_indices - MLA extra keys: is_mla=True, qk_nope_head_dim, qk_rope_head_dim, v_head_dim - n: Parameter name string. - qk_logits: Dict mapping layer indices to logit tensors (or None). - - Returns: - QKClipInfo instance with clipping configuration for this parameter. - """ - if clip_config is None: - return None - - head_dim = clip_config.get('head_dim') - threshold = clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - # In TP settings, qk_logits may be DTensor - # We convert it to full tensor here for simplicity - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - -def compute_scales(p, qk_clip_state): - """Compute per-head scaling factors for QK clipping. - - Returns scales tensor (√γ per head) if any head exceeds threshold, else None. - For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim. - """ - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - # Check if any head exceeds threshold before allocating. - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - - if not head_scales: - return None - - # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows - if qk_clip_state.is_mla and kind == 'wkv_b': - effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim - else: - effective_head_dim = head_dim - - H_global = p.shape[0] // effective_head_dim - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale - return scales_full - - -def qk_clip(p, scales, info): - """Apply per-head scaling to a Q/K projection weight matrix. - - Args: - p: Parameter (nn.Parameter or raw tensor). - scales: [n_heads] tensor, each element = √γ_h. - info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions. - - MLA sub-region scaling per Algorithm 1 (MuonClip): - wq_b: q_nope rows → √γ, q_pe rows → γ - wkv_b: k_nope rows → √γ, v rows → unchanged - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not info.is_mla: - # MHA/GQA: uniform √γ applied to all rows in each head - W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: vectorized sub-region scaling within each head - if info.kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_head_dim = qk_nope + info.qk_rope_head_dim - W_3d = W.view(-1, qk_head_dim, W.shape[1]) # [H, qk_head_dim, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # q_nope → √γ - W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1, - 1)) # q_pe → γ - - elif info.kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - W_3d = W.view(-1, kv_stride, W.shape[1]) # [H, kv_stride, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # k_nope → √γ - # v rows: not touched (k_R shared rotary unchanged) diff --git a/build/torch210-cxx11-cu126-x86_64-linux/__init__.py b/build/torch26-cxx11-cu118-x86_64-linux/optimizer/__init__.py old mode 100644 new mode 100755 similarity index 100% rename from build/torch210-cxx11-cu126-x86_64-linux/__init__.py rename to build/torch26-cxx11-cu118-x86_64-linux/optimizer/__init__.py diff --git a/build/torch26-cxx11-cu118-x86_64-linux/optimizer/_ops.py b/build/torch26-cxx11-cu118-x86_64-linux/optimizer/_ops.py new file mode 100755 index 0000000000000000000000000000000000000000..7cf68eab4638da3512b5e49541c916ebd12301f0 --- /dev/null +++ b/build/torch26-cxx11-cu118-x86_64-linux/optimizer/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _optimizer_036642a_dirty +ops = torch.ops._optimizer_036642a_dirty + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_optimizer_036642a_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx11-cu118-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so b/build/torch26-cxx11-cu118-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..4df57b1d0e99209ec349328d9fa3a61cea7f97da --- /dev/null +++ b/build/torch26-cxx11-cu118-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9c77e5647b6056bfaee25050cca7948c40859db0a88fa4fcf40b67a85c947d8c +size 1787272 diff --git a/build/torch26-cxx11-cu118-x86_64-linux/optimizer/muon.py b/build/torch26-cxx11-cu118-x86_64-linux/optimizer/muon.py new file mode 100755 index 0000000000000000000000000000000000000000..0d614d55d721efac406c147b4f62e6c703a91107 --- /dev/null +++ b/build/torch26-cxx11-cu118-x86_64-linux/optimizer/muon.py @@ -0,0 +1,455 @@ +import math +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch.distributed._tensor import DTensor + + +# This code snippet is a modified version adapted from the following GitHub repositories: +# https://github.com/KellerJordan/Muon/blob/master/muon.py +@torch.no_grad() +def _zeropower_via_newtonschulz5(G, steps): + """ + Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a + quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose + of minimizing steps, it turns out to be empirically effective to keep increasing the slope at + zero even beyond the point where the iteration no longer converges all the way to one everywhere + on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T + where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model + performance at all relative to UV^T, where USV^T = G is the SVD. + """ + assert len(G.shape) == 2 + a, b, c = (3.4445, -4.7750, 2.0315) + X = G # no manual typecast + if G.size(0) > G.size(1): + X = X.T + # Ensure spectral norm is at most 1 + X = X / (X.norm() + 1e-7) + X = X.bfloat16() + # Perform the NS iterations + for _ in range(steps): + A = X @ X.T + # B = ( + # b * A + c * A @ A + # ) + B = torch.addmm(A, A, A, alpha=c, beta=b) + # X = a * X + B @ X + X = torch.addmm(X, B, X, alpha=1.0, beta=a) + + if G.size(0) > G.size(1): + X = X.T + return X.to(G.dtype) + + +@dataclass +class _muon_state: + # TODO: use Optional + worker_rank: int | None = None + gathered_grad: torch.Tensor | None = None + computed_u: torch.Tensor | None = None + gather_event: torch.cuda.Event | None = None + compute_event: torch.cuda.Event | None = None + + +@torch.no_grad() +def _gather(p, state, rank, comm_stream, none_grad): + g = p.grad + mesh = g.device_mesh + + if rank == state.worker_rank: + gather_list = [torch.empty_like(g.to_local()) for _ in range(mesh.mesh.numel())] + else: + gather_list = None + + with torch.cuda.stream(comm_stream): + torch.distributed.gather( + g.to_local(), + dst=state.worker_rank, + gather_list=gather_list, + group=mesh.get_group(), + ) + if rank == state.worker_rank: + if state.gathered_grad is not None: + raise RuntimeError( + "Gather event already exists, which should not happen." + ) + state.gathered_grad = torch.cat(gather_list, dim=0) + state.gather_event = torch.cuda.Event() + state.gather_event.record() + else: + state.gathered_grad = None + state.gather_event = None + if none_grad: + p.grad = None + + +@torch.no_grad() +def _compute_u(state, steps, rank, compute_stream): + with torch.cuda.stream(compute_stream): + if rank == state.worker_rank: + if state.gather_event is None: + raise RuntimeError("Gather event must be set before compute.") + compute_stream.wait_event(state.gather_event) + u = _zeropower_via_newtonschulz5(state.gathered_grad, steps) + state.computed_u = u + state.compute_event = torch.cuda.Event() + state.compute_event.record() + # Clear the gathered gradient to free memory + state.gathered_grad = None + else: + state.computed_u = None + state.compute_event = None + + +@torch.no_grad() +def _scatter(p, state, lr, wd, rank, comm_stream): + u = state.computed_u + mesh = p.device_mesh + + with torch.cuda.stream(comm_stream): + if rank == state.worker_rank: + if state.compute_event is None: + raise RuntimeError("Compute event must be set before scatter.") + comm_stream.wait_event(state.compute_event) + scatter_list = list(torch.split(u, p.size(0) // mesh.mesh.numel(), dim=0)) + else: + scatter_list = None + + u = torch.empty_like(p.to_local()) + torch.distributed.scatter( + u, + scatter_list=scatter_list, + src=state.worker_rank, + group=mesh.get_group(), + ) + if rank == state.worker_rank: + # Clear u to free memory + state.computed_u = None + u = DTensor.from_local( + u, + placements=p.placements, + device_mesh=mesh, + ) + p.data.mul_(1 - lr * wd) + p.data.add_(u, alpha=-lr) + + +class Muon(torch.optim.Optimizer): + """ + Muon - MomentUm Orthogonalized by Newton-schulz + + Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- + processing step, in which each 2D parameter's update is replaced with the nearest orthogonal + matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has + the advantage that it can be stably run in bfloat16 on the GPU. + + Some warnings: + - We believe this optimizer is unlikely to work well for training with small batch size. + - We believe it may not work well for finetuning pretrained models, but we haven't tested this. + + Arguments: + muon_params: The parameters to be optimized by Muon. + lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) + momentum: The momentum used by the internal SGD. (0.95 is a good default) + nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) + ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) + adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are + {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. + adamw_lr: The learning rate for the internal AdamW. + adamw_betas: The betas for the internal AdamW. + adamw_eps: The epsilon for the internal AdamW. + adamw_wd: The weight decay for the internal AdamW. + """ + + def __init__( + self, + model, + is_muon_func, + lr=1e-3, + momentum=0.95, + nesterov=True, + ns_steps=5, + adamw_wd=0.1, + adamw_betas=(0.9, 0.95), + adamw_eps=1e-8, + none_grad=True, + debug=False, + ): + defaults = dict( + lr=lr, + wd=adamw_wd, + momentum=momentum, + nesterov=nesterov, + ns_steps=ns_steps, + adamw_betas=adamw_betas, + adamw_eps=adamw_eps, + none_grad=none_grad, + ) + + super().__init__(model.parameters(), defaults) + self.is_muon_func = is_muon_func + self.model = model + + if not dist.is_initialized(): + raise RuntimeError( + "Muon optimizer requires distributed training to be initialized." + ) + + self.rank = dist.get_rank() + + self.comm_stream = torch.cuda.Stream() + self.compute_stream = torch.cuda.Stream() + self.debug = debug + + def __setstate__(self, state): + # Sort parameters into those for which we will use Muon, and those for which we will not + super().__setstate__(state) + for name, p in self.model.named_parameters(): + if self.is_muon_func(p, name): + # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer + assert p.ndim == 2, p.ndim + self.state[p]["use_muon"] = True + self.state[p]["orig_shape"] = p.shape + else: + # Do not use Muon for parameters in adamw_params + self.state[p]["use_muon"] = False + + def _calc_flops(self, G, steps): + assert len(G.shape) == 2 + M, N = G.shape + if M > N: + M, N = N, M + + return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) + + def adjust_lr_for_muon(self, lr, param_shape): + A, B = param_shape[:2] + # We adjust the learning rate and weight decay based on the size of the parameter matrix + # as describted in the paper + adjusted_ratio = 0.2 * math.sqrt(max(A, B)) + adjusted_lr = lr * adjusted_ratio + return adjusted_lr + + def init_state_and_assign_params(self, params, group): + param_to_state = {} + param_to_flops = {} + + total_flops = 0 + for p in params: + g = p.grad + if g is None: + continue + assert g.ndim == 2, "Muon only supports 2D parameters." + + flops = self._calc_flops(g, group["ns_steps"]) + param_to_flops[id(p)] = flops + total_flops += flops + + if self.debug: + print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", flush=True) + + ordered_params = sorted( + params, key=lambda p: param_to_flops[id(p)], reverse=True + ) + + round_robin = 0 + mesh = None + for p in ordered_params: + if mesh is None: + mesh = p.device_mesh + if mesh.ndim != 1: + raise NotImplementedError( + "Muon requires a 1D mesh for distributed training yet." + ) + elif mesh != p.device_mesh: + raise ValueError("All parameters must be on the same mesh.") + + param_to_state[id(p)] = _muon_state() + param_to_state[id(p)].worker_rank = mesh.mesh[round_robin].item() + + round_robin = (round_robin + 1) % mesh.mesh.numel() + + return param_to_state, ordered_params + + def base(self, params, group, lr, wd, momentum): + # generate weight updates in distributed fashion + for p in params: + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + assert g is not None + + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + + u = _zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) + + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + + # apply weight decay + p.data.mul_(1 - lr * wd) + + # apply update + p.data.add_(u, alpha=-adjusted_lr) + + def _update_g(self, p, g, group, momentum): + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + return g + + def _update_p(self, p, u, lr, wd): + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + # apply weight decay + p.data.mul_(1 - lr * wd) + # apply update + p.data.add_(u, alpha=-adjusted_lr) + + def parallel(self, params, group, lr, wd, momentum): + """ + Perform a parallel optimization step using Muon. + """ + + for p in params: + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + + # Update g in the local rank + g = self._update_g( + p, + g, + group, + momentum=momentum, + ) + p.grad = g + + param_to_state, ordered_params = self.init_state_and_assign_params( + params, group + ) + + def enqueue_gathers(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _gather(p, state, self.rank, self.comm_stream, group["none_grad"]) + + def enqueue_computes(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _compute_u(state, group["ns_steps"], self.rank, self.compute_stream) + + def enqueue_scatters(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream) + + chunk_size = params[0].device_mesh.mesh.numel() + + # Wait grad update + self.comm_stream.wait_stream(torch.cuda.current_stream()) + + enqueue_gathers(0, chunk_size) + for i in range(0, len(params) + chunk_size - 1, chunk_size): + enqueue_computes(i, chunk_size) + enqueue_gathers(i + chunk_size, chunk_size) + enqueue_scatters(i, chunk_size) + + torch.cuda.current_stream().wait_stream(self.comm_stream) + + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + ############################ + # Muon # + ############################ + + params = [p for p in group["params"] if self.state[p]["use_muon"]] + lr = group["lr"] + wd = group["wd"] + momentum = group["momentum"] + + if isinstance(params[0].data, DTensor): + self.parallel( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) + else: + self.base( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) + + ############################ + # AdamW backup # + ############################ + + params = [p for p in group["params"] if not self.state[p]["use_muon"]] + lr = group["lr"] + beta1, beta2 = group["adamw_betas"] + eps = group["adamw_eps"] + weight_decay = group["wd"] + + for p in params: + g = p.grad + if g is None: + continue + state = self.state[p] + if "step" not in state: + state["step"] = 0 + state["moment1"] = torch.zeros_like(g) + state["moment2"] = torch.zeros_like(g) + state["step"] += 1 + step = state["step"] + buf1 = state["moment1"] + buf2 = state["moment2"] + buf1.lerp_(g, 1 - beta1) + buf2.lerp_(g.square(), 1 - beta2) + + g = buf1 / (eps + buf2.sqrt()) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + scale = bias_correction1 / bias_correction2**0.5 + p.data.mul_(1 - lr * weight_decay) + p.data.add_(g, alpha=-lr / scale) + + return loss diff --git a/build/torch210-cxx11-cu128-x86_64-linux/__init__.py b/build/torch26-cxx11-cu124-x86_64-linux/optimizer/__init__.py old mode 100644 new mode 100755 similarity index 100% rename from build/torch210-cxx11-cu128-x86_64-linux/__init__.py rename to build/torch26-cxx11-cu124-x86_64-linux/optimizer/__init__.py diff --git a/build/torch26-cxx11-cu124-x86_64-linux/optimizer/_ops.py b/build/torch26-cxx11-cu124-x86_64-linux/optimizer/_ops.py new file mode 100755 index 0000000000000000000000000000000000000000..7cf68eab4638da3512b5e49541c916ebd12301f0 --- /dev/null +++ b/build/torch26-cxx11-cu124-x86_64-linux/optimizer/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _optimizer_036642a_dirty +ops = torch.ops._optimizer_036642a_dirty + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_optimizer_036642a_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx11-cu124-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so b/build/torch26-cxx11-cu124-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..078af88758fab815617b6fae432c3ce4d18f6271 --- /dev/null +++ b/build/torch26-cxx11-cu124-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:94ea66089cc8d9eda72b017733a9e05e4fee5a2f04c50658b690d2c19f0d3068 +size 1824224 diff --git a/build/torch26-cxx11-cu124-x86_64-linux/optimizer/muon.py b/build/torch26-cxx11-cu124-x86_64-linux/optimizer/muon.py new file mode 100755 index 0000000000000000000000000000000000000000..0d614d55d721efac406c147b4f62e6c703a91107 --- /dev/null +++ b/build/torch26-cxx11-cu124-x86_64-linux/optimizer/muon.py @@ -0,0 +1,455 @@ +import math +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch.distributed._tensor import DTensor + + +# This code snippet is a modified version adapted from the following GitHub repositories: +# https://github.com/KellerJordan/Muon/blob/master/muon.py +@torch.no_grad() +def _zeropower_via_newtonschulz5(G, steps): + """ + Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a + quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose + of minimizing steps, it turns out to be empirically effective to keep increasing the slope at + zero even beyond the point where the iteration no longer converges all the way to one everywhere + on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T + where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model + performance at all relative to UV^T, where USV^T = G is the SVD. + """ + assert len(G.shape) == 2 + a, b, c = (3.4445, -4.7750, 2.0315) + X = G # no manual typecast + if G.size(0) > G.size(1): + X = X.T + # Ensure spectral norm is at most 1 + X = X / (X.norm() + 1e-7) + X = X.bfloat16() + # Perform the NS iterations + for _ in range(steps): + A = X @ X.T + # B = ( + # b * A + c * A @ A + # ) + B = torch.addmm(A, A, A, alpha=c, beta=b) + # X = a * X + B @ X + X = torch.addmm(X, B, X, alpha=1.0, beta=a) + + if G.size(0) > G.size(1): + X = X.T + return X.to(G.dtype) + + +@dataclass +class _muon_state: + # TODO: use Optional + worker_rank: int | None = None + gathered_grad: torch.Tensor | None = None + computed_u: torch.Tensor | None = None + gather_event: torch.cuda.Event | None = None + compute_event: torch.cuda.Event | None = None + + +@torch.no_grad() +def _gather(p, state, rank, comm_stream, none_grad): + g = p.grad + mesh = g.device_mesh + + if rank == state.worker_rank: + gather_list = [torch.empty_like(g.to_local()) for _ in range(mesh.mesh.numel())] + else: + gather_list = None + + with torch.cuda.stream(comm_stream): + torch.distributed.gather( + g.to_local(), + dst=state.worker_rank, + gather_list=gather_list, + group=mesh.get_group(), + ) + if rank == state.worker_rank: + if state.gathered_grad is not None: + raise RuntimeError( + "Gather event already exists, which should not happen." + ) + state.gathered_grad = torch.cat(gather_list, dim=0) + state.gather_event = torch.cuda.Event() + state.gather_event.record() + else: + state.gathered_grad = None + state.gather_event = None + if none_grad: + p.grad = None + + +@torch.no_grad() +def _compute_u(state, steps, rank, compute_stream): + with torch.cuda.stream(compute_stream): + if rank == state.worker_rank: + if state.gather_event is None: + raise RuntimeError("Gather event must be set before compute.") + compute_stream.wait_event(state.gather_event) + u = _zeropower_via_newtonschulz5(state.gathered_grad, steps) + state.computed_u = u + state.compute_event = torch.cuda.Event() + state.compute_event.record() + # Clear the gathered gradient to free memory + state.gathered_grad = None + else: + state.computed_u = None + state.compute_event = None + + +@torch.no_grad() +def _scatter(p, state, lr, wd, rank, comm_stream): + u = state.computed_u + mesh = p.device_mesh + + with torch.cuda.stream(comm_stream): + if rank == state.worker_rank: + if state.compute_event is None: + raise RuntimeError("Compute event must be set before scatter.") + comm_stream.wait_event(state.compute_event) + scatter_list = list(torch.split(u, p.size(0) // mesh.mesh.numel(), dim=0)) + else: + scatter_list = None + + u = torch.empty_like(p.to_local()) + torch.distributed.scatter( + u, + scatter_list=scatter_list, + src=state.worker_rank, + group=mesh.get_group(), + ) + if rank == state.worker_rank: + # Clear u to free memory + state.computed_u = None + u = DTensor.from_local( + u, + placements=p.placements, + device_mesh=mesh, + ) + p.data.mul_(1 - lr * wd) + p.data.add_(u, alpha=-lr) + + +class Muon(torch.optim.Optimizer): + """ + Muon - MomentUm Orthogonalized by Newton-schulz + + Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- + processing step, in which each 2D parameter's update is replaced with the nearest orthogonal + matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has + the advantage that it can be stably run in bfloat16 on the GPU. + + Some warnings: + - We believe this optimizer is unlikely to work well for training with small batch size. + - We believe it may not work well for finetuning pretrained models, but we haven't tested this. + + Arguments: + muon_params: The parameters to be optimized by Muon. + lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) + momentum: The momentum used by the internal SGD. (0.95 is a good default) + nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) + ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) + adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are + {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. + adamw_lr: The learning rate for the internal AdamW. + adamw_betas: The betas for the internal AdamW. + adamw_eps: The epsilon for the internal AdamW. + adamw_wd: The weight decay for the internal AdamW. + """ + + def __init__( + self, + model, + is_muon_func, + lr=1e-3, + momentum=0.95, + nesterov=True, + ns_steps=5, + adamw_wd=0.1, + adamw_betas=(0.9, 0.95), + adamw_eps=1e-8, + none_grad=True, + debug=False, + ): + defaults = dict( + lr=lr, + wd=adamw_wd, + momentum=momentum, + nesterov=nesterov, + ns_steps=ns_steps, + adamw_betas=adamw_betas, + adamw_eps=adamw_eps, + none_grad=none_grad, + ) + + super().__init__(model.parameters(), defaults) + self.is_muon_func = is_muon_func + self.model = model + + if not dist.is_initialized(): + raise RuntimeError( + "Muon optimizer requires distributed training to be initialized." + ) + + self.rank = dist.get_rank() + + self.comm_stream = torch.cuda.Stream() + self.compute_stream = torch.cuda.Stream() + self.debug = debug + + def __setstate__(self, state): + # Sort parameters into those for which we will use Muon, and those for which we will not + super().__setstate__(state) + for name, p in self.model.named_parameters(): + if self.is_muon_func(p, name): + # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer + assert p.ndim == 2, p.ndim + self.state[p]["use_muon"] = True + self.state[p]["orig_shape"] = p.shape + else: + # Do not use Muon for parameters in adamw_params + self.state[p]["use_muon"] = False + + def _calc_flops(self, G, steps): + assert len(G.shape) == 2 + M, N = G.shape + if M > N: + M, N = N, M + + return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) + + def adjust_lr_for_muon(self, lr, param_shape): + A, B = param_shape[:2] + # We adjust the learning rate and weight decay based on the size of the parameter matrix + # as describted in the paper + adjusted_ratio = 0.2 * math.sqrt(max(A, B)) + adjusted_lr = lr * adjusted_ratio + return adjusted_lr + + def init_state_and_assign_params(self, params, group): + param_to_state = {} + param_to_flops = {} + + total_flops = 0 + for p in params: + g = p.grad + if g is None: + continue + assert g.ndim == 2, "Muon only supports 2D parameters." + + flops = self._calc_flops(g, group["ns_steps"]) + param_to_flops[id(p)] = flops + total_flops += flops + + if self.debug: + print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", flush=True) + + ordered_params = sorted( + params, key=lambda p: param_to_flops[id(p)], reverse=True + ) + + round_robin = 0 + mesh = None + for p in ordered_params: + if mesh is None: + mesh = p.device_mesh + if mesh.ndim != 1: + raise NotImplementedError( + "Muon requires a 1D mesh for distributed training yet." + ) + elif mesh != p.device_mesh: + raise ValueError("All parameters must be on the same mesh.") + + param_to_state[id(p)] = _muon_state() + param_to_state[id(p)].worker_rank = mesh.mesh[round_robin].item() + + round_robin = (round_robin + 1) % mesh.mesh.numel() + + return param_to_state, ordered_params + + def base(self, params, group, lr, wd, momentum): + # generate weight updates in distributed fashion + for p in params: + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + assert g is not None + + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + + u = _zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) + + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + + # apply weight decay + p.data.mul_(1 - lr * wd) + + # apply update + p.data.add_(u, alpha=-adjusted_lr) + + def _update_g(self, p, g, group, momentum): + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + return g + + def _update_p(self, p, u, lr, wd): + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + # apply weight decay + p.data.mul_(1 - lr * wd) + # apply update + p.data.add_(u, alpha=-adjusted_lr) + + def parallel(self, params, group, lr, wd, momentum): + """ + Perform a parallel optimization step using Muon. + """ + + for p in params: + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + + # Update g in the local rank + g = self._update_g( + p, + g, + group, + momentum=momentum, + ) + p.grad = g + + param_to_state, ordered_params = self.init_state_and_assign_params( + params, group + ) + + def enqueue_gathers(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _gather(p, state, self.rank, self.comm_stream, group["none_grad"]) + + def enqueue_computes(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _compute_u(state, group["ns_steps"], self.rank, self.compute_stream) + + def enqueue_scatters(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream) + + chunk_size = params[0].device_mesh.mesh.numel() + + # Wait grad update + self.comm_stream.wait_stream(torch.cuda.current_stream()) + + enqueue_gathers(0, chunk_size) + for i in range(0, len(params) + chunk_size - 1, chunk_size): + enqueue_computes(i, chunk_size) + enqueue_gathers(i + chunk_size, chunk_size) + enqueue_scatters(i, chunk_size) + + torch.cuda.current_stream().wait_stream(self.comm_stream) + + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + ############################ + # Muon # + ############################ + + params = [p for p in group["params"] if self.state[p]["use_muon"]] + lr = group["lr"] + wd = group["wd"] + momentum = group["momentum"] + + if isinstance(params[0].data, DTensor): + self.parallel( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) + else: + self.base( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) + + ############################ + # AdamW backup # + ############################ + + params = [p for p in group["params"] if not self.state[p]["use_muon"]] + lr = group["lr"] + beta1, beta2 = group["adamw_betas"] + eps = group["adamw_eps"] + weight_decay = group["wd"] + + for p in params: + g = p.grad + if g is None: + continue + state = self.state[p] + if "step" not in state: + state["step"] = 0 + state["moment1"] = torch.zeros_like(g) + state["moment2"] = torch.zeros_like(g) + state["step"] += 1 + step = state["step"] + buf1 = state["moment1"] + buf2 = state["moment2"] + buf1.lerp_(g, 1 - beta1) + buf2.lerp_(g.square(), 1 - beta2) + + g = buf1 / (eps + buf2.sqrt()) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + scale = bias_correction1 / bias_correction2**0.5 + p.data.mul_(1 - lr * weight_decay) + p.data.add_(g, alpha=-lr / scale) + + return loss diff --git a/build/torch210-cxx11-cu130-x86_64-linux/__init__.py b/build/torch26-cxx11-cu126-x86_64-linux/optimizer/__init__.py old mode 100644 new mode 100755 similarity index 100% rename from build/torch210-cxx11-cu130-x86_64-linux/__init__.py rename to build/torch26-cxx11-cu126-x86_64-linux/optimizer/__init__.py diff --git a/build/torch26-cxx11-cu126-x86_64-linux/optimizer/_ops.py b/build/torch26-cxx11-cu126-x86_64-linux/optimizer/_ops.py new file mode 100755 index 0000000000000000000000000000000000000000..7cf68eab4638da3512b5e49541c916ebd12301f0 --- /dev/null +++ b/build/torch26-cxx11-cu126-x86_64-linux/optimizer/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _optimizer_036642a_dirty +ops = torch.ops._optimizer_036642a_dirty + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_optimizer_036642a_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx11-cu126-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so b/build/torch26-cxx11-cu126-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..d8368d3904638b07e920f72204987d1820114e0a --- /dev/null +++ b/build/torch26-cxx11-cu126-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:46e01e1d957ada2d485b30cd60bc3ef7230b8857dffc59f2e7924339761ec577 +size 1824224 diff --git a/build/torch26-cxx11-cu126-x86_64-linux/optimizer/muon.py b/build/torch26-cxx11-cu126-x86_64-linux/optimizer/muon.py new file mode 100755 index 0000000000000000000000000000000000000000..0d614d55d721efac406c147b4f62e6c703a91107 --- /dev/null +++ b/build/torch26-cxx11-cu126-x86_64-linux/optimizer/muon.py @@ -0,0 +1,455 @@ +import math +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch.distributed._tensor import DTensor + + +# This code snippet is a modified version adapted from the following GitHub repositories: +# https://github.com/KellerJordan/Muon/blob/master/muon.py +@torch.no_grad() +def _zeropower_via_newtonschulz5(G, steps): + """ + Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a + quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose + of minimizing steps, it turns out to be empirically effective to keep increasing the slope at + zero even beyond the point where the iteration no longer converges all the way to one everywhere + on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T + where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model + performance at all relative to UV^T, where USV^T = G is the SVD. + """ + assert len(G.shape) == 2 + a, b, c = (3.4445, -4.7750, 2.0315) + X = G # no manual typecast + if G.size(0) > G.size(1): + X = X.T + # Ensure spectral norm is at most 1 + X = X / (X.norm() + 1e-7) + X = X.bfloat16() + # Perform the NS iterations + for _ in range(steps): + A = X @ X.T + # B = ( + # b * A + c * A @ A + # ) + B = torch.addmm(A, A, A, alpha=c, beta=b) + # X = a * X + B @ X + X = torch.addmm(X, B, X, alpha=1.0, beta=a) + + if G.size(0) > G.size(1): + X = X.T + return X.to(G.dtype) + + +@dataclass +class _muon_state: + # TODO: use Optional + worker_rank: int | None = None + gathered_grad: torch.Tensor | None = None + computed_u: torch.Tensor | None = None + gather_event: torch.cuda.Event | None = None + compute_event: torch.cuda.Event | None = None + + +@torch.no_grad() +def _gather(p, state, rank, comm_stream, none_grad): + g = p.grad + mesh = g.device_mesh + + if rank == state.worker_rank: + gather_list = [torch.empty_like(g.to_local()) for _ in range(mesh.mesh.numel())] + else: + gather_list = None + + with torch.cuda.stream(comm_stream): + torch.distributed.gather( + g.to_local(), + dst=state.worker_rank, + gather_list=gather_list, + group=mesh.get_group(), + ) + if rank == state.worker_rank: + if state.gathered_grad is not None: + raise RuntimeError( + "Gather event already exists, which should not happen." + ) + state.gathered_grad = torch.cat(gather_list, dim=0) + state.gather_event = torch.cuda.Event() + state.gather_event.record() + else: + state.gathered_grad = None + state.gather_event = None + if none_grad: + p.grad = None + + +@torch.no_grad() +def _compute_u(state, steps, rank, compute_stream): + with torch.cuda.stream(compute_stream): + if rank == state.worker_rank: + if state.gather_event is None: + raise RuntimeError("Gather event must be set before compute.") + compute_stream.wait_event(state.gather_event) + u = _zeropower_via_newtonschulz5(state.gathered_grad, steps) + state.computed_u = u + state.compute_event = torch.cuda.Event() + state.compute_event.record() + # Clear the gathered gradient to free memory + state.gathered_grad = None + else: + state.computed_u = None + state.compute_event = None + + +@torch.no_grad() +def _scatter(p, state, lr, wd, rank, comm_stream): + u = state.computed_u + mesh = p.device_mesh + + with torch.cuda.stream(comm_stream): + if rank == state.worker_rank: + if state.compute_event is None: + raise RuntimeError("Compute event must be set before scatter.") + comm_stream.wait_event(state.compute_event) + scatter_list = list(torch.split(u, p.size(0) // mesh.mesh.numel(), dim=0)) + else: + scatter_list = None + + u = torch.empty_like(p.to_local()) + torch.distributed.scatter( + u, + scatter_list=scatter_list, + src=state.worker_rank, + group=mesh.get_group(), + ) + if rank == state.worker_rank: + # Clear u to free memory + state.computed_u = None + u = DTensor.from_local( + u, + placements=p.placements, + device_mesh=mesh, + ) + p.data.mul_(1 - lr * wd) + p.data.add_(u, alpha=-lr) + + +class Muon(torch.optim.Optimizer): + """ + Muon - MomentUm Orthogonalized by Newton-schulz + + Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- + processing step, in which each 2D parameter's update is replaced with the nearest orthogonal + matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has + the advantage that it can be stably run in bfloat16 on the GPU. + + Some warnings: + - We believe this optimizer is unlikely to work well for training with small batch size. + - We believe it may not work well for finetuning pretrained models, but we haven't tested this. + + Arguments: + muon_params: The parameters to be optimized by Muon. + lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) + momentum: The momentum used by the internal SGD. (0.95 is a good default) + nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) + ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) + adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are + {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. + adamw_lr: The learning rate for the internal AdamW. + adamw_betas: The betas for the internal AdamW. + adamw_eps: The epsilon for the internal AdamW. + adamw_wd: The weight decay for the internal AdamW. + """ + + def __init__( + self, + model, + is_muon_func, + lr=1e-3, + momentum=0.95, + nesterov=True, + ns_steps=5, + adamw_wd=0.1, + adamw_betas=(0.9, 0.95), + adamw_eps=1e-8, + none_grad=True, + debug=False, + ): + defaults = dict( + lr=lr, + wd=adamw_wd, + momentum=momentum, + nesterov=nesterov, + ns_steps=ns_steps, + adamw_betas=adamw_betas, + adamw_eps=adamw_eps, + none_grad=none_grad, + ) + + super().__init__(model.parameters(), defaults) + self.is_muon_func = is_muon_func + self.model = model + + if not dist.is_initialized(): + raise RuntimeError( + "Muon optimizer requires distributed training to be initialized." + ) + + self.rank = dist.get_rank() + + self.comm_stream = torch.cuda.Stream() + self.compute_stream = torch.cuda.Stream() + self.debug = debug + + def __setstate__(self, state): + # Sort parameters into those for which we will use Muon, and those for which we will not + super().__setstate__(state) + for name, p in self.model.named_parameters(): + if self.is_muon_func(p, name): + # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer + assert p.ndim == 2, p.ndim + self.state[p]["use_muon"] = True + self.state[p]["orig_shape"] = p.shape + else: + # Do not use Muon for parameters in adamw_params + self.state[p]["use_muon"] = False + + def _calc_flops(self, G, steps): + assert len(G.shape) == 2 + M, N = G.shape + if M > N: + M, N = N, M + + return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) + + def adjust_lr_for_muon(self, lr, param_shape): + A, B = param_shape[:2] + # We adjust the learning rate and weight decay based on the size of the parameter matrix + # as describted in the paper + adjusted_ratio = 0.2 * math.sqrt(max(A, B)) + adjusted_lr = lr * adjusted_ratio + return adjusted_lr + + def init_state_and_assign_params(self, params, group): + param_to_state = {} + param_to_flops = {} + + total_flops = 0 + for p in params: + g = p.grad + if g is None: + continue + assert g.ndim == 2, "Muon only supports 2D parameters." + + flops = self._calc_flops(g, group["ns_steps"]) + param_to_flops[id(p)] = flops + total_flops += flops + + if self.debug: + print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", flush=True) + + ordered_params = sorted( + params, key=lambda p: param_to_flops[id(p)], reverse=True + ) + + round_robin = 0 + mesh = None + for p in ordered_params: + if mesh is None: + mesh = p.device_mesh + if mesh.ndim != 1: + raise NotImplementedError( + "Muon requires a 1D mesh for distributed training yet." + ) + elif mesh != p.device_mesh: + raise ValueError("All parameters must be on the same mesh.") + + param_to_state[id(p)] = _muon_state() + param_to_state[id(p)].worker_rank = mesh.mesh[round_robin].item() + + round_robin = (round_robin + 1) % mesh.mesh.numel() + + return param_to_state, ordered_params + + def base(self, params, group, lr, wd, momentum): + # generate weight updates in distributed fashion + for p in params: + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + assert g is not None + + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + + u = _zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) + + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + + # apply weight decay + p.data.mul_(1 - lr * wd) + + # apply update + p.data.add_(u, alpha=-adjusted_lr) + + def _update_g(self, p, g, group, momentum): + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + return g + + def _update_p(self, p, u, lr, wd): + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + # apply weight decay + p.data.mul_(1 - lr * wd) + # apply update + p.data.add_(u, alpha=-adjusted_lr) + + def parallel(self, params, group, lr, wd, momentum): + """ + Perform a parallel optimization step using Muon. + """ + + for p in params: + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + + # Update g in the local rank + g = self._update_g( + p, + g, + group, + momentum=momentum, + ) + p.grad = g + + param_to_state, ordered_params = self.init_state_and_assign_params( + params, group + ) + + def enqueue_gathers(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _gather(p, state, self.rank, self.comm_stream, group["none_grad"]) + + def enqueue_computes(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _compute_u(state, group["ns_steps"], self.rank, self.compute_stream) + + def enqueue_scatters(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream) + + chunk_size = params[0].device_mesh.mesh.numel() + + # Wait grad update + self.comm_stream.wait_stream(torch.cuda.current_stream()) + + enqueue_gathers(0, chunk_size) + for i in range(0, len(params) + chunk_size - 1, chunk_size): + enqueue_computes(i, chunk_size) + enqueue_gathers(i + chunk_size, chunk_size) + enqueue_scatters(i, chunk_size) + + torch.cuda.current_stream().wait_stream(self.comm_stream) + + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + ############################ + # Muon # + ############################ + + params = [p for p in group["params"] if self.state[p]["use_muon"]] + lr = group["lr"] + wd = group["wd"] + momentum = group["momentum"] + + if isinstance(params[0].data, DTensor): + self.parallel( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) + else: + self.base( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) + + ############################ + # AdamW backup # + ############################ + + params = [p for p in group["params"] if not self.state[p]["use_muon"]] + lr = group["lr"] + beta1, beta2 = group["adamw_betas"] + eps = group["adamw_eps"] + weight_decay = group["wd"] + + for p in params: + g = p.grad + if g is None: + continue + state = self.state[p] + if "step" not in state: + state["step"] = 0 + state["moment1"] = torch.zeros_like(g) + state["moment2"] = torch.zeros_like(g) + state["step"] += 1 + step = state["step"] + buf1 = state["moment1"] + buf2 = state["moment2"] + buf1.lerp_(g, 1 - beta1) + buf2.lerp_(g.square(), 1 - beta2) + + g = buf1 / (eps + buf2.sqrt()) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + scale = bias_correction1 / bias_correction2**0.5 + p.data.mul_(1 - lr * weight_decay) + p.data.add_(g, alpha=-lr / scale) + + return loss diff --git a/build/torch210-cxx11-rocm70-x86_64-linux/__init__.py b/build/torch26-cxx11-rocm62-x86_64-linux/optimizer/__init__.py old mode 100644 new mode 100755 similarity index 100% rename from build/torch210-cxx11-rocm70-x86_64-linux/__init__.py rename to build/torch26-cxx11-rocm62-x86_64-linux/optimizer/__init__.py diff --git a/build/torch26-cxx11-rocm62-x86_64-linux/optimizer/_ops.py b/build/torch26-cxx11-rocm62-x86_64-linux/optimizer/_ops.py new file mode 100755 index 0000000000000000000000000000000000000000..7cf68eab4638da3512b5e49541c916ebd12301f0 --- /dev/null +++ b/build/torch26-cxx11-rocm62-x86_64-linux/optimizer/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _optimizer_036642a_dirty +ops = torch.ops._optimizer_036642a_dirty + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_optimizer_036642a_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx11-rocm62-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so b/build/torch26-cxx11-rocm62-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..873c79d168145a8e956f609a3be376e9f1817b41 --- /dev/null +++ b/build/torch26-cxx11-rocm62-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a825a0cd31d8c1b91aa9db4b24248d7fc0a506615f625a385b40e6002025c7dd +size 1749744 diff --git a/build/torch26-cxx11-rocm62-x86_64-linux/optimizer/muon.py b/build/torch26-cxx11-rocm62-x86_64-linux/optimizer/muon.py new file mode 100755 index 0000000000000000000000000000000000000000..0d614d55d721efac406c147b4f62e6c703a91107 --- /dev/null +++ b/build/torch26-cxx11-rocm62-x86_64-linux/optimizer/muon.py @@ -0,0 +1,455 @@ +import math +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch.distributed._tensor import DTensor + + +# This code snippet is a modified version adapted from the following GitHub repositories: +# https://github.com/KellerJordan/Muon/blob/master/muon.py +@torch.no_grad() +def _zeropower_via_newtonschulz5(G, steps): + """ + Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a + quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose + of minimizing steps, it turns out to be empirically effective to keep increasing the slope at + zero even beyond the point where the iteration no longer converges all the way to one everywhere + on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T + where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model + performance at all relative to UV^T, where USV^T = G is the SVD. + """ + assert len(G.shape) == 2 + a, b, c = (3.4445, -4.7750, 2.0315) + X = G # no manual typecast + if G.size(0) > G.size(1): + X = X.T + # Ensure spectral norm is at most 1 + X = X / (X.norm() + 1e-7) + X = X.bfloat16() + # Perform the NS iterations + for _ in range(steps): + A = X @ X.T + # B = ( + # b * A + c * A @ A + # ) + B = torch.addmm(A, A, A, alpha=c, beta=b) + # X = a * X + B @ X + X = torch.addmm(X, B, X, alpha=1.0, beta=a) + + if G.size(0) > G.size(1): + X = X.T + return X.to(G.dtype) + + +@dataclass +class _muon_state: + # TODO: use Optional + worker_rank: int | None = None + gathered_grad: torch.Tensor | None = None + computed_u: torch.Tensor | None = None + gather_event: torch.cuda.Event | None = None + compute_event: torch.cuda.Event | None = None + + +@torch.no_grad() +def _gather(p, state, rank, comm_stream, none_grad): + g = p.grad + mesh = g.device_mesh + + if rank == state.worker_rank: + gather_list = [torch.empty_like(g.to_local()) for _ in range(mesh.mesh.numel())] + else: + gather_list = None + + with torch.cuda.stream(comm_stream): + torch.distributed.gather( + g.to_local(), + dst=state.worker_rank, + gather_list=gather_list, + group=mesh.get_group(), + ) + if rank == state.worker_rank: + if state.gathered_grad is not None: + raise RuntimeError( + "Gather event already exists, which should not happen." + ) + state.gathered_grad = torch.cat(gather_list, dim=0) + state.gather_event = torch.cuda.Event() + state.gather_event.record() + else: + state.gathered_grad = None + state.gather_event = None + if none_grad: + p.grad = None + + +@torch.no_grad() +def _compute_u(state, steps, rank, compute_stream): + with torch.cuda.stream(compute_stream): + if rank == state.worker_rank: + if state.gather_event is None: + raise RuntimeError("Gather event must be set before compute.") + compute_stream.wait_event(state.gather_event) + u = _zeropower_via_newtonschulz5(state.gathered_grad, steps) + state.computed_u = u + state.compute_event = torch.cuda.Event() + state.compute_event.record() + # Clear the gathered gradient to free memory + state.gathered_grad = None + else: + state.computed_u = None + state.compute_event = None + + +@torch.no_grad() +def _scatter(p, state, lr, wd, rank, comm_stream): + u = state.computed_u + mesh = p.device_mesh + + with torch.cuda.stream(comm_stream): + if rank == state.worker_rank: + if state.compute_event is None: + raise RuntimeError("Compute event must be set before scatter.") + comm_stream.wait_event(state.compute_event) + scatter_list = list(torch.split(u, p.size(0) // mesh.mesh.numel(), dim=0)) + else: + scatter_list = None + + u = torch.empty_like(p.to_local()) + torch.distributed.scatter( + u, + scatter_list=scatter_list, + src=state.worker_rank, + group=mesh.get_group(), + ) + if rank == state.worker_rank: + # Clear u to free memory + state.computed_u = None + u = DTensor.from_local( + u, + placements=p.placements, + device_mesh=mesh, + ) + p.data.mul_(1 - lr * wd) + p.data.add_(u, alpha=-lr) + + +class Muon(torch.optim.Optimizer): + """ + Muon - MomentUm Orthogonalized by Newton-schulz + + Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- + processing step, in which each 2D parameter's update is replaced with the nearest orthogonal + matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has + the advantage that it can be stably run in bfloat16 on the GPU. + + Some warnings: + - We believe this optimizer is unlikely to work well for training with small batch size. + - We believe it may not work well for finetuning pretrained models, but we haven't tested this. + + Arguments: + muon_params: The parameters to be optimized by Muon. + lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) + momentum: The momentum used by the internal SGD. (0.95 is a good default) + nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) + ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) + adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are + {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. + adamw_lr: The learning rate for the internal AdamW. + adamw_betas: The betas for the internal AdamW. + adamw_eps: The epsilon for the internal AdamW. + adamw_wd: The weight decay for the internal AdamW. + """ + + def __init__( + self, + model, + is_muon_func, + lr=1e-3, + momentum=0.95, + nesterov=True, + ns_steps=5, + adamw_wd=0.1, + adamw_betas=(0.9, 0.95), + adamw_eps=1e-8, + none_grad=True, + debug=False, + ): + defaults = dict( + lr=lr, + wd=adamw_wd, + momentum=momentum, + nesterov=nesterov, + ns_steps=ns_steps, + adamw_betas=adamw_betas, + adamw_eps=adamw_eps, + none_grad=none_grad, + ) + + super().__init__(model.parameters(), defaults) + self.is_muon_func = is_muon_func + self.model = model + + if not dist.is_initialized(): + raise RuntimeError( + "Muon optimizer requires distributed training to be initialized." + ) + + self.rank = dist.get_rank() + + self.comm_stream = torch.cuda.Stream() + self.compute_stream = torch.cuda.Stream() + self.debug = debug + + def __setstate__(self, state): + # Sort parameters into those for which we will use Muon, and those for which we will not + super().__setstate__(state) + for name, p in self.model.named_parameters(): + if self.is_muon_func(p, name): + # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer + assert p.ndim == 2, p.ndim + self.state[p]["use_muon"] = True + self.state[p]["orig_shape"] = p.shape + else: + # Do not use Muon for parameters in adamw_params + self.state[p]["use_muon"] = False + + def _calc_flops(self, G, steps): + assert len(G.shape) == 2 + M, N = G.shape + if M > N: + M, N = N, M + + return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) + + def adjust_lr_for_muon(self, lr, param_shape): + A, B = param_shape[:2] + # We adjust the learning rate and weight decay based on the size of the parameter matrix + # as describted in the paper + adjusted_ratio = 0.2 * math.sqrt(max(A, B)) + adjusted_lr = lr * adjusted_ratio + return adjusted_lr + + def init_state_and_assign_params(self, params, group): + param_to_state = {} + param_to_flops = {} + + total_flops = 0 + for p in params: + g = p.grad + if g is None: + continue + assert g.ndim == 2, "Muon only supports 2D parameters." + + flops = self._calc_flops(g, group["ns_steps"]) + param_to_flops[id(p)] = flops + total_flops += flops + + if self.debug: + print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", flush=True) + + ordered_params = sorted( + params, key=lambda p: param_to_flops[id(p)], reverse=True + ) + + round_robin = 0 + mesh = None + for p in ordered_params: + if mesh is None: + mesh = p.device_mesh + if mesh.ndim != 1: + raise NotImplementedError( + "Muon requires a 1D mesh for distributed training yet." + ) + elif mesh != p.device_mesh: + raise ValueError("All parameters must be on the same mesh.") + + param_to_state[id(p)] = _muon_state() + param_to_state[id(p)].worker_rank = mesh.mesh[round_robin].item() + + round_robin = (round_robin + 1) % mesh.mesh.numel() + + return param_to_state, ordered_params + + def base(self, params, group, lr, wd, momentum): + # generate weight updates in distributed fashion + for p in params: + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + assert g is not None + + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + + u = _zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) + + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + + # apply weight decay + p.data.mul_(1 - lr * wd) + + # apply update + p.data.add_(u, alpha=-adjusted_lr) + + def _update_g(self, p, g, group, momentum): + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + return g + + def _update_p(self, p, u, lr, wd): + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + # apply weight decay + p.data.mul_(1 - lr * wd) + # apply update + p.data.add_(u, alpha=-adjusted_lr) + + def parallel(self, params, group, lr, wd, momentum): + """ + Perform a parallel optimization step using Muon. + """ + + for p in params: + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + + # Update g in the local rank + g = self._update_g( + p, + g, + group, + momentum=momentum, + ) + p.grad = g + + param_to_state, ordered_params = self.init_state_and_assign_params( + params, group + ) + + def enqueue_gathers(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _gather(p, state, self.rank, self.comm_stream, group["none_grad"]) + + def enqueue_computes(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _compute_u(state, group["ns_steps"], self.rank, self.compute_stream) + + def enqueue_scatters(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream) + + chunk_size = params[0].device_mesh.mesh.numel() + + # Wait grad update + self.comm_stream.wait_stream(torch.cuda.current_stream()) + + enqueue_gathers(0, chunk_size) + for i in range(0, len(params) + chunk_size - 1, chunk_size): + enqueue_computes(i, chunk_size) + enqueue_gathers(i + chunk_size, chunk_size) + enqueue_scatters(i, chunk_size) + + torch.cuda.current_stream().wait_stream(self.comm_stream) + + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + ############################ + # Muon # + ############################ + + params = [p for p in group["params"] if self.state[p]["use_muon"]] + lr = group["lr"] + wd = group["wd"] + momentum = group["momentum"] + + if isinstance(params[0].data, DTensor): + self.parallel( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) + else: + self.base( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) + + ############################ + # AdamW backup # + ############################ + + params = [p for p in group["params"] if not self.state[p]["use_muon"]] + lr = group["lr"] + beta1, beta2 = group["adamw_betas"] + eps = group["adamw_eps"] + weight_decay = group["wd"] + + for p in params: + g = p.grad + if g is None: + continue + state = self.state[p] + if "step" not in state: + state["step"] = 0 + state["moment1"] = torch.zeros_like(g) + state["moment2"] = torch.zeros_like(g) + state["step"] += 1 + step = state["step"] + buf1 = state["moment1"] + buf2 = state["moment2"] + buf1.lerp_(g, 1 - beta1) + buf2.lerp_(g.square(), 1 - beta2) + + g = buf1 / (eps + buf2.sqrt()) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + scale = bias_correction1 / bias_correction2**0.5 + p.data.mul_(1 - lr * weight_decay) + p.data.add_(g, alpha=-lr / scale) + + return loss diff --git a/build/torch210-cxx11-rocm71-x86_64-linux/__init__.py b/build/torch26-cxx98-cu118-x86_64-linux/optimizer/__init__.py old mode 100644 new mode 100755 similarity index 100% rename from build/torch210-cxx11-rocm71-x86_64-linux/__init__.py rename to build/torch26-cxx98-cu118-x86_64-linux/optimizer/__init__.py diff --git a/build/torch26-cxx98-cu118-x86_64-linux/optimizer/_ops.py b/build/torch26-cxx98-cu118-x86_64-linux/optimizer/_ops.py new file mode 100755 index 0000000000000000000000000000000000000000..7cf68eab4638da3512b5e49541c916ebd12301f0 --- /dev/null +++ b/build/torch26-cxx98-cu118-x86_64-linux/optimizer/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _optimizer_036642a_dirty +ops = torch.ops._optimizer_036642a_dirty + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_optimizer_036642a_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx98-cu118-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so b/build/torch26-cxx98-cu118-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..65af7b83eba70e3218649b947d359fda84b41be0 --- /dev/null +++ b/build/torch26-cxx98-cu118-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:579e9ddf66a4f17ead9232c2f32e6327fe6a3f16dd235e2e73e6cb282de1797e +size 1787192 diff --git a/build/torch26-cxx98-cu118-x86_64-linux/optimizer/muon.py b/build/torch26-cxx98-cu118-x86_64-linux/optimizer/muon.py new file mode 100755 index 0000000000000000000000000000000000000000..0d614d55d721efac406c147b4f62e6c703a91107 --- /dev/null +++ b/build/torch26-cxx98-cu118-x86_64-linux/optimizer/muon.py @@ -0,0 +1,455 @@ +import math +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch.distributed._tensor import DTensor + + +# This code snippet is a modified version adapted from the following GitHub repositories: +# https://github.com/KellerJordan/Muon/blob/master/muon.py +@torch.no_grad() +def _zeropower_via_newtonschulz5(G, steps): + """ + Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a + quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose + of minimizing steps, it turns out to be empirically effective to keep increasing the slope at + zero even beyond the point where the iteration no longer converges all the way to one everywhere + on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T + where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model + performance at all relative to UV^T, where USV^T = G is the SVD. + """ + assert len(G.shape) == 2 + a, b, c = (3.4445, -4.7750, 2.0315) + X = G # no manual typecast + if G.size(0) > G.size(1): + X = X.T + # Ensure spectral norm is at most 1 + X = X / (X.norm() + 1e-7) + X = X.bfloat16() + # Perform the NS iterations + for _ in range(steps): + A = X @ X.T + # B = ( + # b * A + c * A @ A + # ) + B = torch.addmm(A, A, A, alpha=c, beta=b) + # X = a * X + B @ X + X = torch.addmm(X, B, X, alpha=1.0, beta=a) + + if G.size(0) > G.size(1): + X = X.T + return X.to(G.dtype) + + +@dataclass +class _muon_state: + # TODO: use Optional + worker_rank: int | None = None + gathered_grad: torch.Tensor | None = None + computed_u: torch.Tensor | None = None + gather_event: torch.cuda.Event | None = None + compute_event: torch.cuda.Event | None = None + + +@torch.no_grad() +def _gather(p, state, rank, comm_stream, none_grad): + g = p.grad + mesh = g.device_mesh + + if rank == state.worker_rank: + gather_list = [torch.empty_like(g.to_local()) for _ in range(mesh.mesh.numel())] + else: + gather_list = None + + with torch.cuda.stream(comm_stream): + torch.distributed.gather( + g.to_local(), + dst=state.worker_rank, + gather_list=gather_list, + group=mesh.get_group(), + ) + if rank == state.worker_rank: + if state.gathered_grad is not None: + raise RuntimeError( + "Gather event already exists, which should not happen." + ) + state.gathered_grad = torch.cat(gather_list, dim=0) + state.gather_event = torch.cuda.Event() + state.gather_event.record() + else: + state.gathered_grad = None + state.gather_event = None + if none_grad: + p.grad = None + + +@torch.no_grad() +def _compute_u(state, steps, rank, compute_stream): + with torch.cuda.stream(compute_stream): + if rank == state.worker_rank: + if state.gather_event is None: + raise RuntimeError("Gather event must be set before compute.") + compute_stream.wait_event(state.gather_event) + u = _zeropower_via_newtonschulz5(state.gathered_grad, steps) + state.computed_u = u + state.compute_event = torch.cuda.Event() + state.compute_event.record() + # Clear the gathered gradient to free memory + state.gathered_grad = None + else: + state.computed_u = None + state.compute_event = None + + +@torch.no_grad() +def _scatter(p, state, lr, wd, rank, comm_stream): + u = state.computed_u + mesh = p.device_mesh + + with torch.cuda.stream(comm_stream): + if rank == state.worker_rank: + if state.compute_event is None: + raise RuntimeError("Compute event must be set before scatter.") + comm_stream.wait_event(state.compute_event) + scatter_list = list(torch.split(u, p.size(0) // mesh.mesh.numel(), dim=0)) + else: + scatter_list = None + + u = torch.empty_like(p.to_local()) + torch.distributed.scatter( + u, + scatter_list=scatter_list, + src=state.worker_rank, + group=mesh.get_group(), + ) + if rank == state.worker_rank: + # Clear u to free memory + state.computed_u = None + u = DTensor.from_local( + u, + placements=p.placements, + device_mesh=mesh, + ) + p.data.mul_(1 - lr * wd) + p.data.add_(u, alpha=-lr) + + +class Muon(torch.optim.Optimizer): + """ + Muon - MomentUm Orthogonalized by Newton-schulz + + Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- + processing step, in which each 2D parameter's update is replaced with the nearest orthogonal + matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has + the advantage that it can be stably run in bfloat16 on the GPU. + + Some warnings: + - We believe this optimizer is unlikely to work well for training with small batch size. + - We believe it may not work well for finetuning pretrained models, but we haven't tested this. + + Arguments: + muon_params: The parameters to be optimized by Muon. + lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) + momentum: The momentum used by the internal SGD. (0.95 is a good default) + nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) + ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) + adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are + {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. + adamw_lr: The learning rate for the internal AdamW. + adamw_betas: The betas for the internal AdamW. + adamw_eps: The epsilon for the internal AdamW. + adamw_wd: The weight decay for the internal AdamW. + """ + + def __init__( + self, + model, + is_muon_func, + lr=1e-3, + momentum=0.95, + nesterov=True, + ns_steps=5, + adamw_wd=0.1, + adamw_betas=(0.9, 0.95), + adamw_eps=1e-8, + none_grad=True, + debug=False, + ): + defaults = dict( + lr=lr, + wd=adamw_wd, + momentum=momentum, + nesterov=nesterov, + ns_steps=ns_steps, + adamw_betas=adamw_betas, + adamw_eps=adamw_eps, + none_grad=none_grad, + ) + + super().__init__(model.parameters(), defaults) + self.is_muon_func = is_muon_func + self.model = model + + if not dist.is_initialized(): + raise RuntimeError( + "Muon optimizer requires distributed training to be initialized." + ) + + self.rank = dist.get_rank() + + self.comm_stream = torch.cuda.Stream() + self.compute_stream = torch.cuda.Stream() + self.debug = debug + + def __setstate__(self, state): + # Sort parameters into those for which we will use Muon, and those for which we will not + super().__setstate__(state) + for name, p in self.model.named_parameters(): + if self.is_muon_func(p, name): + # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer + assert p.ndim == 2, p.ndim + self.state[p]["use_muon"] = True + self.state[p]["orig_shape"] = p.shape + else: + # Do not use Muon for parameters in adamw_params + self.state[p]["use_muon"] = False + + def _calc_flops(self, G, steps): + assert len(G.shape) == 2 + M, N = G.shape + if M > N: + M, N = N, M + + return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) + + def adjust_lr_for_muon(self, lr, param_shape): + A, B = param_shape[:2] + # We adjust the learning rate and weight decay based on the size of the parameter matrix + # as describted in the paper + adjusted_ratio = 0.2 * math.sqrt(max(A, B)) + adjusted_lr = lr * adjusted_ratio + return adjusted_lr + + def init_state_and_assign_params(self, params, group): + param_to_state = {} + param_to_flops = {} + + total_flops = 0 + for p in params: + g = p.grad + if g is None: + continue + assert g.ndim == 2, "Muon only supports 2D parameters." + + flops = self._calc_flops(g, group["ns_steps"]) + param_to_flops[id(p)] = flops + total_flops += flops + + if self.debug: + print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", flush=True) + + ordered_params = sorted( + params, key=lambda p: param_to_flops[id(p)], reverse=True + ) + + round_robin = 0 + mesh = None + for p in ordered_params: + if mesh is None: + mesh = p.device_mesh + if mesh.ndim != 1: + raise NotImplementedError( + "Muon requires a 1D mesh for distributed training yet." + ) + elif mesh != p.device_mesh: + raise ValueError("All parameters must be on the same mesh.") + + param_to_state[id(p)] = _muon_state() + param_to_state[id(p)].worker_rank = mesh.mesh[round_robin].item() + + round_robin = (round_robin + 1) % mesh.mesh.numel() + + return param_to_state, ordered_params + + def base(self, params, group, lr, wd, momentum): + # generate weight updates in distributed fashion + for p in params: + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + assert g is not None + + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + + u = _zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) + + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + + # apply weight decay + p.data.mul_(1 - lr * wd) + + # apply update + p.data.add_(u, alpha=-adjusted_lr) + + def _update_g(self, p, g, group, momentum): + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + return g + + def _update_p(self, p, u, lr, wd): + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + # apply weight decay + p.data.mul_(1 - lr * wd) + # apply update + p.data.add_(u, alpha=-adjusted_lr) + + def parallel(self, params, group, lr, wd, momentum): + """ + Perform a parallel optimization step using Muon. + """ + + for p in params: + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + + # Update g in the local rank + g = self._update_g( + p, + g, + group, + momentum=momentum, + ) + p.grad = g + + param_to_state, ordered_params = self.init_state_and_assign_params( + params, group + ) + + def enqueue_gathers(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _gather(p, state, self.rank, self.comm_stream, group["none_grad"]) + + def enqueue_computes(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _compute_u(state, group["ns_steps"], self.rank, self.compute_stream) + + def enqueue_scatters(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream) + + chunk_size = params[0].device_mesh.mesh.numel() + + # Wait grad update + self.comm_stream.wait_stream(torch.cuda.current_stream()) + + enqueue_gathers(0, chunk_size) + for i in range(0, len(params) + chunk_size - 1, chunk_size): + enqueue_computes(i, chunk_size) + enqueue_gathers(i + chunk_size, chunk_size) + enqueue_scatters(i, chunk_size) + + torch.cuda.current_stream().wait_stream(self.comm_stream) + + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + ############################ + # Muon # + ############################ + + params = [p for p in group["params"] if self.state[p]["use_muon"]] + lr = group["lr"] + wd = group["wd"] + momentum = group["momentum"] + + if isinstance(params[0].data, DTensor): + self.parallel( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) + else: + self.base( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) + + ############################ + # AdamW backup # + ############################ + + params = [p for p in group["params"] if not self.state[p]["use_muon"]] + lr = group["lr"] + beta1, beta2 = group["adamw_betas"] + eps = group["adamw_eps"] + weight_decay = group["wd"] + + for p in params: + g = p.grad + if g is None: + continue + state = self.state[p] + if "step" not in state: + state["step"] = 0 + state["moment1"] = torch.zeros_like(g) + state["moment2"] = torch.zeros_like(g) + state["step"] += 1 + step = state["step"] + buf1 = state["moment1"] + buf2 = state["moment2"] + buf1.lerp_(g, 1 - beta1) + buf2.lerp_(g.square(), 1 - beta2) + + g = buf1 / (eps + buf2.sqrt()) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + scale = bias_correction1 / bias_correction2**0.5 + p.data.mul_(1 - lr * weight_decay) + p.data.add_(g, alpha=-lr / scale) + + return loss diff --git a/build/torch28-cxx11-cu126-x86_64-linux/__init__.py b/build/torch26-cxx98-cu124-x86_64-linux/optimizer/__init__.py old mode 100644 new mode 100755 similarity index 100% rename from build/torch28-cxx11-cu126-x86_64-linux/__init__.py rename to build/torch26-cxx98-cu124-x86_64-linux/optimizer/__init__.py diff --git a/build/torch26-cxx98-cu124-x86_64-linux/optimizer/_ops.py b/build/torch26-cxx98-cu124-x86_64-linux/optimizer/_ops.py new file mode 100755 index 0000000000000000000000000000000000000000..7cf68eab4638da3512b5e49541c916ebd12301f0 --- /dev/null +++ b/build/torch26-cxx98-cu124-x86_64-linux/optimizer/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _optimizer_036642a_dirty +ops = torch.ops._optimizer_036642a_dirty + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_optimizer_036642a_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx98-cu124-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so b/build/torch26-cxx98-cu124-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..c1a231bc7234f19a04d1c05f491511c3d21ebaca --- /dev/null +++ b/build/torch26-cxx98-cu124-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:beacb4ba2d56463b6d444875728b3462cb3ff6c1449e3c9693cd665bfbbbbb73 +size 1824184 diff --git a/build/torch26-cxx98-cu124-x86_64-linux/optimizer/muon.py b/build/torch26-cxx98-cu124-x86_64-linux/optimizer/muon.py new file mode 100755 index 0000000000000000000000000000000000000000..0d614d55d721efac406c147b4f62e6c703a91107 --- /dev/null +++ b/build/torch26-cxx98-cu124-x86_64-linux/optimizer/muon.py @@ -0,0 +1,455 @@ +import math +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch.distributed._tensor import DTensor + + +# This code snippet is a modified version adapted from the following GitHub repositories: +# https://github.com/KellerJordan/Muon/blob/master/muon.py +@torch.no_grad() +def _zeropower_via_newtonschulz5(G, steps): + """ + Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a + quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose + of minimizing steps, it turns out to be empirically effective to keep increasing the slope at + zero even beyond the point where the iteration no longer converges all the way to one everywhere + on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T + where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model + performance at all relative to UV^T, where USV^T = G is the SVD. + """ + assert len(G.shape) == 2 + a, b, c = (3.4445, -4.7750, 2.0315) + X = G # no manual typecast + if G.size(0) > G.size(1): + X = X.T + # Ensure spectral norm is at most 1 + X = X / (X.norm() + 1e-7) + X = X.bfloat16() + # Perform the NS iterations + for _ in range(steps): + A = X @ X.T + # B = ( + # b * A + c * A @ A + # ) + B = torch.addmm(A, A, A, alpha=c, beta=b) + # X = a * X + B @ X + X = torch.addmm(X, B, X, alpha=1.0, beta=a) + + if G.size(0) > G.size(1): + X = X.T + return X.to(G.dtype) + + +@dataclass +class _muon_state: + # TODO: use Optional + worker_rank: int | None = None + gathered_grad: torch.Tensor | None = None + computed_u: torch.Tensor | None = None + gather_event: torch.cuda.Event | None = None + compute_event: torch.cuda.Event | None = None + + +@torch.no_grad() +def _gather(p, state, rank, comm_stream, none_grad): + g = p.grad + mesh = g.device_mesh + + if rank == state.worker_rank: + gather_list = [torch.empty_like(g.to_local()) for _ in range(mesh.mesh.numel())] + else: + gather_list = None + + with torch.cuda.stream(comm_stream): + torch.distributed.gather( + g.to_local(), + dst=state.worker_rank, + gather_list=gather_list, + group=mesh.get_group(), + ) + if rank == state.worker_rank: + if state.gathered_grad is not None: + raise RuntimeError( + "Gather event already exists, which should not happen." + ) + state.gathered_grad = torch.cat(gather_list, dim=0) + state.gather_event = torch.cuda.Event() + state.gather_event.record() + else: + state.gathered_grad = None + state.gather_event = None + if none_grad: + p.grad = None + + +@torch.no_grad() +def _compute_u(state, steps, rank, compute_stream): + with torch.cuda.stream(compute_stream): + if rank == state.worker_rank: + if state.gather_event is None: + raise RuntimeError("Gather event must be set before compute.") + compute_stream.wait_event(state.gather_event) + u = _zeropower_via_newtonschulz5(state.gathered_grad, steps) + state.computed_u = u + state.compute_event = torch.cuda.Event() + state.compute_event.record() + # Clear the gathered gradient to free memory + state.gathered_grad = None + else: + state.computed_u = None + state.compute_event = None + + +@torch.no_grad() +def _scatter(p, state, lr, wd, rank, comm_stream): + u = state.computed_u + mesh = p.device_mesh + + with torch.cuda.stream(comm_stream): + if rank == state.worker_rank: + if state.compute_event is None: + raise RuntimeError("Compute event must be set before scatter.") + comm_stream.wait_event(state.compute_event) + scatter_list = list(torch.split(u, p.size(0) // mesh.mesh.numel(), dim=0)) + else: + scatter_list = None + + u = torch.empty_like(p.to_local()) + torch.distributed.scatter( + u, + scatter_list=scatter_list, + src=state.worker_rank, + group=mesh.get_group(), + ) + if rank == state.worker_rank: + # Clear u to free memory + state.computed_u = None + u = DTensor.from_local( + u, + placements=p.placements, + device_mesh=mesh, + ) + p.data.mul_(1 - lr * wd) + p.data.add_(u, alpha=-lr) + + +class Muon(torch.optim.Optimizer): + """ + Muon - MomentUm Orthogonalized by Newton-schulz + + Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- + processing step, in which each 2D parameter's update is replaced with the nearest orthogonal + matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has + the advantage that it can be stably run in bfloat16 on the GPU. + + Some warnings: + - We believe this optimizer is unlikely to work well for training with small batch size. + - We believe it may not work well for finetuning pretrained models, but we haven't tested this. + + Arguments: + muon_params: The parameters to be optimized by Muon. + lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) + momentum: The momentum used by the internal SGD. (0.95 is a good default) + nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) + ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) + adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are + {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. + adamw_lr: The learning rate for the internal AdamW. + adamw_betas: The betas for the internal AdamW. + adamw_eps: The epsilon for the internal AdamW. + adamw_wd: The weight decay for the internal AdamW. + """ + + def __init__( + self, + model, + is_muon_func, + lr=1e-3, + momentum=0.95, + nesterov=True, + ns_steps=5, + adamw_wd=0.1, + adamw_betas=(0.9, 0.95), + adamw_eps=1e-8, + none_grad=True, + debug=False, + ): + defaults = dict( + lr=lr, + wd=adamw_wd, + momentum=momentum, + nesterov=nesterov, + ns_steps=ns_steps, + adamw_betas=adamw_betas, + adamw_eps=adamw_eps, + none_grad=none_grad, + ) + + super().__init__(model.parameters(), defaults) + self.is_muon_func = is_muon_func + self.model = model + + if not dist.is_initialized(): + raise RuntimeError( + "Muon optimizer requires distributed training to be initialized." + ) + + self.rank = dist.get_rank() + + self.comm_stream = torch.cuda.Stream() + self.compute_stream = torch.cuda.Stream() + self.debug = debug + + def __setstate__(self, state): + # Sort parameters into those for which we will use Muon, and those for which we will not + super().__setstate__(state) + for name, p in self.model.named_parameters(): + if self.is_muon_func(p, name): + # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer + assert p.ndim == 2, p.ndim + self.state[p]["use_muon"] = True + self.state[p]["orig_shape"] = p.shape + else: + # Do not use Muon for parameters in adamw_params + self.state[p]["use_muon"] = False + + def _calc_flops(self, G, steps): + assert len(G.shape) == 2 + M, N = G.shape + if M > N: + M, N = N, M + + return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) + + def adjust_lr_for_muon(self, lr, param_shape): + A, B = param_shape[:2] + # We adjust the learning rate and weight decay based on the size of the parameter matrix + # as describted in the paper + adjusted_ratio = 0.2 * math.sqrt(max(A, B)) + adjusted_lr = lr * adjusted_ratio + return adjusted_lr + + def init_state_and_assign_params(self, params, group): + param_to_state = {} + param_to_flops = {} + + total_flops = 0 + for p in params: + g = p.grad + if g is None: + continue + assert g.ndim == 2, "Muon only supports 2D parameters." + + flops = self._calc_flops(g, group["ns_steps"]) + param_to_flops[id(p)] = flops + total_flops += flops + + if self.debug: + print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", flush=True) + + ordered_params = sorted( + params, key=lambda p: param_to_flops[id(p)], reverse=True + ) + + round_robin = 0 + mesh = None + for p in ordered_params: + if mesh is None: + mesh = p.device_mesh + if mesh.ndim != 1: + raise NotImplementedError( + "Muon requires a 1D mesh for distributed training yet." + ) + elif mesh != p.device_mesh: + raise ValueError("All parameters must be on the same mesh.") + + param_to_state[id(p)] = _muon_state() + param_to_state[id(p)].worker_rank = mesh.mesh[round_robin].item() + + round_robin = (round_robin + 1) % mesh.mesh.numel() + + return param_to_state, ordered_params + + def base(self, params, group, lr, wd, momentum): + # generate weight updates in distributed fashion + for p in params: + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + assert g is not None + + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + + u = _zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) + + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + + # apply weight decay + p.data.mul_(1 - lr * wd) + + # apply update + p.data.add_(u, alpha=-adjusted_lr) + + def _update_g(self, p, g, group, momentum): + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + return g + + def _update_p(self, p, u, lr, wd): + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + # apply weight decay + p.data.mul_(1 - lr * wd) + # apply update + p.data.add_(u, alpha=-adjusted_lr) + + def parallel(self, params, group, lr, wd, momentum): + """ + Perform a parallel optimization step using Muon. + """ + + for p in params: + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + + # Update g in the local rank + g = self._update_g( + p, + g, + group, + momentum=momentum, + ) + p.grad = g + + param_to_state, ordered_params = self.init_state_and_assign_params( + params, group + ) + + def enqueue_gathers(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _gather(p, state, self.rank, self.comm_stream, group["none_grad"]) + + def enqueue_computes(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _compute_u(state, group["ns_steps"], self.rank, self.compute_stream) + + def enqueue_scatters(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream) + + chunk_size = params[0].device_mesh.mesh.numel() + + # Wait grad update + self.comm_stream.wait_stream(torch.cuda.current_stream()) + + enqueue_gathers(0, chunk_size) + for i in range(0, len(params) + chunk_size - 1, chunk_size): + enqueue_computes(i, chunk_size) + enqueue_gathers(i + chunk_size, chunk_size) + enqueue_scatters(i, chunk_size) + + torch.cuda.current_stream().wait_stream(self.comm_stream) + + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + ############################ + # Muon # + ############################ + + params = [p for p in group["params"] if self.state[p]["use_muon"]] + lr = group["lr"] + wd = group["wd"] + momentum = group["momentum"] + + if isinstance(params[0].data, DTensor): + self.parallel( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) + else: + self.base( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) + + ############################ + # AdamW backup # + ############################ + + params = [p for p in group["params"] if not self.state[p]["use_muon"]] + lr = group["lr"] + beta1, beta2 = group["adamw_betas"] + eps = group["adamw_eps"] + weight_decay = group["wd"] + + for p in params: + g = p.grad + if g is None: + continue + state = self.state[p] + if "step" not in state: + state["step"] = 0 + state["moment1"] = torch.zeros_like(g) + state["moment2"] = torch.zeros_like(g) + state["step"] += 1 + step = state["step"] + buf1 = state["moment1"] + buf2 = state["moment2"] + buf1.lerp_(g, 1 - beta1) + buf2.lerp_(g.square(), 1 - beta2) + + g = buf1 / (eps + buf2.sqrt()) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + scale = bias_correction1 / bias_correction2**0.5 + p.data.mul_(1 - lr * weight_decay) + p.data.add_(g, alpha=-lr / scale) + + return loss diff --git a/build/torch28-cxx11-cu128-x86_64-linux/__init__.py b/build/torch26-cxx98-cu126-x86_64-linux/optimizer/__init__.py old mode 100644 new mode 100755 similarity index 100% rename from build/torch28-cxx11-cu128-x86_64-linux/__init__.py rename to build/torch26-cxx98-cu126-x86_64-linux/optimizer/__init__.py diff --git a/build/torch26-cxx98-cu126-x86_64-linux/optimizer/_ops.py b/build/torch26-cxx98-cu126-x86_64-linux/optimizer/_ops.py new file mode 100755 index 0000000000000000000000000000000000000000..7cf68eab4638da3512b5e49541c916ebd12301f0 --- /dev/null +++ b/build/torch26-cxx98-cu126-x86_64-linux/optimizer/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _optimizer_036642a_dirty +ops = torch.ops._optimizer_036642a_dirty + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_optimizer_036642a_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx98-cu126-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so b/build/torch26-cxx98-cu126-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..ac6d1fa4c9a0f286c2b7597d1f2f337aff89a570 --- /dev/null +++ b/build/torch26-cxx98-cu126-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9b04b011803d328d8dcd2edcf4c3840ddbb1bb2f093464c208f0ba2faf4f16bc +size 1824184 diff --git a/build/torch26-cxx98-cu126-x86_64-linux/optimizer/muon.py b/build/torch26-cxx98-cu126-x86_64-linux/optimizer/muon.py new file mode 100755 index 0000000000000000000000000000000000000000..0d614d55d721efac406c147b4f62e6c703a91107 --- /dev/null +++ b/build/torch26-cxx98-cu126-x86_64-linux/optimizer/muon.py @@ -0,0 +1,455 @@ +import math +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch.distributed._tensor import DTensor + + +# This code snippet is a modified version adapted from the following GitHub repositories: +# https://github.com/KellerJordan/Muon/blob/master/muon.py +@torch.no_grad() +def _zeropower_via_newtonschulz5(G, steps): + """ + Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a + quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose + of minimizing steps, it turns out to be empirically effective to keep increasing the slope at + zero even beyond the point where the iteration no longer converges all the way to one everywhere + on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T + where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model + performance at all relative to UV^T, where USV^T = G is the SVD. + """ + assert len(G.shape) == 2 + a, b, c = (3.4445, -4.7750, 2.0315) + X = G # no manual typecast + if G.size(0) > G.size(1): + X = X.T + # Ensure spectral norm is at most 1 + X = X / (X.norm() + 1e-7) + X = X.bfloat16() + # Perform the NS iterations + for _ in range(steps): + A = X @ X.T + # B = ( + # b * A + c * A @ A + # ) + B = torch.addmm(A, A, A, alpha=c, beta=b) + # X = a * X + B @ X + X = torch.addmm(X, B, X, alpha=1.0, beta=a) + + if G.size(0) > G.size(1): + X = X.T + return X.to(G.dtype) + + +@dataclass +class _muon_state: + # TODO: use Optional + worker_rank: int | None = None + gathered_grad: torch.Tensor | None = None + computed_u: torch.Tensor | None = None + gather_event: torch.cuda.Event | None = None + compute_event: torch.cuda.Event | None = None + + +@torch.no_grad() +def _gather(p, state, rank, comm_stream, none_grad): + g = p.grad + mesh = g.device_mesh + + if rank == state.worker_rank: + gather_list = [torch.empty_like(g.to_local()) for _ in range(mesh.mesh.numel())] + else: + gather_list = None + + with torch.cuda.stream(comm_stream): + torch.distributed.gather( + g.to_local(), + dst=state.worker_rank, + gather_list=gather_list, + group=mesh.get_group(), + ) + if rank == state.worker_rank: + if state.gathered_grad is not None: + raise RuntimeError( + "Gather event already exists, which should not happen." + ) + state.gathered_grad = torch.cat(gather_list, dim=0) + state.gather_event = torch.cuda.Event() + state.gather_event.record() + else: + state.gathered_grad = None + state.gather_event = None + if none_grad: + p.grad = None + + +@torch.no_grad() +def _compute_u(state, steps, rank, compute_stream): + with torch.cuda.stream(compute_stream): + if rank == state.worker_rank: + if state.gather_event is None: + raise RuntimeError("Gather event must be set before compute.") + compute_stream.wait_event(state.gather_event) + u = _zeropower_via_newtonschulz5(state.gathered_grad, steps) + state.computed_u = u + state.compute_event = torch.cuda.Event() + state.compute_event.record() + # Clear the gathered gradient to free memory + state.gathered_grad = None + else: + state.computed_u = None + state.compute_event = None + + +@torch.no_grad() +def _scatter(p, state, lr, wd, rank, comm_stream): + u = state.computed_u + mesh = p.device_mesh + + with torch.cuda.stream(comm_stream): + if rank == state.worker_rank: + if state.compute_event is None: + raise RuntimeError("Compute event must be set before scatter.") + comm_stream.wait_event(state.compute_event) + scatter_list = list(torch.split(u, p.size(0) // mesh.mesh.numel(), dim=0)) + else: + scatter_list = None + + u = torch.empty_like(p.to_local()) + torch.distributed.scatter( + u, + scatter_list=scatter_list, + src=state.worker_rank, + group=mesh.get_group(), + ) + if rank == state.worker_rank: + # Clear u to free memory + state.computed_u = None + u = DTensor.from_local( + u, + placements=p.placements, + device_mesh=mesh, + ) + p.data.mul_(1 - lr * wd) + p.data.add_(u, alpha=-lr) + + +class Muon(torch.optim.Optimizer): + """ + Muon - MomentUm Orthogonalized by Newton-schulz + + Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- + processing step, in which each 2D parameter's update is replaced with the nearest orthogonal + matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has + the advantage that it can be stably run in bfloat16 on the GPU. + + Some warnings: + - We believe this optimizer is unlikely to work well for training with small batch size. + - We believe it may not work well for finetuning pretrained models, but we haven't tested this. + + Arguments: + muon_params: The parameters to be optimized by Muon. + lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) + momentum: The momentum used by the internal SGD. (0.95 is a good default) + nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) + ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) + adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are + {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. + adamw_lr: The learning rate for the internal AdamW. + adamw_betas: The betas for the internal AdamW. + adamw_eps: The epsilon for the internal AdamW. + adamw_wd: The weight decay for the internal AdamW. + """ + + def __init__( + self, + model, + is_muon_func, + lr=1e-3, + momentum=0.95, + nesterov=True, + ns_steps=5, + adamw_wd=0.1, + adamw_betas=(0.9, 0.95), + adamw_eps=1e-8, + none_grad=True, + debug=False, + ): + defaults = dict( + lr=lr, + wd=adamw_wd, + momentum=momentum, + nesterov=nesterov, + ns_steps=ns_steps, + adamw_betas=adamw_betas, + adamw_eps=adamw_eps, + none_grad=none_grad, + ) + + super().__init__(model.parameters(), defaults) + self.is_muon_func = is_muon_func + self.model = model + + if not dist.is_initialized(): + raise RuntimeError( + "Muon optimizer requires distributed training to be initialized." + ) + + self.rank = dist.get_rank() + + self.comm_stream = torch.cuda.Stream() + self.compute_stream = torch.cuda.Stream() + self.debug = debug + + def __setstate__(self, state): + # Sort parameters into those for which we will use Muon, and those for which we will not + super().__setstate__(state) + for name, p in self.model.named_parameters(): + if self.is_muon_func(p, name): + # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer + assert p.ndim == 2, p.ndim + self.state[p]["use_muon"] = True + self.state[p]["orig_shape"] = p.shape + else: + # Do not use Muon for parameters in adamw_params + self.state[p]["use_muon"] = False + + def _calc_flops(self, G, steps): + assert len(G.shape) == 2 + M, N = G.shape + if M > N: + M, N = N, M + + return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) + + def adjust_lr_for_muon(self, lr, param_shape): + A, B = param_shape[:2] + # We adjust the learning rate and weight decay based on the size of the parameter matrix + # as describted in the paper + adjusted_ratio = 0.2 * math.sqrt(max(A, B)) + adjusted_lr = lr * adjusted_ratio + return adjusted_lr + + def init_state_and_assign_params(self, params, group): + param_to_state = {} + param_to_flops = {} + + total_flops = 0 + for p in params: + g = p.grad + if g is None: + continue + assert g.ndim == 2, "Muon only supports 2D parameters." + + flops = self._calc_flops(g, group["ns_steps"]) + param_to_flops[id(p)] = flops + total_flops += flops + + if self.debug: + print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", flush=True) + + ordered_params = sorted( + params, key=lambda p: param_to_flops[id(p)], reverse=True + ) + + round_robin = 0 + mesh = None + for p in ordered_params: + if mesh is None: + mesh = p.device_mesh + if mesh.ndim != 1: + raise NotImplementedError( + "Muon requires a 1D mesh for distributed training yet." + ) + elif mesh != p.device_mesh: + raise ValueError("All parameters must be on the same mesh.") + + param_to_state[id(p)] = _muon_state() + param_to_state[id(p)].worker_rank = mesh.mesh[round_robin].item() + + round_robin = (round_robin + 1) % mesh.mesh.numel() + + return param_to_state, ordered_params + + def base(self, params, group, lr, wd, momentum): + # generate weight updates in distributed fashion + for p in params: + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + assert g is not None + + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + + u = _zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) + + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + + # apply weight decay + p.data.mul_(1 - lr * wd) + + # apply update + p.data.add_(u, alpha=-adjusted_lr) + + def _update_g(self, p, g, group, momentum): + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + return g + + def _update_p(self, p, u, lr, wd): + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + # apply weight decay + p.data.mul_(1 - lr * wd) + # apply update + p.data.add_(u, alpha=-adjusted_lr) + + def parallel(self, params, group, lr, wd, momentum): + """ + Perform a parallel optimization step using Muon. + """ + + for p in params: + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + + # Update g in the local rank + g = self._update_g( + p, + g, + group, + momentum=momentum, + ) + p.grad = g + + param_to_state, ordered_params = self.init_state_and_assign_params( + params, group + ) + + def enqueue_gathers(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _gather(p, state, self.rank, self.comm_stream, group["none_grad"]) + + def enqueue_computes(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _compute_u(state, group["ns_steps"], self.rank, self.compute_stream) + + def enqueue_scatters(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream) + + chunk_size = params[0].device_mesh.mesh.numel() + + # Wait grad update + self.comm_stream.wait_stream(torch.cuda.current_stream()) + + enqueue_gathers(0, chunk_size) + for i in range(0, len(params) + chunk_size - 1, chunk_size): + enqueue_computes(i, chunk_size) + enqueue_gathers(i + chunk_size, chunk_size) + enqueue_scatters(i, chunk_size) + + torch.cuda.current_stream().wait_stream(self.comm_stream) + + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + ############################ + # Muon # + ############################ + + params = [p for p in group["params"] if self.state[p]["use_muon"]] + lr = group["lr"] + wd = group["wd"] + momentum = group["momentum"] + + if isinstance(params[0].data, DTensor): + self.parallel( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) + else: + self.base( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) + + ############################ + # AdamW backup # + ############################ + + params = [p for p in group["params"] if not self.state[p]["use_muon"]] + lr = group["lr"] + beta1, beta2 = group["adamw_betas"] + eps = group["adamw_eps"] + weight_decay = group["wd"] + + for p in params: + g = p.grad + if g is None: + continue + state = self.state[p] + if "step" not in state: + state["step"] = 0 + state["moment1"] = torch.zeros_like(g) + state["moment2"] = torch.zeros_like(g) + state["step"] += 1 + step = state["step"] + buf1 = state["moment1"] + buf2 = state["moment2"] + buf1.lerp_(g, 1 - beta1) + buf2.lerp_(g.square(), 1 - beta2) + + g = buf1 / (eps + buf2.sqrt()) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + scale = bias_correction1 / bias_correction2**0.5 + p.data.mul_(1 - lr * weight_decay) + p.data.add_(g, alpha=-lr / scale) + + return loss diff --git a/build/torch27-cxx11-cu118-x86_64-linux/optimizer/__init__.py b/build/torch27-cxx11-cu118-x86_64-linux/optimizer/__init__.py old mode 100644 new mode 100755 diff --git a/build/torch27-cxx11-cu118-x86_64-linux/optimizer/_ops.py b/build/torch27-cxx11-cu118-x86_64-linux/optimizer/_ops.py old mode 100644 new mode 100755 index cb9efd677b388ebc299d6c4747eee701c96211f6..7cf68eab4638da3512b5e49541c916ebd12301f0 --- a/build/torch27-cxx11-cu118-x86_64-linux/optimizer/_ops.py +++ b/build/torch27-cxx11-cu118-x86_64-linux/optimizer/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _optimizer_b0230e7_dirty -ops = torch.ops._optimizer_b0230e7_dirty +from . import _optimizer_036642a_dirty +ops = torch.ops._optimizer_036642a_dirty def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_optimizer_b0230e7_dirty::{op_name}" \ No newline at end of file + return f"_optimizer_036642a_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu118-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so b/build/torch27-cxx11-cu118-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..72737f2570d97f57b8f15140e6fd9e12623152f2 --- /dev/null +++ b/build/torch27-cxx11-cu118-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ad6c725009f2e776b99d3134c75f15e11dd7fe75fe4ba1fa94779018c7871f8c +size 1787368 diff --git a/build/torch27-cxx11-cu118-x86_64-linux/optimizer/_optimizer_b0230e7_dirty.abi3.so b/build/torch27-cxx11-cu118-x86_64-linux/optimizer/_optimizer_b0230e7_dirty.abi3.so deleted file mode 100755 index ac384d7194105cc1fd531bed4212a63de4d9be00..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu118-x86_64-linux/optimizer/_optimizer_b0230e7_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:236bb0d67cbb2718b076637569923cf240de1c7a074790623ecb9c049fca9732 -size 1787368 diff --git a/build/torch27-cxx11-cu118-x86_64-linux/optimizer/matmul_transpose_triton.py b/build/torch27-cxx11-cu118-x86_64-linux/optimizer/matmul_transpose_triton.py deleted file mode 100644 index 4565b2c4fd506a4218340d380d6c962b16774b1d..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu118-x86_64-linux/optimizer/matmul_transpose_triton.py +++ /dev/null @@ -1,128 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -def matmul_transpose_assign(d_in, d_out): - assert d_in.is_cuda, "Input `d_in` must be a CUDA tensor" - assert d_out.is_cuda, "Input `d_out` must be a CUDA tensor" - assert d_in.device == d_out.device, "Inputs `d_in` and `d_out` must be on the same CUDA device" - assert d_in.dtype == d_out.dtype, "Inputs must have the same data type" - assert d_in.ndim == 2, "Input `d_in` must be a 2D tensor" - assert d_out.ndim == 2, "Input `d_out` must be a 2D tensor" - assert d_in.size(0) == d_out.size(0) == d_out.size(0), \ - "First dimension of `d_in` must match first and second dimension of `d_out`" - - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -def matmul_transpose(d_in): - M, _ = d_in.shape - d_out = torch.empty((M, M), device=d_in.device, dtype=d_in.dtype) - matmul_transpose_assign(d_in, d_out) - return d_out diff --git a/build/torch27-cxx11-cu118-x86_64-linux/optimizer/muon.py b/build/torch27-cxx11-cu118-x86_64-linux/optimizer/muon.py old mode 100644 new mode 100755 index 4af25d55c528fb0db2272540838ab70eb3619194..0d614d55d721efac406c147b4f62e6c703a91107 --- a/build/torch27-cxx11-cu118-x86_64-linux/optimizer/muon.py +++ b/build/torch27-cxx11-cu118-x86_64-linux/optimizer/muon.py @@ -1,26 +1,14 @@ -import logging import math -import types from dataclasses import dataclass -from typing import List, Optional, Union, cast import torch import torch.distributed as dist -from torch.distributed._tensor import DTensor, Replicate, Shard - -from .matmul_transpose_triton import matmul_transpose_assign - -logger = logging.getLogger(__name__) - -COMM_DTYPE = torch.bfloat16 +from torch.distributed._tensor import DTensor # This code snippet is a modified version adapted from the following GitHub repositories: # https://github.com/KellerJordan/Muon/blob/master/muon.py -# Muon's Newton–Schulz iteration causes high variance in singular values -# Idea: give each iteration its own 3 coefficients and optimize them via gradient descent. @torch.no_grad() -# matmul_transpose_assign from : https://github.com/nil0x9/flash-muon def _zeropower_via_newtonschulz5(G, steps): """ Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a @@ -32,31 +20,26 @@ def _zeropower_via_newtonschulz5(G, steps): performance at all relative to UV^T, where USV^T = G is the SVD. """ assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE + a, b, c = (3.4445, -4.7750, 2.0315) X = G # no manual typecast - if G.size(0) > G.size(1): X = X.T # Ensure spectral norm is at most 1 X = X / (X.norm() + 1e-7) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) + X = X.bfloat16() # Perform the NS iterations - for a, b, c in [ - (4.0848, -6.8946, 2.9270), - (3.9505, -6.3029, 2.6377), - (3.7418, -5.5913, 2.3037), - (2.8769, -3.1427, 1.2046), - (2.8366, -3.0525, 1.2012), - ]: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) + for _ in range(steps): + A = X @ X.T + # B = ( + # b * A + c * A @ A + # ) + B = torch.addmm(A, A, A, alpha=c, beta=b) + # X = a * X + B @ X + X = torch.addmm(X, B, X, alpha=1.0, beta=a) if G.size(0) > G.size(1): X = X.T - return X + return X.to(G.dtype) @dataclass @@ -64,425 +47,92 @@ class _muon_state: # TODO: use Optional worker_rank: int | None = None gathered_grad: torch.Tensor | None = None - scattered_u: DTensor | None = None computed_u: torch.Tensor | None = None gather_event: torch.cuda.Event | None = None compute_event: torch.cuda.Event | None = None - scatter_event: torch.cuda.Event | None = None - process_group = None - qk_clip_state = None - - -def split_elems_for_src(param, src_rank, num_ranks) -> int: - rows = param.shape[0] - cols = int(param.numel() // rows) - base, rem = divmod(rows, num_ranks) - my_rows = base + (1 if src_rank < rem else 0) - return my_rows * cols @torch.no_grad() -def _alloc_gathered_grad(params, param_to_state, rank, compute_stream): - """ - Pre-allocate gathered_grad buffer on compute_stream - before launching all2all gather - """ - with torch.cuda.stream(compute_stream): - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - num_ranks = dist.get_world_size(group=state.process_group) - state.gathered_grad = torch.empty(p.grad.numel(), - dtype=COMM_DTYPE, - device="cuda") - else: - state.gathered_grad = None - - alloc_event = torch.cuda.Event() - alloc_event.record(compute_stream) - return alloc_event +def _gather(p, state, rank, comm_stream, none_grad): + g = p.grad + mesh = g.device_mesh + if rank == state.worker_rank: + gather_list = [torch.empty_like(g.to_local()) for _ in range(mesh.mesh.numel())] + else: + gather_list = None -@torch.no_grad() -def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad, - alloc_event): - """ - All2all gathers shards so each owner rank reconstructs its full gradient - """ with torch.cuda.stream(comm_stream): - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - - # Construct sending buffers - per_dst = [[] for _ in range(num_ranks)] - send_counts = [0] * num_ranks - - for p in params: - state = param_to_state[id(p)] - dst = state.worker_rank - assert dst < num_ranks - shard_elems = split_elems_for_src(p, rank, num_ranks) - g = p.grad - g = g.to_local().to(COMM_DTYPE).contiguous().view(-1) - assert g.numel() == shard_elems - per_dst[dst].append(g) - send_counts[dst] += shard_elems - - assert any( - len(v) > 0 for v in per_dst - ), "At least one destination rank must receive a sharded tensor" - # list[list[Tensor]] -> list[Tensor] - per_dst = [t for dst in per_dst for t in dst] - - send_buf = torch.cat(per_dst, dim=0) - - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - # Compute receive sizes and allocate receiving buffers - recv_counts = [0] * num_ranks - - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += split_elems_for_src(p, src, num_ranks) - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - #All2All - dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, + torch.distributed.gather( + g.to_local(), + dst=state.worker_rank, + gather_list=gather_list, + group=mesh.get_group(), ) - - # Reconstructs gathered grad from the received buffer - # - # recv_buf (num ranks = 3) - # - # From rank 0 From rank 1 From rank 2 - # | p1_0, p2_0, p3_0 | p1_1, p2_1, p3_1 | p1_2, p2_2, p3_2 | - # - # Outer loop: - # rank 0 -> rank 1 -> rank2 - # - # Inner loop: - # p1_n -> p2_n -> p3_n - - comm_stream.wait_event(alloc_event) - - off = 0 - write_offsets = {id(p): 0 for p in owned_params} - for src in range(num_ranks): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - n = split_elems_for_src(p, src, num_ranks) - assert n > 0 - - sg = recv_buf.narrow(0, off + inner_off, n) - woff = write_offsets[id(p)] - dst = state.gathered_grad.narrow(0, woff, n) - dst.copy_(sg) - - write_offsets[id(p)] += n - inner_off += n - off += block - - for p in params: - state = param_to_state[id(p)] - if state.worker_rank == rank: - state.gathered_grad = state.gathered_grad.view_as(p) - state.gather_event = torch.cuda.Event() - state.gather_event.record(comm_stream) - else: - state.gathered_grad = None - state.gather_event = None - if none_grad: - p.grad = None + if rank == state.worker_rank: + if state.gathered_grad is not None: + raise RuntimeError( + "Gather event already exists, which should not happen." + ) + state.gathered_grad = torch.cat(gather_list, dim=0) + state.gather_event = torch.cuda.Event() + state.gather_event.record() + else: + state.gathered_grad = None + state.gather_event = None + if none_grad: + p.grad = None @torch.no_grad() -def _compute_u(p, state, steps, rank, compute_stream): - """ - On worker_rank, compute the orthogonalized update using Newton-Schulz iteration. - """ +def _compute_u(state, steps, rank, compute_stream): with torch.cuda.stream(compute_stream): if rank == state.worker_rank: if state.gather_event is None: raise RuntimeError("Gather event must be set before compute.") compute_stream.wait_event(state.gather_event) u = _zeropower_via_newtonschulz5(state.gathered_grad, steps) - state.gathered_grad = None state.computed_u = u state.compute_event = torch.cuda.Event() state.compute_event.record() + # Clear the gathered gradient to free memory + state.gathered_grad = None else: state.computed_u = None state.compute_event = None @torch.no_grad() -def _alloc_scattered_u(params, param_to_state, rank, compute_stream): - """ - Pre-allocate scattered_u buffer on compute_stream - before launching all2all gather - """ - with torch.cuda.stream(compute_stream): - for p in params: - state = param_to_state[id(p)] - state.scattered_u = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - alloc_event = torch.cuda.Event() - alloc_event.record(compute_stream) - return alloc_event - +def _scatter(p, state, lr, wd, rank, comm_stream): + u = state.computed_u + mesh = p.device_mesh -def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event): - """ - All2all scatters full gradients to all ranks - """ with torch.cuda.stream(comm_stream): - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - # Construct sending buffer - per_dst = [[] for _ in range(num_ranks)] - send_counts = [0] * num_ranks - - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - if state.compute_event is None: - raise RuntimeError( - "Compute event must be set before scatter.") - comm_stream.wait_event(state.compute_event) - state.gathered_grad = None - - assert state.computed_u is not None - - u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1) - - offset = 0 - for dst in range(num_ranks): - n = split_elems_for_src(p, dst, num_ranks) - assert n > 0 - - su = u_full.narrow(0, offset, n) - per_dst[dst].append(su) - send_counts[dst] += n - offset += n - - assert offset == u_full.numel() - - lengths = [len(v) for v in per_dst] - if all(l > 0 for l in lengths): - assert all( - l == lengths[0] for l in lengths - ), "All destination ranks must have the same number of sharded tensor" - # list[list[Tensor]] -> list[Tensor] - per_dst = [t for dst in per_dst for t in dst] - send_buf = torch.cat(per_dst, dim=0) + if rank == state.worker_rank: + if state.compute_event is None: + raise RuntimeError("Compute event must be set before scatter.") + comm_stream.wait_event(state.compute_event) + scatter_list = list(torch.split(u, p.size(0) // mesh.mesh.numel(), dim=0)) else: - # all_to_all requires participation from all ranks - # Even non-owner ranks must join the collective call - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Compute receive sizes and allocate receiving buffers - recv_counts = [0] * num_ranks - - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += split_elems_for_src(p, rank, num_ranks) - recv_counts[src] = total - - recv_total = sum(recv_counts) - assert recv_total > 0 - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - #All2All - dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, + scatter_list = None + + u = torch.empty_like(p.to_local()) + torch.distributed.scatter( + u, + scatter_list=scatter_list, + src=state.worker_rank, + group=mesh.get_group(), ) - - # Copy to pre-allocated scattered_u buffer from the received buffer - # - # recv_buf (num ranks = 3, local_rank = 0) - # - # From rank 0 From rank 1 From rank 2 - # | p1_0, p2_0, p3_0 | p4_0 | p5_0, p6_0 | - # - # Outer loop: - # rank 0 -> rank 1 -> rank2 - # - # Inner loop: - # src(0) : p1_0 -> p2_0 -> p3_0 - # src(1) : p4_0 - # src(2) : p5_0 -> p6_0 - - comm_stream.wait_event(alloc_event) - - off = 0 - for src in range(num_ranks): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = split_elems_for_src(p, rank, num_ranks) - assert n > 0 - - flat_local = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - state.scattered_u.copy_(flat_local) - - state.scatter_event = torch.cuda.Event() - state.scatter_event.record(comm_stream) - inner_off += n - - assert inner_off == block - off += block - - -def _update_param(p, state, lr, adjusted_lr, weight_decay, rank, - compute_stream): - """ - Update sharded parameter p with the scattered_u. - Only worker_rank frees computed_u. - """ - with torch.cuda.stream(compute_stream): - if state.scatter_event is None: - raise RuntimeError("Scatter event must be set before update") - compute_stream.wait_event(state.scatter_event) - u_dtensor = DTensor.from_local( - state.scattered_u, - placements=p.placements, - device_mesh=p.device_mesh, - ) - - state.scattered_u = u_dtensor - if rank == state.worker_rank: - # Free computed_u + # Clear u to free memory state.computed_u = None - - Muon._update_p(p, state.scattered_u, lr, adjusted_lr, weight_decay) - state.scattered_u = None - u_dtensor = None - - scales_full = Muon._compute_scales(p, state.qk_clip_state) - if scales_full is not None: - num_ranks = dist.get_world_size(group=state.process_group) - local_rank = dist.get_rank(group=state.process_group) - scales_local = scales_full.chunk(num_ranks, dim=0)[local_rank] - scales_local = DTensor.from_local( - scales_local, - placements=p.placements, - device_mesh=p.device_mesh, - ) - Muon._qk_clip(p, scales_local, state.qk_clip_state.head_dim) - - -def default_is_muon(name, x): - skip_keys = ["embed_tokens", "lm_head", "tok_embeddings", "output"] - return x.ndim >= 2 and not any(key in name for key in skip_keys) - - -def get_default_muon_param_groups(model, is_muon_func=default_is_muon): - muon_params, muon_names = [], [] - non_muon_params = [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - ('wq', 'wk', 'q_proj', 'k_proj') and return (kind, layer_index). - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = name.split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: Optional[str] # 'wq'/'q_proj' or 'wk'/'k_proj' or None - indices: List[int] # which heads to consider for clipping - head_dim: int # from config - threshold: float # from config - logit: Optional[torch.Tensor] + u = DTensor.from_local( + u, + placements=p.placements, + device_mesh=mesh, + ) + p.data.mul_(1 - lr * wd) + p.data.add_(u, alpha=-lr) class Muon(torch.optim.Optimizer): @@ -499,87 +149,71 @@ class Muon(torch.optim.Optimizer): - We believe it may not work well for finetuning pretrained models, but we haven't tested this. Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. + muon_params: The parameters to be optimized by Muon. lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) momentum: The momentum used by the internal SGD. (0.95 is a good default) nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. + adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. adamw_lr: The learning rate for the internal AdamW. adamw_betas: The betas for the internal AdamW. adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - overlap_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher overlap_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. + adamw_wd: The weight decay for the internal AdamW. """ - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config={ - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - }, - overlap_step=5): + def __init__( + self, + model, + is_muon_func, + lr=1e-3, + momentum=0.95, + nesterov=True, + ns_steps=5, + adamw_wd=0.1, + adamw_betas=(0.9, 0.95), + adamw_eps=1e-8, + none_grad=True, + debug=False, + ): defaults = dict( lr=lr, - weight_decay=weight_decay, + wd=adamw_wd, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps, adamw_betas=adamw_betas, adamw_eps=adamw_eps, none_grad=none_grad, - use_muon=True, ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) + super().__init__(model.parameters(), defaults) + self.is_muon_func = is_muon_func + self.model = model - super().__init__(params, defaults) + if not dist.is_initialized(): + raise RuntimeError( + "Muon optimizer requires distributed training to be initialized." + ) - self.rank = None + self.rank = dist.get_rank() self.comm_stream = torch.cuda.Stream() self.compute_stream = torch.cuda.Stream() self.debug = debug - self.clip_config = clip_config - self.overlap_step = overlap_step + + def __setstate__(self, state): + # Sort parameters into those for which we will use Muon, and those for which we will not + super().__setstate__(state) + for name, p in self.model.named_parameters(): + if self.is_muon_func(p, name): + # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer + assert p.ndim == 2, p.ndim + self.state[p]["use_muon"] = True + self.state[p]["orig_shape"] = p.shape + else: + # Do not use Muon for parameters in adamw_params + self.state[p]["use_muon"] = False def _calc_flops(self, G, steps): assert len(G.shape) == 2 @@ -597,30 +231,7 @@ class Muon(torch.optim.Optimizer): adjusted_lr = lr * adjusted_ratio return adjusted_lr - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - if p.placements == (Shard(dim=0), ): - # Case for FSDP - return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0) - elif p.placements == (Replicate(), Shard(dim=0)): - # Case for HSDP - process_group = p.device_mesh.get_group(mesh_dim=1) - if self.rank is None: - self.rank = dist.get_rank(group=process_group) - else: - assert self.rank == dist.get_rank(group=process_group) - for i, shard_mesh in enumerate(p.device_mesh.mesh): - if self.rank in shard_mesh: - return shard_mesh, p.device_mesh.get_group(mesh_dim=1) - else: - raise ValueError(f"Unsupported placements ({p.placements}).") - - def init_state_and_assign_params(self, names, params, group, qk_logits): + def init_state_and_assign_params(self, params, group): param_to_state = {} param_to_flops = {} @@ -636,44 +247,34 @@ class Muon(torch.optim.Optimizer): total_flops += flops if self.debug: - print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", - flush=True) - - paired = list(zip(names, params)) + print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", flush=True) - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) + ordered_params = sorted( + params, key=lambda p: param_to_flops[id(p)], reverse=True + ) round_robin = 0 mesh = None - shard_mesh = None - process_group = None - for n, p in zip(ordered_names, ordered_params): + for p in ordered_params: if mesh is None: mesh = p.device_mesh - shard_mesh, process_group = self.get_shard_mesh(p) + if mesh.ndim != 1: + raise NotImplementedError( + "Muon requires a 1D mesh for distributed training yet." + ) elif mesh != p.device_mesh: raise ValueError("All parameters must be on the same mesh.") - num_ranks = dist.get_world_size(group=process_group) + param_to_state[id(p)] = _muon_state() - param_to_state[id( - p)].worker_rank = shard_mesh[round_robin].item() % num_ranks - param_to_state[id(p)].process_group = process_group - qk_clip_state = self.get_qk_clip_info(n, qk_logits) - param_to_state[id(p)].qk_clip_state = qk_clip_state - round_robin = (round_robin + 1) % len(shard_mesh) + param_to_state[id(p)].worker_rank = mesh.mesh[round_robin].item() + + round_robin = (round_robin + 1) % mesh.mesh.numel() return param_to_state, ordered_params - def base(self, names, params, group, lr, weight_decay, momentum, - qk_logits): + def base(self, params, group, lr, wd, momentum): # generate weight updates in distributed fashion - for n, p in zip(names, params): + for p in params: g = p.grad if g is None: continue @@ -692,87 +293,39 @@ class Muon(torch.optim.Optimizer): else: g = buf - u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) + u = _zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) + # scale update adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) - Muon._update_p(p, u, lr, adjusted_lr, weight_decay) - qk_clip_state = self.get_qk_clip_info(n, qk_logits) + # apply weight decay + p.data.mul_(1 - lr * wd) - scales_full = self._compute_scales(p, qk_clip_state) - if scales_full is not None: - Muon._qk_clip(p, scales_full, qk_clip_state.head_dim) + # apply update + p.data.add_(u, alpha=-adjusted_lr) def _update_g(self, p, g, group, momentum): # calc update state = self.state[p] - buf = state.setdefault("momentum_buffer", torch.zeros_like(g)) - torch.add(g, buf, alpha=momentum, out=buf) + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) if group["nesterov"]: - g.add_(buf, alpha=momentum) - return g - return buf + g = g.add(buf, alpha=momentum) + else: + g = buf + return g - @staticmethod - def _update_p(p, u, lr, adjusted_lr, weight_decay): + def _update_p(self, p, u, lr, wd): + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) # apply weight decay - p.data.mul_(1 - lr * weight_decay) + p.data.mul_(1 - lr * wd) # apply update p.data.add_(u, alpha=-adjusted_lr) - def get_qk_clip_info(self, n, qk_logits): - head_dim = self.clip_config.get('head_dim') - threshold = self.clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - indices_key = 'q_indices' if 'q' in kind else 'k_indices' - indices = self.clip_config.get(indices_key, []) or [] - - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - @staticmethod - def _compute_scales(p, qk_clip_state): - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - H_global = p.shape[0] // head_dim - scales_full = torch.ones(H_global, device=p.data.device) - scaling = 0 - - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if new_scale < scales_full[head_idx]: - scales_full[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - scaling += 1 - - return scales_full if scaling > 0 else None - - @staticmethod - def _qk_clip(p, scales, head_dim): - W = p.data.view(-1, head_dim, p.data.shape[1]) - W.mul_(scales.view(-1, 1, 1)) - - def parallel(self, names, params, group, lr, weight_decay, momentum, - qk_logits): + def parallel(self, params, group, lr, wd, momentum): """ Perform a parallel optimization step using Muon. """ @@ -794,143 +347,44 @@ class Muon(torch.optim.Optimizer): p.grad = g param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - assert self.rank is not None + params, group + ) - def enqueue_all2all_gather(start_idx, chunk_size): - target_params = ordered_params[start_idx:start_idx + chunk_size] - if target_params: - alloc_event = _alloc_gathered_grad(target_params, - param_to_state, self.rank, - self.compute_stream) - _all2all_gather(target_params, param_to_state, self.rank, - self.comm_stream, group["none_grad"], - alloc_event) + def enqueue_gathers(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _gather(p, state, self.rank, self.comm_stream, group["none_grad"]) def enqueue_computes(start_idx, chunk_size): - for p in ordered_params[start_idx:start_idx + chunk_size]: + for p in ordered_params[start_idx : start_idx + chunk_size]: state = param_to_state[id(p)] - _compute_u(p, state, group["ns_steps"], self.rank, - self.compute_stream) - - def enqueue_all2all_scatter(start_idx, chunk_size): - target_params = ordered_params[start_idx:start_idx + chunk_size] - if target_params: - alloc_event = _alloc_scattered_u(target_params, param_to_state, - self.rank, - self.compute_stream) - _all2all_scatter(target_params, param_to_state, self.rank, - self.comm_stream, alloc_event) - - def enqueue_update_param(start_idx, chunk_size): - for p in ordered_params[start_idx:start_idx + chunk_size]: + _compute_u(state, group["ns_steps"], self.rank, self.compute_stream) + + def enqueue_scatters(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: state = param_to_state[id(p)] adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) - _update_param(p, state, lr, adjusted_lr, weight_decay, - self.rank, self.compute_stream) + _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream) - chunk_size = dist.get_world_size(param_to_state[id( - params[0])].process_group) + chunk_size = params[0].device_mesh.mesh.numel() # Wait grad update self.comm_stream.wait_stream(torch.cuda.current_stream()) - overlap_step = self.overlap_step - for i in range(0, overlap_step): - enqueue_all2all_gather(i * chunk_size, chunk_size) - enqueue_computes(i * chunk_size, chunk_size) - + enqueue_gathers(0, chunk_size) for i in range(0, len(params) + chunk_size - 1, chunk_size): - enqueue_all2all_scatter(i, chunk_size) - enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size) - enqueue_update_param(i, chunk_size) - enqueue_computes(i + overlap_step * chunk_size, chunk_size) - - # Wait the last update_param to finish - torch.cuda.current_stream().wait_stream(self.compute_stream) - - @staticmethod - def _fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: Union[float, torch.Tensor], - weight_decay: float, - eps: float, - maximize: bool, - ) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: Optional[DeviceDict] = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else - None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [ - params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps - ] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, - non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) + enqueue_computes(i, chunk_size) + enqueue_gathers(i + chunk_size, chunk_size) + enqueue_scatters(i, chunk_size) - def step(self, closure=None, qk_logits=None): + torch.cuda.current_stream().wait_stream(self.comm_stream) + + def step(self, closure=None): """Perform a single optimization step. Args: closure (Callable, optional): A closure that reevaluates the model and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). """ loss = None if closure is not None: @@ -938,127 +392,64 @@ class Muon(torch.optim.Optimizer): loss = closure() for group in self.param_groups: - params = group["params"] - - if group["use_muon"]: - ############################ - # Muon # - ############################ - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - param_dtensors = [] - param_tensors = [] - name_dtensors = [] - name_tensors = [] - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - param_tensors.append(p) - name_tensors.append(n) - else: - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError( - f"Unsupported parameter type: {type(p.data)}") - - if self.debug: - print( - f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors", - flush=True, - ) - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - self.parallel( - name_dtensors, - param_dtensors, - group, - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - qk_logits=qk_logits, - ) - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - qk_logits=qk_logits, - ) - + ############################ + # Muon # + ############################ + + params = [p for p in group["params"] if self.state[p]["use_muon"]] + lr = group["lr"] + wd = group["wd"] + momentum = group["momentum"] + + if isinstance(params[0].data, DTensor): + self.parallel( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) else: - ############################ - # AdamW backup # - ############################ - - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - max_exp_avg_sqs = [] - state_steps = [] - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - for p in params: - g = p.grad - if g is None: - continue - state = self.state[p] - params_with_grads.append(p) - grads.append(g) - if "step" not in state: - state["step"] = (torch.zeros((), - dtype=torch.float32, - device=p.device)) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(state["moment1"]) - moment2.append(state["moment2"]) - if not isinstance(state["step"], torch.Tensor): - step_tensor = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - else: - step_tensor = state["step"] - state_steps.append(step_tensor) - - self._fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - max_exp_avg_sqs, - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, + self.base( + params, + group, lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, + wd=wd, + momentum=momentum, ) + ############################ + # AdamW backup # + ############################ + + params = [p for p in group["params"] if not self.state[p]["use_muon"]] + lr = group["lr"] + beta1, beta2 = group["adamw_betas"] + eps = group["adamw_eps"] + weight_decay = group["wd"] + + for p in params: + g = p.grad + if g is None: + continue + state = self.state[p] + if "step" not in state: + state["step"] = 0 + state["moment1"] = torch.zeros_like(g) + state["moment2"] = torch.zeros_like(g) + state["step"] += 1 + step = state["step"] + buf1 = state["moment1"] + buf2 = state["moment2"] + buf1.lerp_(g, 1 - beta1) + buf2.lerp_(g.square(), 1 - beta2) + + g = buf1 / (eps + buf2.sqrt()) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + scale = bias_correction1 / bias_correction2**0.5 + p.data.mul_(1 - lr * weight_decay) + p.data.add_(g, alpha=-lr / scale) + return loss diff --git a/build/torch27-cxx11-cu126-x86_64-linux/optimizer/__init__.py b/build/torch27-cxx11-cu126-x86_64-linux/optimizer/__init__.py old mode 100644 new mode 100755 diff --git a/build/torch27-cxx11-cu126-x86_64-linux/optimizer/_ops.py b/build/torch27-cxx11-cu126-x86_64-linux/optimizer/_ops.py old mode 100644 new mode 100755 index cb9efd677b388ebc299d6c4747eee701c96211f6..7cf68eab4638da3512b5e49541c916ebd12301f0 --- a/build/torch27-cxx11-cu126-x86_64-linux/optimizer/_ops.py +++ b/build/torch27-cxx11-cu126-x86_64-linux/optimizer/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _optimizer_b0230e7_dirty -ops = torch.ops._optimizer_b0230e7_dirty +from . import _optimizer_036642a_dirty +ops = torch.ops._optimizer_036642a_dirty def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_optimizer_b0230e7_dirty::{op_name}" \ No newline at end of file + return f"_optimizer_036642a_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu126-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so b/build/torch27-cxx11-cu126-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..6c71757d9ed8f7ccfae12cd8eb6837ff15c9f773 --- /dev/null +++ b/build/torch27-cxx11-cu126-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:50cb5819ff08a2179d78cd98164d07fd3cef1b66ee7703d599a310dfb140b9d1 +size 1824256 diff --git a/build/torch27-cxx11-cu126-x86_64-linux/optimizer/_optimizer_b0230e7_dirty.abi3.so b/build/torch27-cxx11-cu126-x86_64-linux/optimizer/_optimizer_b0230e7_dirty.abi3.so deleted file mode 100755 index 7ccc064e07fe7031403165bb3c78e31e84cbdf19..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu126-x86_64-linux/optimizer/_optimizer_b0230e7_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:69525fcbfbe640264f4d52c9843b395b17f1828d38e1eceb97cec6bf46b0d8d0 -size 1824256 diff --git a/build/torch27-cxx11-cu126-x86_64-linux/optimizer/matmul_transpose_triton.py b/build/torch27-cxx11-cu126-x86_64-linux/optimizer/matmul_transpose_triton.py deleted file mode 100644 index 4565b2c4fd506a4218340d380d6c962b16774b1d..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu126-x86_64-linux/optimizer/matmul_transpose_triton.py +++ /dev/null @@ -1,128 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -def matmul_transpose_assign(d_in, d_out): - assert d_in.is_cuda, "Input `d_in` must be a CUDA tensor" - assert d_out.is_cuda, "Input `d_out` must be a CUDA tensor" - assert d_in.device == d_out.device, "Inputs `d_in` and `d_out` must be on the same CUDA device" - assert d_in.dtype == d_out.dtype, "Inputs must have the same data type" - assert d_in.ndim == 2, "Input `d_in` must be a 2D tensor" - assert d_out.ndim == 2, "Input `d_out` must be a 2D tensor" - assert d_in.size(0) == d_out.size(0) == d_out.size(0), \ - "First dimension of `d_in` must match first and second dimension of `d_out`" - - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -def matmul_transpose(d_in): - M, _ = d_in.shape - d_out = torch.empty((M, M), device=d_in.device, dtype=d_in.dtype) - matmul_transpose_assign(d_in, d_out) - return d_out diff --git a/build/torch27-cxx11-cu126-x86_64-linux/optimizer/muon.py b/build/torch27-cxx11-cu126-x86_64-linux/optimizer/muon.py old mode 100644 new mode 100755 index 4af25d55c528fb0db2272540838ab70eb3619194..0d614d55d721efac406c147b4f62e6c703a91107 --- a/build/torch27-cxx11-cu126-x86_64-linux/optimizer/muon.py +++ b/build/torch27-cxx11-cu126-x86_64-linux/optimizer/muon.py @@ -1,26 +1,14 @@ -import logging import math -import types from dataclasses import dataclass -from typing import List, Optional, Union, cast import torch import torch.distributed as dist -from torch.distributed._tensor import DTensor, Replicate, Shard - -from .matmul_transpose_triton import matmul_transpose_assign - -logger = logging.getLogger(__name__) - -COMM_DTYPE = torch.bfloat16 +from torch.distributed._tensor import DTensor # This code snippet is a modified version adapted from the following GitHub repositories: # https://github.com/KellerJordan/Muon/blob/master/muon.py -# Muon's Newton–Schulz iteration causes high variance in singular values -# Idea: give each iteration its own 3 coefficients and optimize them via gradient descent. @torch.no_grad() -# matmul_transpose_assign from : https://github.com/nil0x9/flash-muon def _zeropower_via_newtonschulz5(G, steps): """ Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a @@ -32,31 +20,26 @@ def _zeropower_via_newtonschulz5(G, steps): performance at all relative to UV^T, where USV^T = G is the SVD. """ assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE + a, b, c = (3.4445, -4.7750, 2.0315) X = G # no manual typecast - if G.size(0) > G.size(1): X = X.T # Ensure spectral norm is at most 1 X = X / (X.norm() + 1e-7) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) + X = X.bfloat16() # Perform the NS iterations - for a, b, c in [ - (4.0848, -6.8946, 2.9270), - (3.9505, -6.3029, 2.6377), - (3.7418, -5.5913, 2.3037), - (2.8769, -3.1427, 1.2046), - (2.8366, -3.0525, 1.2012), - ]: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) + for _ in range(steps): + A = X @ X.T + # B = ( + # b * A + c * A @ A + # ) + B = torch.addmm(A, A, A, alpha=c, beta=b) + # X = a * X + B @ X + X = torch.addmm(X, B, X, alpha=1.0, beta=a) if G.size(0) > G.size(1): X = X.T - return X + return X.to(G.dtype) @dataclass @@ -64,425 +47,92 @@ class _muon_state: # TODO: use Optional worker_rank: int | None = None gathered_grad: torch.Tensor | None = None - scattered_u: DTensor | None = None computed_u: torch.Tensor | None = None gather_event: torch.cuda.Event | None = None compute_event: torch.cuda.Event | None = None - scatter_event: torch.cuda.Event | None = None - process_group = None - qk_clip_state = None - - -def split_elems_for_src(param, src_rank, num_ranks) -> int: - rows = param.shape[0] - cols = int(param.numel() // rows) - base, rem = divmod(rows, num_ranks) - my_rows = base + (1 if src_rank < rem else 0) - return my_rows * cols @torch.no_grad() -def _alloc_gathered_grad(params, param_to_state, rank, compute_stream): - """ - Pre-allocate gathered_grad buffer on compute_stream - before launching all2all gather - """ - with torch.cuda.stream(compute_stream): - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - num_ranks = dist.get_world_size(group=state.process_group) - state.gathered_grad = torch.empty(p.grad.numel(), - dtype=COMM_DTYPE, - device="cuda") - else: - state.gathered_grad = None - - alloc_event = torch.cuda.Event() - alloc_event.record(compute_stream) - return alloc_event +def _gather(p, state, rank, comm_stream, none_grad): + g = p.grad + mesh = g.device_mesh + if rank == state.worker_rank: + gather_list = [torch.empty_like(g.to_local()) for _ in range(mesh.mesh.numel())] + else: + gather_list = None -@torch.no_grad() -def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad, - alloc_event): - """ - All2all gathers shards so each owner rank reconstructs its full gradient - """ with torch.cuda.stream(comm_stream): - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - - # Construct sending buffers - per_dst = [[] for _ in range(num_ranks)] - send_counts = [0] * num_ranks - - for p in params: - state = param_to_state[id(p)] - dst = state.worker_rank - assert dst < num_ranks - shard_elems = split_elems_for_src(p, rank, num_ranks) - g = p.grad - g = g.to_local().to(COMM_DTYPE).contiguous().view(-1) - assert g.numel() == shard_elems - per_dst[dst].append(g) - send_counts[dst] += shard_elems - - assert any( - len(v) > 0 for v in per_dst - ), "At least one destination rank must receive a sharded tensor" - # list[list[Tensor]] -> list[Tensor] - per_dst = [t for dst in per_dst for t in dst] - - send_buf = torch.cat(per_dst, dim=0) - - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - # Compute receive sizes and allocate receiving buffers - recv_counts = [0] * num_ranks - - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += split_elems_for_src(p, src, num_ranks) - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - #All2All - dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, + torch.distributed.gather( + g.to_local(), + dst=state.worker_rank, + gather_list=gather_list, + group=mesh.get_group(), ) - - # Reconstructs gathered grad from the received buffer - # - # recv_buf (num ranks = 3) - # - # From rank 0 From rank 1 From rank 2 - # | p1_0, p2_0, p3_0 | p1_1, p2_1, p3_1 | p1_2, p2_2, p3_2 | - # - # Outer loop: - # rank 0 -> rank 1 -> rank2 - # - # Inner loop: - # p1_n -> p2_n -> p3_n - - comm_stream.wait_event(alloc_event) - - off = 0 - write_offsets = {id(p): 0 for p in owned_params} - for src in range(num_ranks): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - n = split_elems_for_src(p, src, num_ranks) - assert n > 0 - - sg = recv_buf.narrow(0, off + inner_off, n) - woff = write_offsets[id(p)] - dst = state.gathered_grad.narrow(0, woff, n) - dst.copy_(sg) - - write_offsets[id(p)] += n - inner_off += n - off += block - - for p in params: - state = param_to_state[id(p)] - if state.worker_rank == rank: - state.gathered_grad = state.gathered_grad.view_as(p) - state.gather_event = torch.cuda.Event() - state.gather_event.record(comm_stream) - else: - state.gathered_grad = None - state.gather_event = None - if none_grad: - p.grad = None + if rank == state.worker_rank: + if state.gathered_grad is not None: + raise RuntimeError( + "Gather event already exists, which should not happen." + ) + state.gathered_grad = torch.cat(gather_list, dim=0) + state.gather_event = torch.cuda.Event() + state.gather_event.record() + else: + state.gathered_grad = None + state.gather_event = None + if none_grad: + p.grad = None @torch.no_grad() -def _compute_u(p, state, steps, rank, compute_stream): - """ - On worker_rank, compute the orthogonalized update using Newton-Schulz iteration. - """ +def _compute_u(state, steps, rank, compute_stream): with torch.cuda.stream(compute_stream): if rank == state.worker_rank: if state.gather_event is None: raise RuntimeError("Gather event must be set before compute.") compute_stream.wait_event(state.gather_event) u = _zeropower_via_newtonschulz5(state.gathered_grad, steps) - state.gathered_grad = None state.computed_u = u state.compute_event = torch.cuda.Event() state.compute_event.record() + # Clear the gathered gradient to free memory + state.gathered_grad = None else: state.computed_u = None state.compute_event = None @torch.no_grad() -def _alloc_scattered_u(params, param_to_state, rank, compute_stream): - """ - Pre-allocate scattered_u buffer on compute_stream - before launching all2all gather - """ - with torch.cuda.stream(compute_stream): - for p in params: - state = param_to_state[id(p)] - state.scattered_u = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - alloc_event = torch.cuda.Event() - alloc_event.record(compute_stream) - return alloc_event - +def _scatter(p, state, lr, wd, rank, comm_stream): + u = state.computed_u + mesh = p.device_mesh -def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event): - """ - All2all scatters full gradients to all ranks - """ with torch.cuda.stream(comm_stream): - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - # Construct sending buffer - per_dst = [[] for _ in range(num_ranks)] - send_counts = [0] * num_ranks - - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - if state.compute_event is None: - raise RuntimeError( - "Compute event must be set before scatter.") - comm_stream.wait_event(state.compute_event) - state.gathered_grad = None - - assert state.computed_u is not None - - u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1) - - offset = 0 - for dst in range(num_ranks): - n = split_elems_for_src(p, dst, num_ranks) - assert n > 0 - - su = u_full.narrow(0, offset, n) - per_dst[dst].append(su) - send_counts[dst] += n - offset += n - - assert offset == u_full.numel() - - lengths = [len(v) for v in per_dst] - if all(l > 0 for l in lengths): - assert all( - l == lengths[0] for l in lengths - ), "All destination ranks must have the same number of sharded tensor" - # list[list[Tensor]] -> list[Tensor] - per_dst = [t for dst in per_dst for t in dst] - send_buf = torch.cat(per_dst, dim=0) + if rank == state.worker_rank: + if state.compute_event is None: + raise RuntimeError("Compute event must be set before scatter.") + comm_stream.wait_event(state.compute_event) + scatter_list = list(torch.split(u, p.size(0) // mesh.mesh.numel(), dim=0)) else: - # all_to_all requires participation from all ranks - # Even non-owner ranks must join the collective call - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Compute receive sizes and allocate receiving buffers - recv_counts = [0] * num_ranks - - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += split_elems_for_src(p, rank, num_ranks) - recv_counts[src] = total - - recv_total = sum(recv_counts) - assert recv_total > 0 - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - #All2All - dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, + scatter_list = None + + u = torch.empty_like(p.to_local()) + torch.distributed.scatter( + u, + scatter_list=scatter_list, + src=state.worker_rank, + group=mesh.get_group(), ) - - # Copy to pre-allocated scattered_u buffer from the received buffer - # - # recv_buf (num ranks = 3, local_rank = 0) - # - # From rank 0 From rank 1 From rank 2 - # | p1_0, p2_0, p3_0 | p4_0 | p5_0, p6_0 | - # - # Outer loop: - # rank 0 -> rank 1 -> rank2 - # - # Inner loop: - # src(0) : p1_0 -> p2_0 -> p3_0 - # src(1) : p4_0 - # src(2) : p5_0 -> p6_0 - - comm_stream.wait_event(alloc_event) - - off = 0 - for src in range(num_ranks): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = split_elems_for_src(p, rank, num_ranks) - assert n > 0 - - flat_local = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - state.scattered_u.copy_(flat_local) - - state.scatter_event = torch.cuda.Event() - state.scatter_event.record(comm_stream) - inner_off += n - - assert inner_off == block - off += block - - -def _update_param(p, state, lr, adjusted_lr, weight_decay, rank, - compute_stream): - """ - Update sharded parameter p with the scattered_u. - Only worker_rank frees computed_u. - """ - with torch.cuda.stream(compute_stream): - if state.scatter_event is None: - raise RuntimeError("Scatter event must be set before update") - compute_stream.wait_event(state.scatter_event) - u_dtensor = DTensor.from_local( - state.scattered_u, - placements=p.placements, - device_mesh=p.device_mesh, - ) - - state.scattered_u = u_dtensor - if rank == state.worker_rank: - # Free computed_u + # Clear u to free memory state.computed_u = None - - Muon._update_p(p, state.scattered_u, lr, adjusted_lr, weight_decay) - state.scattered_u = None - u_dtensor = None - - scales_full = Muon._compute_scales(p, state.qk_clip_state) - if scales_full is not None: - num_ranks = dist.get_world_size(group=state.process_group) - local_rank = dist.get_rank(group=state.process_group) - scales_local = scales_full.chunk(num_ranks, dim=0)[local_rank] - scales_local = DTensor.from_local( - scales_local, - placements=p.placements, - device_mesh=p.device_mesh, - ) - Muon._qk_clip(p, scales_local, state.qk_clip_state.head_dim) - - -def default_is_muon(name, x): - skip_keys = ["embed_tokens", "lm_head", "tok_embeddings", "output"] - return x.ndim >= 2 and not any(key in name for key in skip_keys) - - -def get_default_muon_param_groups(model, is_muon_func=default_is_muon): - muon_params, muon_names = [], [] - non_muon_params = [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - ('wq', 'wk', 'q_proj', 'k_proj') and return (kind, layer_index). - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = name.split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: Optional[str] # 'wq'/'q_proj' or 'wk'/'k_proj' or None - indices: List[int] # which heads to consider for clipping - head_dim: int # from config - threshold: float # from config - logit: Optional[torch.Tensor] + u = DTensor.from_local( + u, + placements=p.placements, + device_mesh=mesh, + ) + p.data.mul_(1 - lr * wd) + p.data.add_(u, alpha=-lr) class Muon(torch.optim.Optimizer): @@ -499,87 +149,71 @@ class Muon(torch.optim.Optimizer): - We believe it may not work well for finetuning pretrained models, but we haven't tested this. Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. + muon_params: The parameters to be optimized by Muon. lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) momentum: The momentum used by the internal SGD. (0.95 is a good default) nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. + adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. adamw_lr: The learning rate for the internal AdamW. adamw_betas: The betas for the internal AdamW. adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - overlap_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher overlap_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. + adamw_wd: The weight decay for the internal AdamW. """ - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config={ - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - }, - overlap_step=5): + def __init__( + self, + model, + is_muon_func, + lr=1e-3, + momentum=0.95, + nesterov=True, + ns_steps=5, + adamw_wd=0.1, + adamw_betas=(0.9, 0.95), + adamw_eps=1e-8, + none_grad=True, + debug=False, + ): defaults = dict( lr=lr, - weight_decay=weight_decay, + wd=adamw_wd, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps, adamw_betas=adamw_betas, adamw_eps=adamw_eps, none_grad=none_grad, - use_muon=True, ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) + super().__init__(model.parameters(), defaults) + self.is_muon_func = is_muon_func + self.model = model - super().__init__(params, defaults) + if not dist.is_initialized(): + raise RuntimeError( + "Muon optimizer requires distributed training to be initialized." + ) - self.rank = None + self.rank = dist.get_rank() self.comm_stream = torch.cuda.Stream() self.compute_stream = torch.cuda.Stream() self.debug = debug - self.clip_config = clip_config - self.overlap_step = overlap_step + + def __setstate__(self, state): + # Sort parameters into those for which we will use Muon, and those for which we will not + super().__setstate__(state) + for name, p in self.model.named_parameters(): + if self.is_muon_func(p, name): + # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer + assert p.ndim == 2, p.ndim + self.state[p]["use_muon"] = True + self.state[p]["orig_shape"] = p.shape + else: + # Do not use Muon for parameters in adamw_params + self.state[p]["use_muon"] = False def _calc_flops(self, G, steps): assert len(G.shape) == 2 @@ -597,30 +231,7 @@ class Muon(torch.optim.Optimizer): adjusted_lr = lr * adjusted_ratio return adjusted_lr - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - if p.placements == (Shard(dim=0), ): - # Case for FSDP - return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0) - elif p.placements == (Replicate(), Shard(dim=0)): - # Case for HSDP - process_group = p.device_mesh.get_group(mesh_dim=1) - if self.rank is None: - self.rank = dist.get_rank(group=process_group) - else: - assert self.rank == dist.get_rank(group=process_group) - for i, shard_mesh in enumerate(p.device_mesh.mesh): - if self.rank in shard_mesh: - return shard_mesh, p.device_mesh.get_group(mesh_dim=1) - else: - raise ValueError(f"Unsupported placements ({p.placements}).") - - def init_state_and_assign_params(self, names, params, group, qk_logits): + def init_state_and_assign_params(self, params, group): param_to_state = {} param_to_flops = {} @@ -636,44 +247,34 @@ class Muon(torch.optim.Optimizer): total_flops += flops if self.debug: - print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", - flush=True) - - paired = list(zip(names, params)) + print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", flush=True) - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) + ordered_params = sorted( + params, key=lambda p: param_to_flops[id(p)], reverse=True + ) round_robin = 0 mesh = None - shard_mesh = None - process_group = None - for n, p in zip(ordered_names, ordered_params): + for p in ordered_params: if mesh is None: mesh = p.device_mesh - shard_mesh, process_group = self.get_shard_mesh(p) + if mesh.ndim != 1: + raise NotImplementedError( + "Muon requires a 1D mesh for distributed training yet." + ) elif mesh != p.device_mesh: raise ValueError("All parameters must be on the same mesh.") - num_ranks = dist.get_world_size(group=process_group) + param_to_state[id(p)] = _muon_state() - param_to_state[id( - p)].worker_rank = shard_mesh[round_robin].item() % num_ranks - param_to_state[id(p)].process_group = process_group - qk_clip_state = self.get_qk_clip_info(n, qk_logits) - param_to_state[id(p)].qk_clip_state = qk_clip_state - round_robin = (round_robin + 1) % len(shard_mesh) + param_to_state[id(p)].worker_rank = mesh.mesh[round_robin].item() + + round_robin = (round_robin + 1) % mesh.mesh.numel() return param_to_state, ordered_params - def base(self, names, params, group, lr, weight_decay, momentum, - qk_logits): + def base(self, params, group, lr, wd, momentum): # generate weight updates in distributed fashion - for n, p in zip(names, params): + for p in params: g = p.grad if g is None: continue @@ -692,87 +293,39 @@ class Muon(torch.optim.Optimizer): else: g = buf - u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) + u = _zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) + # scale update adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) - Muon._update_p(p, u, lr, adjusted_lr, weight_decay) - qk_clip_state = self.get_qk_clip_info(n, qk_logits) + # apply weight decay + p.data.mul_(1 - lr * wd) - scales_full = self._compute_scales(p, qk_clip_state) - if scales_full is not None: - Muon._qk_clip(p, scales_full, qk_clip_state.head_dim) + # apply update + p.data.add_(u, alpha=-adjusted_lr) def _update_g(self, p, g, group, momentum): # calc update state = self.state[p] - buf = state.setdefault("momentum_buffer", torch.zeros_like(g)) - torch.add(g, buf, alpha=momentum, out=buf) + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) if group["nesterov"]: - g.add_(buf, alpha=momentum) - return g - return buf + g = g.add(buf, alpha=momentum) + else: + g = buf + return g - @staticmethod - def _update_p(p, u, lr, adjusted_lr, weight_decay): + def _update_p(self, p, u, lr, wd): + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) # apply weight decay - p.data.mul_(1 - lr * weight_decay) + p.data.mul_(1 - lr * wd) # apply update p.data.add_(u, alpha=-adjusted_lr) - def get_qk_clip_info(self, n, qk_logits): - head_dim = self.clip_config.get('head_dim') - threshold = self.clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - indices_key = 'q_indices' if 'q' in kind else 'k_indices' - indices = self.clip_config.get(indices_key, []) or [] - - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - @staticmethod - def _compute_scales(p, qk_clip_state): - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - H_global = p.shape[0] // head_dim - scales_full = torch.ones(H_global, device=p.data.device) - scaling = 0 - - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if new_scale < scales_full[head_idx]: - scales_full[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - scaling += 1 - - return scales_full if scaling > 0 else None - - @staticmethod - def _qk_clip(p, scales, head_dim): - W = p.data.view(-1, head_dim, p.data.shape[1]) - W.mul_(scales.view(-1, 1, 1)) - - def parallel(self, names, params, group, lr, weight_decay, momentum, - qk_logits): + def parallel(self, params, group, lr, wd, momentum): """ Perform a parallel optimization step using Muon. """ @@ -794,143 +347,44 @@ class Muon(torch.optim.Optimizer): p.grad = g param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - assert self.rank is not None + params, group + ) - def enqueue_all2all_gather(start_idx, chunk_size): - target_params = ordered_params[start_idx:start_idx + chunk_size] - if target_params: - alloc_event = _alloc_gathered_grad(target_params, - param_to_state, self.rank, - self.compute_stream) - _all2all_gather(target_params, param_to_state, self.rank, - self.comm_stream, group["none_grad"], - alloc_event) + def enqueue_gathers(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _gather(p, state, self.rank, self.comm_stream, group["none_grad"]) def enqueue_computes(start_idx, chunk_size): - for p in ordered_params[start_idx:start_idx + chunk_size]: + for p in ordered_params[start_idx : start_idx + chunk_size]: state = param_to_state[id(p)] - _compute_u(p, state, group["ns_steps"], self.rank, - self.compute_stream) - - def enqueue_all2all_scatter(start_idx, chunk_size): - target_params = ordered_params[start_idx:start_idx + chunk_size] - if target_params: - alloc_event = _alloc_scattered_u(target_params, param_to_state, - self.rank, - self.compute_stream) - _all2all_scatter(target_params, param_to_state, self.rank, - self.comm_stream, alloc_event) - - def enqueue_update_param(start_idx, chunk_size): - for p in ordered_params[start_idx:start_idx + chunk_size]: + _compute_u(state, group["ns_steps"], self.rank, self.compute_stream) + + def enqueue_scatters(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: state = param_to_state[id(p)] adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) - _update_param(p, state, lr, adjusted_lr, weight_decay, - self.rank, self.compute_stream) + _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream) - chunk_size = dist.get_world_size(param_to_state[id( - params[0])].process_group) + chunk_size = params[0].device_mesh.mesh.numel() # Wait grad update self.comm_stream.wait_stream(torch.cuda.current_stream()) - overlap_step = self.overlap_step - for i in range(0, overlap_step): - enqueue_all2all_gather(i * chunk_size, chunk_size) - enqueue_computes(i * chunk_size, chunk_size) - + enqueue_gathers(0, chunk_size) for i in range(0, len(params) + chunk_size - 1, chunk_size): - enqueue_all2all_scatter(i, chunk_size) - enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size) - enqueue_update_param(i, chunk_size) - enqueue_computes(i + overlap_step * chunk_size, chunk_size) - - # Wait the last update_param to finish - torch.cuda.current_stream().wait_stream(self.compute_stream) - - @staticmethod - def _fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: Union[float, torch.Tensor], - weight_decay: float, - eps: float, - maximize: bool, - ) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: Optional[DeviceDict] = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else - None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [ - params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps - ] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, - non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) + enqueue_computes(i, chunk_size) + enqueue_gathers(i + chunk_size, chunk_size) + enqueue_scatters(i, chunk_size) - def step(self, closure=None, qk_logits=None): + torch.cuda.current_stream().wait_stream(self.comm_stream) + + def step(self, closure=None): """Perform a single optimization step. Args: closure (Callable, optional): A closure that reevaluates the model and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). """ loss = None if closure is not None: @@ -938,127 +392,64 @@ class Muon(torch.optim.Optimizer): loss = closure() for group in self.param_groups: - params = group["params"] - - if group["use_muon"]: - ############################ - # Muon # - ############################ - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - param_dtensors = [] - param_tensors = [] - name_dtensors = [] - name_tensors = [] - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - param_tensors.append(p) - name_tensors.append(n) - else: - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError( - f"Unsupported parameter type: {type(p.data)}") - - if self.debug: - print( - f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors", - flush=True, - ) - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - self.parallel( - name_dtensors, - param_dtensors, - group, - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - qk_logits=qk_logits, - ) - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - qk_logits=qk_logits, - ) - + ############################ + # Muon # + ############################ + + params = [p for p in group["params"] if self.state[p]["use_muon"]] + lr = group["lr"] + wd = group["wd"] + momentum = group["momentum"] + + if isinstance(params[0].data, DTensor): + self.parallel( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) else: - ############################ - # AdamW backup # - ############################ - - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - max_exp_avg_sqs = [] - state_steps = [] - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - for p in params: - g = p.grad - if g is None: - continue - state = self.state[p] - params_with_grads.append(p) - grads.append(g) - if "step" not in state: - state["step"] = (torch.zeros((), - dtype=torch.float32, - device=p.device)) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(state["moment1"]) - moment2.append(state["moment2"]) - if not isinstance(state["step"], torch.Tensor): - step_tensor = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - else: - step_tensor = state["step"] - state_steps.append(step_tensor) - - self._fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - max_exp_avg_sqs, - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, + self.base( + params, + group, lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, + wd=wd, + momentum=momentum, ) + ############################ + # AdamW backup # + ############################ + + params = [p for p in group["params"] if not self.state[p]["use_muon"]] + lr = group["lr"] + beta1, beta2 = group["adamw_betas"] + eps = group["adamw_eps"] + weight_decay = group["wd"] + + for p in params: + g = p.grad + if g is None: + continue + state = self.state[p] + if "step" not in state: + state["step"] = 0 + state["moment1"] = torch.zeros_like(g) + state["moment2"] = torch.zeros_like(g) + state["step"] += 1 + step = state["step"] + buf1 = state["moment1"] + buf2 = state["moment2"] + buf1.lerp_(g, 1 - beta1) + buf2.lerp_(g.square(), 1 - beta2) + + g = buf1 / (eps + buf2.sqrt()) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + scale = bias_correction1 / bias_correction2**0.5 + p.data.mul_(1 - lr * weight_decay) + p.data.add_(g, alpha=-lr / scale) + return loss diff --git a/build/torch27-cxx11-cu128-x86_64-linux/optimizer/__init__.py b/build/torch27-cxx11-cu128-x86_64-linux/optimizer/__init__.py old mode 100644 new mode 100755 diff --git a/build/torch27-cxx11-cu128-x86_64-linux/optimizer/_ops.py b/build/torch27-cxx11-cu128-x86_64-linux/optimizer/_ops.py old mode 100644 new mode 100755 index cb9efd677b388ebc299d6c4747eee701c96211f6..7cf68eab4638da3512b5e49541c916ebd12301f0 --- a/build/torch27-cxx11-cu128-x86_64-linux/optimizer/_ops.py +++ b/build/torch27-cxx11-cu128-x86_64-linux/optimizer/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _optimizer_b0230e7_dirty -ops = torch.ops._optimizer_b0230e7_dirty +from . import _optimizer_036642a_dirty +ops = torch.ops._optimizer_036642a_dirty def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_optimizer_b0230e7_dirty::{op_name}" \ No newline at end of file + return f"_optimizer_036642a_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu128-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so b/build/torch27-cxx11-cu128-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..18e4199f161721ea032af3872fa5f7d3cd95a862 --- /dev/null +++ b/build/torch27-cxx11-cu128-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9c75e42265f382addc71327ad5628e8a2414da5872791c975e384708c4acd549 +size 1883352 diff --git a/build/torch27-cxx11-cu128-x86_64-linux/optimizer/_optimizer_b0230e7_dirty.abi3.so b/build/torch27-cxx11-cu128-x86_64-linux/optimizer/_optimizer_b0230e7_dirty.abi3.so deleted file mode 100755 index b49ddd7a4dd2c980f9693a477e26221042cc85c5..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-x86_64-linux/optimizer/_optimizer_b0230e7_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:331cc0bc5ee469afdfe0fc590bf52910c118cd0cec62ccbf85778c12ae367a95 -size 1883344 diff --git a/build/torch27-cxx11-cu128-x86_64-linux/optimizer/matmul_transpose_triton.py b/build/torch27-cxx11-cu128-x86_64-linux/optimizer/matmul_transpose_triton.py deleted file mode 100644 index 4565b2c4fd506a4218340d380d6c962b16774b1d..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-x86_64-linux/optimizer/matmul_transpose_triton.py +++ /dev/null @@ -1,128 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -def matmul_transpose_assign(d_in, d_out): - assert d_in.is_cuda, "Input `d_in` must be a CUDA tensor" - assert d_out.is_cuda, "Input `d_out` must be a CUDA tensor" - assert d_in.device == d_out.device, "Inputs `d_in` and `d_out` must be on the same CUDA device" - assert d_in.dtype == d_out.dtype, "Inputs must have the same data type" - assert d_in.ndim == 2, "Input `d_in` must be a 2D tensor" - assert d_out.ndim == 2, "Input `d_out` must be a 2D tensor" - assert d_in.size(0) == d_out.size(0) == d_out.size(0), \ - "First dimension of `d_in` must match first and second dimension of `d_out`" - - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -def matmul_transpose(d_in): - M, _ = d_in.shape - d_out = torch.empty((M, M), device=d_in.device, dtype=d_in.dtype) - matmul_transpose_assign(d_in, d_out) - return d_out diff --git a/build/torch27-cxx11-cu128-x86_64-linux/optimizer/muon.py b/build/torch27-cxx11-cu128-x86_64-linux/optimizer/muon.py old mode 100644 new mode 100755 index 4af25d55c528fb0db2272540838ab70eb3619194..0d614d55d721efac406c147b4f62e6c703a91107 --- a/build/torch27-cxx11-cu128-x86_64-linux/optimizer/muon.py +++ b/build/torch27-cxx11-cu128-x86_64-linux/optimizer/muon.py @@ -1,26 +1,14 @@ -import logging import math -import types from dataclasses import dataclass -from typing import List, Optional, Union, cast import torch import torch.distributed as dist -from torch.distributed._tensor import DTensor, Replicate, Shard - -from .matmul_transpose_triton import matmul_transpose_assign - -logger = logging.getLogger(__name__) - -COMM_DTYPE = torch.bfloat16 +from torch.distributed._tensor import DTensor # This code snippet is a modified version adapted from the following GitHub repositories: # https://github.com/KellerJordan/Muon/blob/master/muon.py -# Muon's Newton–Schulz iteration causes high variance in singular values -# Idea: give each iteration its own 3 coefficients and optimize them via gradient descent. @torch.no_grad() -# matmul_transpose_assign from : https://github.com/nil0x9/flash-muon def _zeropower_via_newtonschulz5(G, steps): """ Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a @@ -32,31 +20,26 @@ def _zeropower_via_newtonschulz5(G, steps): performance at all relative to UV^T, where USV^T = G is the SVD. """ assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE + a, b, c = (3.4445, -4.7750, 2.0315) X = G # no manual typecast - if G.size(0) > G.size(1): X = X.T # Ensure spectral norm is at most 1 X = X / (X.norm() + 1e-7) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) + X = X.bfloat16() # Perform the NS iterations - for a, b, c in [ - (4.0848, -6.8946, 2.9270), - (3.9505, -6.3029, 2.6377), - (3.7418, -5.5913, 2.3037), - (2.8769, -3.1427, 1.2046), - (2.8366, -3.0525, 1.2012), - ]: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) + for _ in range(steps): + A = X @ X.T + # B = ( + # b * A + c * A @ A + # ) + B = torch.addmm(A, A, A, alpha=c, beta=b) + # X = a * X + B @ X + X = torch.addmm(X, B, X, alpha=1.0, beta=a) if G.size(0) > G.size(1): X = X.T - return X + return X.to(G.dtype) @dataclass @@ -64,425 +47,92 @@ class _muon_state: # TODO: use Optional worker_rank: int | None = None gathered_grad: torch.Tensor | None = None - scattered_u: DTensor | None = None computed_u: torch.Tensor | None = None gather_event: torch.cuda.Event | None = None compute_event: torch.cuda.Event | None = None - scatter_event: torch.cuda.Event | None = None - process_group = None - qk_clip_state = None - - -def split_elems_for_src(param, src_rank, num_ranks) -> int: - rows = param.shape[0] - cols = int(param.numel() // rows) - base, rem = divmod(rows, num_ranks) - my_rows = base + (1 if src_rank < rem else 0) - return my_rows * cols @torch.no_grad() -def _alloc_gathered_grad(params, param_to_state, rank, compute_stream): - """ - Pre-allocate gathered_grad buffer on compute_stream - before launching all2all gather - """ - with torch.cuda.stream(compute_stream): - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - num_ranks = dist.get_world_size(group=state.process_group) - state.gathered_grad = torch.empty(p.grad.numel(), - dtype=COMM_DTYPE, - device="cuda") - else: - state.gathered_grad = None - - alloc_event = torch.cuda.Event() - alloc_event.record(compute_stream) - return alloc_event +def _gather(p, state, rank, comm_stream, none_grad): + g = p.grad + mesh = g.device_mesh + if rank == state.worker_rank: + gather_list = [torch.empty_like(g.to_local()) for _ in range(mesh.mesh.numel())] + else: + gather_list = None -@torch.no_grad() -def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad, - alloc_event): - """ - All2all gathers shards so each owner rank reconstructs its full gradient - """ with torch.cuda.stream(comm_stream): - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - - # Construct sending buffers - per_dst = [[] for _ in range(num_ranks)] - send_counts = [0] * num_ranks - - for p in params: - state = param_to_state[id(p)] - dst = state.worker_rank - assert dst < num_ranks - shard_elems = split_elems_for_src(p, rank, num_ranks) - g = p.grad - g = g.to_local().to(COMM_DTYPE).contiguous().view(-1) - assert g.numel() == shard_elems - per_dst[dst].append(g) - send_counts[dst] += shard_elems - - assert any( - len(v) > 0 for v in per_dst - ), "At least one destination rank must receive a sharded tensor" - # list[list[Tensor]] -> list[Tensor] - per_dst = [t for dst in per_dst for t in dst] - - send_buf = torch.cat(per_dst, dim=0) - - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - # Compute receive sizes and allocate receiving buffers - recv_counts = [0] * num_ranks - - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += split_elems_for_src(p, src, num_ranks) - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - #All2All - dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, + torch.distributed.gather( + g.to_local(), + dst=state.worker_rank, + gather_list=gather_list, + group=mesh.get_group(), ) - - # Reconstructs gathered grad from the received buffer - # - # recv_buf (num ranks = 3) - # - # From rank 0 From rank 1 From rank 2 - # | p1_0, p2_0, p3_0 | p1_1, p2_1, p3_1 | p1_2, p2_2, p3_2 | - # - # Outer loop: - # rank 0 -> rank 1 -> rank2 - # - # Inner loop: - # p1_n -> p2_n -> p3_n - - comm_stream.wait_event(alloc_event) - - off = 0 - write_offsets = {id(p): 0 for p in owned_params} - for src in range(num_ranks): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - n = split_elems_for_src(p, src, num_ranks) - assert n > 0 - - sg = recv_buf.narrow(0, off + inner_off, n) - woff = write_offsets[id(p)] - dst = state.gathered_grad.narrow(0, woff, n) - dst.copy_(sg) - - write_offsets[id(p)] += n - inner_off += n - off += block - - for p in params: - state = param_to_state[id(p)] - if state.worker_rank == rank: - state.gathered_grad = state.gathered_grad.view_as(p) - state.gather_event = torch.cuda.Event() - state.gather_event.record(comm_stream) - else: - state.gathered_grad = None - state.gather_event = None - if none_grad: - p.grad = None + if rank == state.worker_rank: + if state.gathered_grad is not None: + raise RuntimeError( + "Gather event already exists, which should not happen." + ) + state.gathered_grad = torch.cat(gather_list, dim=0) + state.gather_event = torch.cuda.Event() + state.gather_event.record() + else: + state.gathered_grad = None + state.gather_event = None + if none_grad: + p.grad = None @torch.no_grad() -def _compute_u(p, state, steps, rank, compute_stream): - """ - On worker_rank, compute the orthogonalized update using Newton-Schulz iteration. - """ +def _compute_u(state, steps, rank, compute_stream): with torch.cuda.stream(compute_stream): if rank == state.worker_rank: if state.gather_event is None: raise RuntimeError("Gather event must be set before compute.") compute_stream.wait_event(state.gather_event) u = _zeropower_via_newtonschulz5(state.gathered_grad, steps) - state.gathered_grad = None state.computed_u = u state.compute_event = torch.cuda.Event() state.compute_event.record() + # Clear the gathered gradient to free memory + state.gathered_grad = None else: state.computed_u = None state.compute_event = None @torch.no_grad() -def _alloc_scattered_u(params, param_to_state, rank, compute_stream): - """ - Pre-allocate scattered_u buffer on compute_stream - before launching all2all gather - """ - with torch.cuda.stream(compute_stream): - for p in params: - state = param_to_state[id(p)] - state.scattered_u = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - alloc_event = torch.cuda.Event() - alloc_event.record(compute_stream) - return alloc_event - +def _scatter(p, state, lr, wd, rank, comm_stream): + u = state.computed_u + mesh = p.device_mesh -def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event): - """ - All2all scatters full gradients to all ranks - """ with torch.cuda.stream(comm_stream): - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - # Construct sending buffer - per_dst = [[] for _ in range(num_ranks)] - send_counts = [0] * num_ranks - - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - if state.compute_event is None: - raise RuntimeError( - "Compute event must be set before scatter.") - comm_stream.wait_event(state.compute_event) - state.gathered_grad = None - - assert state.computed_u is not None - - u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1) - - offset = 0 - for dst in range(num_ranks): - n = split_elems_for_src(p, dst, num_ranks) - assert n > 0 - - su = u_full.narrow(0, offset, n) - per_dst[dst].append(su) - send_counts[dst] += n - offset += n - - assert offset == u_full.numel() - - lengths = [len(v) for v in per_dst] - if all(l > 0 for l in lengths): - assert all( - l == lengths[0] for l in lengths - ), "All destination ranks must have the same number of sharded tensor" - # list[list[Tensor]] -> list[Tensor] - per_dst = [t for dst in per_dst for t in dst] - send_buf = torch.cat(per_dst, dim=0) + if rank == state.worker_rank: + if state.compute_event is None: + raise RuntimeError("Compute event must be set before scatter.") + comm_stream.wait_event(state.compute_event) + scatter_list = list(torch.split(u, p.size(0) // mesh.mesh.numel(), dim=0)) else: - # all_to_all requires participation from all ranks - # Even non-owner ranks must join the collective call - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Compute receive sizes and allocate receiving buffers - recv_counts = [0] * num_ranks - - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += split_elems_for_src(p, rank, num_ranks) - recv_counts[src] = total - - recv_total = sum(recv_counts) - assert recv_total > 0 - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - #All2All - dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, + scatter_list = None + + u = torch.empty_like(p.to_local()) + torch.distributed.scatter( + u, + scatter_list=scatter_list, + src=state.worker_rank, + group=mesh.get_group(), ) - - # Copy to pre-allocated scattered_u buffer from the received buffer - # - # recv_buf (num ranks = 3, local_rank = 0) - # - # From rank 0 From rank 1 From rank 2 - # | p1_0, p2_0, p3_0 | p4_0 | p5_0, p6_0 | - # - # Outer loop: - # rank 0 -> rank 1 -> rank2 - # - # Inner loop: - # src(0) : p1_0 -> p2_0 -> p3_0 - # src(1) : p4_0 - # src(2) : p5_0 -> p6_0 - - comm_stream.wait_event(alloc_event) - - off = 0 - for src in range(num_ranks): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = split_elems_for_src(p, rank, num_ranks) - assert n > 0 - - flat_local = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - state.scattered_u.copy_(flat_local) - - state.scatter_event = torch.cuda.Event() - state.scatter_event.record(comm_stream) - inner_off += n - - assert inner_off == block - off += block - - -def _update_param(p, state, lr, adjusted_lr, weight_decay, rank, - compute_stream): - """ - Update sharded parameter p with the scattered_u. - Only worker_rank frees computed_u. - """ - with torch.cuda.stream(compute_stream): - if state.scatter_event is None: - raise RuntimeError("Scatter event must be set before update") - compute_stream.wait_event(state.scatter_event) - u_dtensor = DTensor.from_local( - state.scattered_u, - placements=p.placements, - device_mesh=p.device_mesh, - ) - - state.scattered_u = u_dtensor - if rank == state.worker_rank: - # Free computed_u + # Clear u to free memory state.computed_u = None - - Muon._update_p(p, state.scattered_u, lr, adjusted_lr, weight_decay) - state.scattered_u = None - u_dtensor = None - - scales_full = Muon._compute_scales(p, state.qk_clip_state) - if scales_full is not None: - num_ranks = dist.get_world_size(group=state.process_group) - local_rank = dist.get_rank(group=state.process_group) - scales_local = scales_full.chunk(num_ranks, dim=0)[local_rank] - scales_local = DTensor.from_local( - scales_local, - placements=p.placements, - device_mesh=p.device_mesh, - ) - Muon._qk_clip(p, scales_local, state.qk_clip_state.head_dim) - - -def default_is_muon(name, x): - skip_keys = ["embed_tokens", "lm_head", "tok_embeddings", "output"] - return x.ndim >= 2 and not any(key in name for key in skip_keys) - - -def get_default_muon_param_groups(model, is_muon_func=default_is_muon): - muon_params, muon_names = [], [] - non_muon_params = [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - ('wq', 'wk', 'q_proj', 'k_proj') and return (kind, layer_index). - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = name.split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: Optional[str] # 'wq'/'q_proj' or 'wk'/'k_proj' or None - indices: List[int] # which heads to consider for clipping - head_dim: int # from config - threshold: float # from config - logit: Optional[torch.Tensor] + u = DTensor.from_local( + u, + placements=p.placements, + device_mesh=mesh, + ) + p.data.mul_(1 - lr * wd) + p.data.add_(u, alpha=-lr) class Muon(torch.optim.Optimizer): @@ -499,87 +149,71 @@ class Muon(torch.optim.Optimizer): - We believe it may not work well for finetuning pretrained models, but we haven't tested this. Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. + muon_params: The parameters to be optimized by Muon. lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) momentum: The momentum used by the internal SGD. (0.95 is a good default) nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. + adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. adamw_lr: The learning rate for the internal AdamW. adamw_betas: The betas for the internal AdamW. adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - overlap_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher overlap_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. + adamw_wd: The weight decay for the internal AdamW. """ - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config={ - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - }, - overlap_step=5): + def __init__( + self, + model, + is_muon_func, + lr=1e-3, + momentum=0.95, + nesterov=True, + ns_steps=5, + adamw_wd=0.1, + adamw_betas=(0.9, 0.95), + adamw_eps=1e-8, + none_grad=True, + debug=False, + ): defaults = dict( lr=lr, - weight_decay=weight_decay, + wd=adamw_wd, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps, adamw_betas=adamw_betas, adamw_eps=adamw_eps, none_grad=none_grad, - use_muon=True, ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) + super().__init__(model.parameters(), defaults) + self.is_muon_func = is_muon_func + self.model = model - super().__init__(params, defaults) + if not dist.is_initialized(): + raise RuntimeError( + "Muon optimizer requires distributed training to be initialized." + ) - self.rank = None + self.rank = dist.get_rank() self.comm_stream = torch.cuda.Stream() self.compute_stream = torch.cuda.Stream() self.debug = debug - self.clip_config = clip_config - self.overlap_step = overlap_step + + def __setstate__(self, state): + # Sort parameters into those for which we will use Muon, and those for which we will not + super().__setstate__(state) + for name, p in self.model.named_parameters(): + if self.is_muon_func(p, name): + # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer + assert p.ndim == 2, p.ndim + self.state[p]["use_muon"] = True + self.state[p]["orig_shape"] = p.shape + else: + # Do not use Muon for parameters in adamw_params + self.state[p]["use_muon"] = False def _calc_flops(self, G, steps): assert len(G.shape) == 2 @@ -597,30 +231,7 @@ class Muon(torch.optim.Optimizer): adjusted_lr = lr * adjusted_ratio return adjusted_lr - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - if p.placements == (Shard(dim=0), ): - # Case for FSDP - return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0) - elif p.placements == (Replicate(), Shard(dim=0)): - # Case for HSDP - process_group = p.device_mesh.get_group(mesh_dim=1) - if self.rank is None: - self.rank = dist.get_rank(group=process_group) - else: - assert self.rank == dist.get_rank(group=process_group) - for i, shard_mesh in enumerate(p.device_mesh.mesh): - if self.rank in shard_mesh: - return shard_mesh, p.device_mesh.get_group(mesh_dim=1) - else: - raise ValueError(f"Unsupported placements ({p.placements}).") - - def init_state_and_assign_params(self, names, params, group, qk_logits): + def init_state_and_assign_params(self, params, group): param_to_state = {} param_to_flops = {} @@ -636,44 +247,34 @@ class Muon(torch.optim.Optimizer): total_flops += flops if self.debug: - print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", - flush=True) - - paired = list(zip(names, params)) + print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", flush=True) - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) + ordered_params = sorted( + params, key=lambda p: param_to_flops[id(p)], reverse=True + ) round_robin = 0 mesh = None - shard_mesh = None - process_group = None - for n, p in zip(ordered_names, ordered_params): + for p in ordered_params: if mesh is None: mesh = p.device_mesh - shard_mesh, process_group = self.get_shard_mesh(p) + if mesh.ndim != 1: + raise NotImplementedError( + "Muon requires a 1D mesh for distributed training yet." + ) elif mesh != p.device_mesh: raise ValueError("All parameters must be on the same mesh.") - num_ranks = dist.get_world_size(group=process_group) + param_to_state[id(p)] = _muon_state() - param_to_state[id( - p)].worker_rank = shard_mesh[round_robin].item() % num_ranks - param_to_state[id(p)].process_group = process_group - qk_clip_state = self.get_qk_clip_info(n, qk_logits) - param_to_state[id(p)].qk_clip_state = qk_clip_state - round_robin = (round_robin + 1) % len(shard_mesh) + param_to_state[id(p)].worker_rank = mesh.mesh[round_robin].item() + + round_robin = (round_robin + 1) % mesh.mesh.numel() return param_to_state, ordered_params - def base(self, names, params, group, lr, weight_decay, momentum, - qk_logits): + def base(self, params, group, lr, wd, momentum): # generate weight updates in distributed fashion - for n, p in zip(names, params): + for p in params: g = p.grad if g is None: continue @@ -692,87 +293,39 @@ class Muon(torch.optim.Optimizer): else: g = buf - u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) + u = _zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) + # scale update adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) - Muon._update_p(p, u, lr, adjusted_lr, weight_decay) - qk_clip_state = self.get_qk_clip_info(n, qk_logits) + # apply weight decay + p.data.mul_(1 - lr * wd) - scales_full = self._compute_scales(p, qk_clip_state) - if scales_full is not None: - Muon._qk_clip(p, scales_full, qk_clip_state.head_dim) + # apply update + p.data.add_(u, alpha=-adjusted_lr) def _update_g(self, p, g, group, momentum): # calc update state = self.state[p] - buf = state.setdefault("momentum_buffer", torch.zeros_like(g)) - torch.add(g, buf, alpha=momentum, out=buf) + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) if group["nesterov"]: - g.add_(buf, alpha=momentum) - return g - return buf + g = g.add(buf, alpha=momentum) + else: + g = buf + return g - @staticmethod - def _update_p(p, u, lr, adjusted_lr, weight_decay): + def _update_p(self, p, u, lr, wd): + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) # apply weight decay - p.data.mul_(1 - lr * weight_decay) + p.data.mul_(1 - lr * wd) # apply update p.data.add_(u, alpha=-adjusted_lr) - def get_qk_clip_info(self, n, qk_logits): - head_dim = self.clip_config.get('head_dim') - threshold = self.clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - indices_key = 'q_indices' if 'q' in kind else 'k_indices' - indices = self.clip_config.get(indices_key, []) or [] - - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - @staticmethod - def _compute_scales(p, qk_clip_state): - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - H_global = p.shape[0] // head_dim - scales_full = torch.ones(H_global, device=p.data.device) - scaling = 0 - - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if new_scale < scales_full[head_idx]: - scales_full[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - scaling += 1 - - return scales_full if scaling > 0 else None - - @staticmethod - def _qk_clip(p, scales, head_dim): - W = p.data.view(-1, head_dim, p.data.shape[1]) - W.mul_(scales.view(-1, 1, 1)) - - def parallel(self, names, params, group, lr, weight_decay, momentum, - qk_logits): + def parallel(self, params, group, lr, wd, momentum): """ Perform a parallel optimization step using Muon. """ @@ -794,143 +347,44 @@ class Muon(torch.optim.Optimizer): p.grad = g param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - assert self.rank is not None + params, group + ) - def enqueue_all2all_gather(start_idx, chunk_size): - target_params = ordered_params[start_idx:start_idx + chunk_size] - if target_params: - alloc_event = _alloc_gathered_grad(target_params, - param_to_state, self.rank, - self.compute_stream) - _all2all_gather(target_params, param_to_state, self.rank, - self.comm_stream, group["none_grad"], - alloc_event) + def enqueue_gathers(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _gather(p, state, self.rank, self.comm_stream, group["none_grad"]) def enqueue_computes(start_idx, chunk_size): - for p in ordered_params[start_idx:start_idx + chunk_size]: + for p in ordered_params[start_idx : start_idx + chunk_size]: state = param_to_state[id(p)] - _compute_u(p, state, group["ns_steps"], self.rank, - self.compute_stream) - - def enqueue_all2all_scatter(start_idx, chunk_size): - target_params = ordered_params[start_idx:start_idx + chunk_size] - if target_params: - alloc_event = _alloc_scattered_u(target_params, param_to_state, - self.rank, - self.compute_stream) - _all2all_scatter(target_params, param_to_state, self.rank, - self.comm_stream, alloc_event) - - def enqueue_update_param(start_idx, chunk_size): - for p in ordered_params[start_idx:start_idx + chunk_size]: + _compute_u(state, group["ns_steps"], self.rank, self.compute_stream) + + def enqueue_scatters(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: state = param_to_state[id(p)] adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) - _update_param(p, state, lr, adjusted_lr, weight_decay, - self.rank, self.compute_stream) + _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream) - chunk_size = dist.get_world_size(param_to_state[id( - params[0])].process_group) + chunk_size = params[0].device_mesh.mesh.numel() # Wait grad update self.comm_stream.wait_stream(torch.cuda.current_stream()) - overlap_step = self.overlap_step - for i in range(0, overlap_step): - enqueue_all2all_gather(i * chunk_size, chunk_size) - enqueue_computes(i * chunk_size, chunk_size) - + enqueue_gathers(0, chunk_size) for i in range(0, len(params) + chunk_size - 1, chunk_size): - enqueue_all2all_scatter(i, chunk_size) - enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size) - enqueue_update_param(i, chunk_size) - enqueue_computes(i + overlap_step * chunk_size, chunk_size) - - # Wait the last update_param to finish - torch.cuda.current_stream().wait_stream(self.compute_stream) - - @staticmethod - def _fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: Union[float, torch.Tensor], - weight_decay: float, - eps: float, - maximize: bool, - ) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: Optional[DeviceDict] = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else - None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [ - params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps - ] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, - non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) + enqueue_computes(i, chunk_size) + enqueue_gathers(i + chunk_size, chunk_size) + enqueue_scatters(i, chunk_size) - def step(self, closure=None, qk_logits=None): + torch.cuda.current_stream().wait_stream(self.comm_stream) + + def step(self, closure=None): """Perform a single optimization step. Args: closure (Callable, optional): A closure that reevaluates the model and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). """ loss = None if closure is not None: @@ -938,127 +392,64 @@ class Muon(torch.optim.Optimizer): loss = closure() for group in self.param_groups: - params = group["params"] - - if group["use_muon"]: - ############################ - # Muon # - ############################ - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - param_dtensors = [] - param_tensors = [] - name_dtensors = [] - name_tensors = [] - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - param_tensors.append(p) - name_tensors.append(n) - else: - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError( - f"Unsupported parameter type: {type(p.data)}") - - if self.debug: - print( - f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors", - flush=True, - ) - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - self.parallel( - name_dtensors, - param_dtensors, - group, - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - qk_logits=qk_logits, - ) - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - qk_logits=qk_logits, - ) - + ############################ + # Muon # + ############################ + + params = [p for p in group["params"] if self.state[p]["use_muon"]] + lr = group["lr"] + wd = group["wd"] + momentum = group["momentum"] + + if isinstance(params[0].data, DTensor): + self.parallel( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) else: - ############################ - # AdamW backup # - ############################ - - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - max_exp_avg_sqs = [] - state_steps = [] - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - for p in params: - g = p.grad - if g is None: - continue - state = self.state[p] - params_with_grads.append(p) - grads.append(g) - if "step" not in state: - state["step"] = (torch.zeros((), - dtype=torch.float32, - device=p.device)) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(state["moment1"]) - moment2.append(state["moment2"]) - if not isinstance(state["step"], torch.Tensor): - step_tensor = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - else: - step_tensor = state["step"] - state_steps.append(step_tensor) - - self._fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - max_exp_avg_sqs, - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, + self.base( + params, + group, lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, + wd=wd, + momentum=momentum, ) + ############################ + # AdamW backup # + ############################ + + params = [p for p in group["params"] if not self.state[p]["use_muon"]] + lr = group["lr"] + beta1, beta2 = group["adamw_betas"] + eps = group["adamw_eps"] + weight_decay = group["wd"] + + for p in params: + g = p.grad + if g is None: + continue + state = self.state[p] + if "step" not in state: + state["step"] = 0 + state["moment1"] = torch.zeros_like(g) + state["moment2"] = torch.zeros_like(g) + state["step"] += 1 + step = state["step"] + buf1 = state["moment1"] + buf2 = state["moment2"] + buf1.lerp_(g, 1 - beta1) + buf2.lerp_(g.square(), 1 - beta2) + + g = buf1 / (eps + buf2.sqrt()) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + scale = bias_correction1 / bias_correction2**0.5 + p.data.mul_(1 - lr * weight_decay) + p.data.add_(g, alpha=-lr / scale) + return loss diff --git a/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/__init__.py b/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/__init__.py old mode 100644 new mode 100755 diff --git a/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/_ops.py b/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/_ops.py old mode 100644 new mode 100755 index cb9efd677b388ebc299d6c4747eee701c96211f6..7cf68eab4638da3512b5e49541c916ebd12301f0 --- a/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/_ops.py +++ b/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _optimizer_b0230e7_dirty -ops = torch.ops._optimizer_b0230e7_dirty +from . import _optimizer_036642a_dirty +ops = torch.ops._optimizer_036642a_dirty def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_optimizer_b0230e7_dirty::{op_name}" \ No newline at end of file + return f"_optimizer_036642a_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so b/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..73f6a2b6170eaf2e12e73fa4cf70c620ae925bdd --- /dev/null +++ b/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9a2363d4311d6a75fbcc03e6d4a71c73dae4d54e00a30135d25198d4078c6b0f +size 1749648 diff --git a/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/_optimizer_b0230e7_dirty.abi3.so b/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/_optimizer_b0230e7_dirty.abi3.so deleted file mode 100755 index 52c84f7911385aeed54b2abce05b257b4b498f14..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/_optimizer_b0230e7_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:21d5da3673206b979eaba9dd6d8918d7745ecd3bd3715e55105fd57c234a3a42 -size 1749776 diff --git a/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/matmul_transpose_triton.py b/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/matmul_transpose_triton.py deleted file mode 100644 index 4565b2c4fd506a4218340d380d6c962b16774b1d..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/matmul_transpose_triton.py +++ /dev/null @@ -1,128 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -def matmul_transpose_assign(d_in, d_out): - assert d_in.is_cuda, "Input `d_in` must be a CUDA tensor" - assert d_out.is_cuda, "Input `d_out` must be a CUDA tensor" - assert d_in.device == d_out.device, "Inputs `d_in` and `d_out` must be on the same CUDA device" - assert d_in.dtype == d_out.dtype, "Inputs must have the same data type" - assert d_in.ndim == 2, "Input `d_in` must be a 2D tensor" - assert d_out.ndim == 2, "Input `d_out` must be a 2D tensor" - assert d_in.size(0) == d_out.size(0) == d_out.size(0), \ - "First dimension of `d_in` must match first and second dimension of `d_out`" - - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -def matmul_transpose(d_in): - M, _ = d_in.shape - d_out = torch.empty((M, M), device=d_in.device, dtype=d_in.dtype) - matmul_transpose_assign(d_in, d_out) - return d_out diff --git a/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/muon.py b/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/muon.py old mode 100644 new mode 100755 index 4af25d55c528fb0db2272540838ab70eb3619194..0d614d55d721efac406c147b4f62e6c703a91107 --- a/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/muon.py +++ b/build/torch27-cxx11-rocm63-x86_64-linux/optimizer/muon.py @@ -1,26 +1,14 @@ -import logging import math -import types from dataclasses import dataclass -from typing import List, Optional, Union, cast import torch import torch.distributed as dist -from torch.distributed._tensor import DTensor, Replicate, Shard - -from .matmul_transpose_triton import matmul_transpose_assign - -logger = logging.getLogger(__name__) - -COMM_DTYPE = torch.bfloat16 +from torch.distributed._tensor import DTensor # This code snippet is a modified version adapted from the following GitHub repositories: # https://github.com/KellerJordan/Muon/blob/master/muon.py -# Muon's Newton–Schulz iteration causes high variance in singular values -# Idea: give each iteration its own 3 coefficients and optimize them via gradient descent. @torch.no_grad() -# matmul_transpose_assign from : https://github.com/nil0x9/flash-muon def _zeropower_via_newtonschulz5(G, steps): """ Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a @@ -32,31 +20,26 @@ def _zeropower_via_newtonschulz5(G, steps): performance at all relative to UV^T, where USV^T = G is the SVD. """ assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE + a, b, c = (3.4445, -4.7750, 2.0315) X = G # no manual typecast - if G.size(0) > G.size(1): X = X.T # Ensure spectral norm is at most 1 X = X / (X.norm() + 1e-7) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) + X = X.bfloat16() # Perform the NS iterations - for a, b, c in [ - (4.0848, -6.8946, 2.9270), - (3.9505, -6.3029, 2.6377), - (3.7418, -5.5913, 2.3037), - (2.8769, -3.1427, 1.2046), - (2.8366, -3.0525, 1.2012), - ]: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) + for _ in range(steps): + A = X @ X.T + # B = ( + # b * A + c * A @ A + # ) + B = torch.addmm(A, A, A, alpha=c, beta=b) + # X = a * X + B @ X + X = torch.addmm(X, B, X, alpha=1.0, beta=a) if G.size(0) > G.size(1): X = X.T - return X + return X.to(G.dtype) @dataclass @@ -64,425 +47,92 @@ class _muon_state: # TODO: use Optional worker_rank: int | None = None gathered_grad: torch.Tensor | None = None - scattered_u: DTensor | None = None computed_u: torch.Tensor | None = None gather_event: torch.cuda.Event | None = None compute_event: torch.cuda.Event | None = None - scatter_event: torch.cuda.Event | None = None - process_group = None - qk_clip_state = None - - -def split_elems_for_src(param, src_rank, num_ranks) -> int: - rows = param.shape[0] - cols = int(param.numel() // rows) - base, rem = divmod(rows, num_ranks) - my_rows = base + (1 if src_rank < rem else 0) - return my_rows * cols @torch.no_grad() -def _alloc_gathered_grad(params, param_to_state, rank, compute_stream): - """ - Pre-allocate gathered_grad buffer on compute_stream - before launching all2all gather - """ - with torch.cuda.stream(compute_stream): - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - num_ranks = dist.get_world_size(group=state.process_group) - state.gathered_grad = torch.empty(p.grad.numel(), - dtype=COMM_DTYPE, - device="cuda") - else: - state.gathered_grad = None - - alloc_event = torch.cuda.Event() - alloc_event.record(compute_stream) - return alloc_event +def _gather(p, state, rank, comm_stream, none_grad): + g = p.grad + mesh = g.device_mesh + if rank == state.worker_rank: + gather_list = [torch.empty_like(g.to_local()) for _ in range(mesh.mesh.numel())] + else: + gather_list = None -@torch.no_grad() -def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad, - alloc_event): - """ - All2all gathers shards so each owner rank reconstructs its full gradient - """ with torch.cuda.stream(comm_stream): - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - - # Construct sending buffers - per_dst = [[] for _ in range(num_ranks)] - send_counts = [0] * num_ranks - - for p in params: - state = param_to_state[id(p)] - dst = state.worker_rank - assert dst < num_ranks - shard_elems = split_elems_for_src(p, rank, num_ranks) - g = p.grad - g = g.to_local().to(COMM_DTYPE).contiguous().view(-1) - assert g.numel() == shard_elems - per_dst[dst].append(g) - send_counts[dst] += shard_elems - - assert any( - len(v) > 0 for v in per_dst - ), "At least one destination rank must receive a sharded tensor" - # list[list[Tensor]] -> list[Tensor] - per_dst = [t for dst in per_dst for t in dst] - - send_buf = torch.cat(per_dst, dim=0) - - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - # Compute receive sizes and allocate receiving buffers - recv_counts = [0] * num_ranks - - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += split_elems_for_src(p, src, num_ranks) - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - #All2All - dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, + torch.distributed.gather( + g.to_local(), + dst=state.worker_rank, + gather_list=gather_list, + group=mesh.get_group(), ) - - # Reconstructs gathered grad from the received buffer - # - # recv_buf (num ranks = 3) - # - # From rank 0 From rank 1 From rank 2 - # | p1_0, p2_0, p3_0 | p1_1, p2_1, p3_1 | p1_2, p2_2, p3_2 | - # - # Outer loop: - # rank 0 -> rank 1 -> rank2 - # - # Inner loop: - # p1_n -> p2_n -> p3_n - - comm_stream.wait_event(alloc_event) - - off = 0 - write_offsets = {id(p): 0 for p in owned_params} - for src in range(num_ranks): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - n = split_elems_for_src(p, src, num_ranks) - assert n > 0 - - sg = recv_buf.narrow(0, off + inner_off, n) - woff = write_offsets[id(p)] - dst = state.gathered_grad.narrow(0, woff, n) - dst.copy_(sg) - - write_offsets[id(p)] += n - inner_off += n - off += block - - for p in params: - state = param_to_state[id(p)] - if state.worker_rank == rank: - state.gathered_grad = state.gathered_grad.view_as(p) - state.gather_event = torch.cuda.Event() - state.gather_event.record(comm_stream) - else: - state.gathered_grad = None - state.gather_event = None - if none_grad: - p.grad = None + if rank == state.worker_rank: + if state.gathered_grad is not None: + raise RuntimeError( + "Gather event already exists, which should not happen." + ) + state.gathered_grad = torch.cat(gather_list, dim=0) + state.gather_event = torch.cuda.Event() + state.gather_event.record() + else: + state.gathered_grad = None + state.gather_event = None + if none_grad: + p.grad = None @torch.no_grad() -def _compute_u(p, state, steps, rank, compute_stream): - """ - On worker_rank, compute the orthogonalized update using Newton-Schulz iteration. - """ +def _compute_u(state, steps, rank, compute_stream): with torch.cuda.stream(compute_stream): if rank == state.worker_rank: if state.gather_event is None: raise RuntimeError("Gather event must be set before compute.") compute_stream.wait_event(state.gather_event) u = _zeropower_via_newtonschulz5(state.gathered_grad, steps) - state.gathered_grad = None state.computed_u = u state.compute_event = torch.cuda.Event() state.compute_event.record() + # Clear the gathered gradient to free memory + state.gathered_grad = None else: state.computed_u = None state.compute_event = None @torch.no_grad() -def _alloc_scattered_u(params, param_to_state, rank, compute_stream): - """ - Pre-allocate scattered_u buffer on compute_stream - before launching all2all gather - """ - with torch.cuda.stream(compute_stream): - for p in params: - state = param_to_state[id(p)] - state.scattered_u = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - alloc_event = torch.cuda.Event() - alloc_event.record(compute_stream) - return alloc_event - +def _scatter(p, state, lr, wd, rank, comm_stream): + u = state.computed_u + mesh = p.device_mesh -def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event): - """ - All2all scatters full gradients to all ranks - """ with torch.cuda.stream(comm_stream): - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - # Construct sending buffer - per_dst = [[] for _ in range(num_ranks)] - send_counts = [0] * num_ranks - - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - if state.compute_event is None: - raise RuntimeError( - "Compute event must be set before scatter.") - comm_stream.wait_event(state.compute_event) - state.gathered_grad = None - - assert state.computed_u is not None - - u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1) - - offset = 0 - for dst in range(num_ranks): - n = split_elems_for_src(p, dst, num_ranks) - assert n > 0 - - su = u_full.narrow(0, offset, n) - per_dst[dst].append(su) - send_counts[dst] += n - offset += n - - assert offset == u_full.numel() - - lengths = [len(v) for v in per_dst] - if all(l > 0 for l in lengths): - assert all( - l == lengths[0] for l in lengths - ), "All destination ranks must have the same number of sharded tensor" - # list[list[Tensor]] -> list[Tensor] - per_dst = [t for dst in per_dst for t in dst] - send_buf = torch.cat(per_dst, dim=0) + if rank == state.worker_rank: + if state.compute_event is None: + raise RuntimeError("Compute event must be set before scatter.") + comm_stream.wait_event(state.compute_event) + scatter_list = list(torch.split(u, p.size(0) // mesh.mesh.numel(), dim=0)) else: - # all_to_all requires participation from all ranks - # Even non-owner ranks must join the collective call - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Compute receive sizes and allocate receiving buffers - recv_counts = [0] * num_ranks - - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += split_elems_for_src(p, rank, num_ranks) - recv_counts[src] = total - - recv_total = sum(recv_counts) - assert recv_total > 0 - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - #All2All - dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, + scatter_list = None + + u = torch.empty_like(p.to_local()) + torch.distributed.scatter( + u, + scatter_list=scatter_list, + src=state.worker_rank, + group=mesh.get_group(), ) - - # Copy to pre-allocated scattered_u buffer from the received buffer - # - # recv_buf (num ranks = 3, local_rank = 0) - # - # From rank 0 From rank 1 From rank 2 - # | p1_0, p2_0, p3_0 | p4_0 | p5_0, p6_0 | - # - # Outer loop: - # rank 0 -> rank 1 -> rank2 - # - # Inner loop: - # src(0) : p1_0 -> p2_0 -> p3_0 - # src(1) : p4_0 - # src(2) : p5_0 -> p6_0 - - comm_stream.wait_event(alloc_event) - - off = 0 - for src in range(num_ranks): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = split_elems_for_src(p, rank, num_ranks) - assert n > 0 - - flat_local = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - state.scattered_u.copy_(flat_local) - - state.scatter_event = torch.cuda.Event() - state.scatter_event.record(comm_stream) - inner_off += n - - assert inner_off == block - off += block - - -def _update_param(p, state, lr, adjusted_lr, weight_decay, rank, - compute_stream): - """ - Update sharded parameter p with the scattered_u. - Only worker_rank frees computed_u. - """ - with torch.cuda.stream(compute_stream): - if state.scatter_event is None: - raise RuntimeError("Scatter event must be set before update") - compute_stream.wait_event(state.scatter_event) - u_dtensor = DTensor.from_local( - state.scattered_u, - placements=p.placements, - device_mesh=p.device_mesh, - ) - - state.scattered_u = u_dtensor - if rank == state.worker_rank: - # Free computed_u + # Clear u to free memory state.computed_u = None - - Muon._update_p(p, state.scattered_u, lr, adjusted_lr, weight_decay) - state.scattered_u = None - u_dtensor = None - - scales_full = Muon._compute_scales(p, state.qk_clip_state) - if scales_full is not None: - num_ranks = dist.get_world_size(group=state.process_group) - local_rank = dist.get_rank(group=state.process_group) - scales_local = scales_full.chunk(num_ranks, dim=0)[local_rank] - scales_local = DTensor.from_local( - scales_local, - placements=p.placements, - device_mesh=p.device_mesh, - ) - Muon._qk_clip(p, scales_local, state.qk_clip_state.head_dim) - - -def default_is_muon(name, x): - skip_keys = ["embed_tokens", "lm_head", "tok_embeddings", "output"] - return x.ndim >= 2 and not any(key in name for key in skip_keys) - - -def get_default_muon_param_groups(model, is_muon_func=default_is_muon): - muon_params, muon_names = [], [] - non_muon_params = [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - ('wq', 'wk', 'q_proj', 'k_proj') and return (kind, layer_index). - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = name.split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: Optional[str] # 'wq'/'q_proj' or 'wk'/'k_proj' or None - indices: List[int] # which heads to consider for clipping - head_dim: int # from config - threshold: float # from config - logit: Optional[torch.Tensor] + u = DTensor.from_local( + u, + placements=p.placements, + device_mesh=mesh, + ) + p.data.mul_(1 - lr * wd) + p.data.add_(u, alpha=-lr) class Muon(torch.optim.Optimizer): @@ -499,87 +149,71 @@ class Muon(torch.optim.Optimizer): - We believe it may not work well for finetuning pretrained models, but we haven't tested this. Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. + muon_params: The parameters to be optimized by Muon. lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) momentum: The momentum used by the internal SGD. (0.95 is a good default) nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. + adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. adamw_lr: The learning rate for the internal AdamW. adamw_betas: The betas for the internal AdamW. adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - overlap_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher overlap_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. + adamw_wd: The weight decay for the internal AdamW. """ - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config={ - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - }, - overlap_step=5): + def __init__( + self, + model, + is_muon_func, + lr=1e-3, + momentum=0.95, + nesterov=True, + ns_steps=5, + adamw_wd=0.1, + adamw_betas=(0.9, 0.95), + adamw_eps=1e-8, + none_grad=True, + debug=False, + ): defaults = dict( lr=lr, - weight_decay=weight_decay, + wd=adamw_wd, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps, adamw_betas=adamw_betas, adamw_eps=adamw_eps, none_grad=none_grad, - use_muon=True, ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) + super().__init__(model.parameters(), defaults) + self.is_muon_func = is_muon_func + self.model = model - super().__init__(params, defaults) + if not dist.is_initialized(): + raise RuntimeError( + "Muon optimizer requires distributed training to be initialized." + ) - self.rank = None + self.rank = dist.get_rank() self.comm_stream = torch.cuda.Stream() self.compute_stream = torch.cuda.Stream() self.debug = debug - self.clip_config = clip_config - self.overlap_step = overlap_step + + def __setstate__(self, state): + # Sort parameters into those for which we will use Muon, and those for which we will not + super().__setstate__(state) + for name, p in self.model.named_parameters(): + if self.is_muon_func(p, name): + # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer + assert p.ndim == 2, p.ndim + self.state[p]["use_muon"] = True + self.state[p]["orig_shape"] = p.shape + else: + # Do not use Muon for parameters in adamw_params + self.state[p]["use_muon"] = False def _calc_flops(self, G, steps): assert len(G.shape) == 2 @@ -597,30 +231,7 @@ class Muon(torch.optim.Optimizer): adjusted_lr = lr * adjusted_ratio return adjusted_lr - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - if p.placements == (Shard(dim=0), ): - # Case for FSDP - return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0) - elif p.placements == (Replicate(), Shard(dim=0)): - # Case for HSDP - process_group = p.device_mesh.get_group(mesh_dim=1) - if self.rank is None: - self.rank = dist.get_rank(group=process_group) - else: - assert self.rank == dist.get_rank(group=process_group) - for i, shard_mesh in enumerate(p.device_mesh.mesh): - if self.rank in shard_mesh: - return shard_mesh, p.device_mesh.get_group(mesh_dim=1) - else: - raise ValueError(f"Unsupported placements ({p.placements}).") - - def init_state_and_assign_params(self, names, params, group, qk_logits): + def init_state_and_assign_params(self, params, group): param_to_state = {} param_to_flops = {} @@ -636,44 +247,34 @@ class Muon(torch.optim.Optimizer): total_flops += flops if self.debug: - print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", - flush=True) - - paired = list(zip(names, params)) + print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", flush=True) - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) + ordered_params = sorted( + params, key=lambda p: param_to_flops[id(p)], reverse=True + ) round_robin = 0 mesh = None - shard_mesh = None - process_group = None - for n, p in zip(ordered_names, ordered_params): + for p in ordered_params: if mesh is None: mesh = p.device_mesh - shard_mesh, process_group = self.get_shard_mesh(p) + if mesh.ndim != 1: + raise NotImplementedError( + "Muon requires a 1D mesh for distributed training yet." + ) elif mesh != p.device_mesh: raise ValueError("All parameters must be on the same mesh.") - num_ranks = dist.get_world_size(group=process_group) + param_to_state[id(p)] = _muon_state() - param_to_state[id( - p)].worker_rank = shard_mesh[round_robin].item() % num_ranks - param_to_state[id(p)].process_group = process_group - qk_clip_state = self.get_qk_clip_info(n, qk_logits) - param_to_state[id(p)].qk_clip_state = qk_clip_state - round_robin = (round_robin + 1) % len(shard_mesh) + param_to_state[id(p)].worker_rank = mesh.mesh[round_robin].item() + + round_robin = (round_robin + 1) % mesh.mesh.numel() return param_to_state, ordered_params - def base(self, names, params, group, lr, weight_decay, momentum, - qk_logits): + def base(self, params, group, lr, wd, momentum): # generate weight updates in distributed fashion - for n, p in zip(names, params): + for p in params: g = p.grad if g is None: continue @@ -692,87 +293,39 @@ class Muon(torch.optim.Optimizer): else: g = buf - u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) + u = _zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) + # scale update adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) - Muon._update_p(p, u, lr, adjusted_lr, weight_decay) - qk_clip_state = self.get_qk_clip_info(n, qk_logits) + # apply weight decay + p.data.mul_(1 - lr * wd) - scales_full = self._compute_scales(p, qk_clip_state) - if scales_full is not None: - Muon._qk_clip(p, scales_full, qk_clip_state.head_dim) + # apply update + p.data.add_(u, alpha=-adjusted_lr) def _update_g(self, p, g, group, momentum): # calc update state = self.state[p] - buf = state.setdefault("momentum_buffer", torch.zeros_like(g)) - torch.add(g, buf, alpha=momentum, out=buf) + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) if group["nesterov"]: - g.add_(buf, alpha=momentum) - return g - return buf + g = g.add(buf, alpha=momentum) + else: + g = buf + return g - @staticmethod - def _update_p(p, u, lr, adjusted_lr, weight_decay): + def _update_p(self, p, u, lr, wd): + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) # apply weight decay - p.data.mul_(1 - lr * weight_decay) + p.data.mul_(1 - lr * wd) # apply update p.data.add_(u, alpha=-adjusted_lr) - def get_qk_clip_info(self, n, qk_logits): - head_dim = self.clip_config.get('head_dim') - threshold = self.clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - indices_key = 'q_indices' if 'q' in kind else 'k_indices' - indices = self.clip_config.get(indices_key, []) or [] - - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - @staticmethod - def _compute_scales(p, qk_clip_state): - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - H_global = p.shape[0] // head_dim - scales_full = torch.ones(H_global, device=p.data.device) - scaling = 0 - - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if new_scale < scales_full[head_idx]: - scales_full[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - scaling += 1 - - return scales_full if scaling > 0 else None - - @staticmethod - def _qk_clip(p, scales, head_dim): - W = p.data.view(-1, head_dim, p.data.shape[1]) - W.mul_(scales.view(-1, 1, 1)) - - def parallel(self, names, params, group, lr, weight_decay, momentum, - qk_logits): + def parallel(self, params, group, lr, wd, momentum): """ Perform a parallel optimization step using Muon. """ @@ -794,143 +347,44 @@ class Muon(torch.optim.Optimizer): p.grad = g param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - assert self.rank is not None + params, group + ) - def enqueue_all2all_gather(start_idx, chunk_size): - target_params = ordered_params[start_idx:start_idx + chunk_size] - if target_params: - alloc_event = _alloc_gathered_grad(target_params, - param_to_state, self.rank, - self.compute_stream) - _all2all_gather(target_params, param_to_state, self.rank, - self.comm_stream, group["none_grad"], - alloc_event) + def enqueue_gathers(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _gather(p, state, self.rank, self.comm_stream, group["none_grad"]) def enqueue_computes(start_idx, chunk_size): - for p in ordered_params[start_idx:start_idx + chunk_size]: + for p in ordered_params[start_idx : start_idx + chunk_size]: state = param_to_state[id(p)] - _compute_u(p, state, group["ns_steps"], self.rank, - self.compute_stream) - - def enqueue_all2all_scatter(start_idx, chunk_size): - target_params = ordered_params[start_idx:start_idx + chunk_size] - if target_params: - alloc_event = _alloc_scattered_u(target_params, param_to_state, - self.rank, - self.compute_stream) - _all2all_scatter(target_params, param_to_state, self.rank, - self.comm_stream, alloc_event) - - def enqueue_update_param(start_idx, chunk_size): - for p in ordered_params[start_idx:start_idx + chunk_size]: + _compute_u(state, group["ns_steps"], self.rank, self.compute_stream) + + def enqueue_scatters(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: state = param_to_state[id(p)] adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) - _update_param(p, state, lr, adjusted_lr, weight_decay, - self.rank, self.compute_stream) + _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream) - chunk_size = dist.get_world_size(param_to_state[id( - params[0])].process_group) + chunk_size = params[0].device_mesh.mesh.numel() # Wait grad update self.comm_stream.wait_stream(torch.cuda.current_stream()) - overlap_step = self.overlap_step - for i in range(0, overlap_step): - enqueue_all2all_gather(i * chunk_size, chunk_size) - enqueue_computes(i * chunk_size, chunk_size) - + enqueue_gathers(0, chunk_size) for i in range(0, len(params) + chunk_size - 1, chunk_size): - enqueue_all2all_scatter(i, chunk_size) - enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size) - enqueue_update_param(i, chunk_size) - enqueue_computes(i + overlap_step * chunk_size, chunk_size) - - # Wait the last update_param to finish - torch.cuda.current_stream().wait_stream(self.compute_stream) - - @staticmethod - def _fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: Union[float, torch.Tensor], - weight_decay: float, - eps: float, - maximize: bool, - ) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: Optional[DeviceDict] = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else - None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [ - params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps - ] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, - non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) + enqueue_computes(i, chunk_size) + enqueue_gathers(i + chunk_size, chunk_size) + enqueue_scatters(i, chunk_size) - def step(self, closure=None, qk_logits=None): + torch.cuda.current_stream().wait_stream(self.comm_stream) + + def step(self, closure=None): """Perform a single optimization step. Args: closure (Callable, optional): A closure that reevaluates the model and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). """ loss = None if closure is not None: @@ -938,127 +392,64 @@ class Muon(torch.optim.Optimizer): loss = closure() for group in self.param_groups: - params = group["params"] - - if group["use_muon"]: - ############################ - # Muon # - ############################ - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - param_dtensors = [] - param_tensors = [] - name_dtensors = [] - name_tensors = [] - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - param_tensors.append(p) - name_tensors.append(n) - else: - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError( - f"Unsupported parameter type: {type(p.data)}") - - if self.debug: - print( - f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors", - flush=True, - ) - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - self.parallel( - name_dtensors, - param_dtensors, - group, - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - qk_logits=qk_logits, - ) - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - qk_logits=qk_logits, - ) - + ############################ + # Muon # + ############################ + + params = [p for p in group["params"] if self.state[p]["use_muon"]] + lr = group["lr"] + wd = group["wd"] + momentum = group["momentum"] + + if isinstance(params[0].data, DTensor): + self.parallel( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) else: - ############################ - # AdamW backup # - ############################ - - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - max_exp_avg_sqs = [] - state_steps = [] - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - for p in params: - g = p.grad - if g is None: - continue - state = self.state[p] - params_with_grads.append(p) - grads.append(g) - if "step" not in state: - state["step"] = (torch.zeros((), - dtype=torch.float32, - device=p.device)) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(state["moment1"]) - moment2.append(state["moment2"]) - if not isinstance(state["step"], torch.Tensor): - step_tensor = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - else: - step_tensor = state["step"] - state_steps.append(step_tensor) - - self._fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - max_exp_avg_sqs, - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, + self.base( + params, + group, lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, + wd=wd, + momentum=momentum, ) + ############################ + # AdamW backup # + ############################ + + params = [p for p in group["params"] if not self.state[p]["use_muon"]] + lr = group["lr"] + beta1, beta2 = group["adamw_betas"] + eps = group["adamw_eps"] + weight_decay = group["wd"] + + for p in params: + g = p.grad + if g is None: + continue + state = self.state[p] + if "step" not in state: + state["step"] = 0 + state["moment1"] = torch.zeros_like(g) + state["moment2"] = torch.zeros_like(g) + state["step"] += 1 + step = state["step"] + buf1 = state["moment1"] + buf2 = state["moment2"] + buf1.lerp_(g, 1 - beta1) + buf2.lerp_(g.square(), 1 - beta2) + + g = buf1 / (eps + buf2.sqrt()) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + scale = bias_correction1 / bias_correction2**0.5 + p.data.mul_(1 - lr * weight_decay) + p.data.add_(g, alpha=-lr / scale) + return loss diff --git a/build/torch28-cxx11-cu126-x86_64-linux/_ops.py b/build/torch28-cxx11-cu126-x86_64-linux/_ops.py deleted file mode 100644 index 034bff088659b5df6f6d401feb18c89dc5f33b29..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _optimizer_8d53b78_dirty -ops = torch.ops._optimizer_8d53b78_dirty - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_optimizer_8d53b78_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-cu126-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so b/build/torch28-cxx11-cu126-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so deleted file mode 100755 index 49cb725bbd4c9011ecc4ad53c60007d4b39e93c4..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a9a7c1beffbad405ef7d6f46f44cf9c6671d119e04a340b54c8f4c8f9d699caf -size 1936664 diff --git a/build/torch28-cxx11-cu126-x86_64-linux/adamw.py b/build/torch28-cxx11-cu126-x86_64-linux/adamw.py deleted file mode 100644 index b5a95816a9f5b9e1889eaadae65373bfbced809a..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/adamw.py +++ /dev/null @@ -1,271 +0,0 @@ -import logging -from collections import defaultdict -from typing import cast - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -def fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: float | torch.Tensor, - weight_decay: float, - eps: float, - maximize: bool, -) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: dict | None = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) - - -def _to_local(t): - """Unwrap DTensor to local tensor for fused ops.""" - return t._local_tensor if isinstance(t, DTensor) else t - - -# --------------------------------------------------------------------------- -# Caches for eliminating per-step Python overhead. -# -# Placement grouping and tensor list assembly are identical every step -# (params don't change placement, moment/step tensors are the same objects -# after initialisation). We cache them keyed by id() of the param list -# stored in param_groups (stable across steps). -# -# Only gradients change each step and must be collected fresh. -# --------------------------------------------------------------------------- - -# id(group["params"]) → dict[placement_key, list[param]] -_placement_cache: dict[int, dict[tuple, list]] = {} - -# id(placement_group_list) → (params_local, moment1, moment2, state_steps) -_tensor_cache: dict[int, tuple[list, list, list, list]] = {} - - -def _step_adamw_params_slow(optimizer_state, params, group): - """Uncached fallback for the rare case where some params lack grads.""" - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - g = p.grad - if g is None: - continue - state = optimizer_state[p] - params_with_grads.append(_to_local(p)) - grads.append(_to_local(g)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - if not params_with_grads: - return - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw_params(optimizer_state, params, group): - """Run fused AdamW on a list of parameters sharing the same placement. - - After the first call, cached tensor lists (params_local, moment1, - moment2, state_steps) are reused — only gradients are collected fresh. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - params: List of parameters to update. - group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay. - """ - # Collect grads — the only thing that changes each step. - with record_function("adamw::collect_grads"): - grads = [] - for p in params: - g = p.grad - if g is None: - # Rare: fall back to slow path that filters per-param. - _step_adamw_params_slow(optimizer_state, params, group) - return - grads.append(_to_local(g)) - - tensor_key = id(params) - if tensor_key not in _tensor_cache: - with record_function("adamw::init_tensor_cache"): - params_local = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - state = optimizer_state[p] - params_local.append(_to_local(p)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(p.grad) - state["moment2"] = torch.zeros_like(p.grad) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - _tensor_cache[tensor_key] = (params_local, moment1, moment2, - state_steps) - - params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key] - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - with record_function("adamw::fused_adamw"): - fused_adamw( - params_local, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw(optimizer_state, group): - """Dispatch AdamW step, grouping parameters by type and placement. - - Placement grouping is cached after the first call since params never - change their placement between steps. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - group: Parameter group dict. - """ - params = group["params"] - placement_key = id(params) - - if placement_key not in _placement_cache: - with record_function("adamw::group_by_placement"): - placement_to_params: dict[tuple, - list[torch.Tensor]] = defaultdict(list) - for p in params: - match p: - case DTensor(): - logger.debug( - "[AdamW] DTensor param: shape=%s, placements=%s, " - "mesh=%s, grad=%s", p.shape, p.placements, - p.device_mesh.mesh_dim_names, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple( - [p.placements, p.device_mesh])].append(p) - case torch.Tensor(): - logger.debug( - "[AdamW] plain param: shape=%s, grad=%s", p.shape, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple([torch.Tensor, - None])].append(p) - - logger.debug("[AdamW] %d placement groups, %d total params", - len(placement_to_params), len(params)) - - _placement_cache[placement_key] = dict(placement_to_params) - - for group_params in _placement_cache[placement_key].values(): - step_adamw_params(optimizer_state, group_params, group) diff --git a/build/torch28-cxx11-cu126-x86_64-linux/async_utils.py b/build/torch28-cxx11-cu126-x86_64-linux/async_utils.py deleted file mode 100644 index a45c530ac9cad88e3555ec1047a6aa59f225347e..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/async_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import logging -from typing import Generator - -logger = logging.getLogger(__name__) - - -class _Task: - """Internal: wraps a generator, advances one yield at a time.""" - - def __init__(self, generator: Generator[None, None, None], index: int): - self._generator = generator - self._index = index - self._steps_completed = 0 - self.step() # run to first yield - - def step(self) -> bool: - try: - next(self._generator) - self._steps_completed += 1 - logger.debug("pipeline[%d] completed stage %d", self._index, - self._steps_completed) - return True - except StopIteration: - logger.debug("pipeline[%d] finished after %d stages", self._index, - self._steps_completed) - return False - - def close(self): - self._generator.close() - - -def run_pipeline( - pipelines: Generator[Generator[None, None, None], None, None], - max_concurrent: int, -) -> None: - """Run generator-based pipelines with bounded concurrency. - - Each pipeline is a generator that yields at stage boundaries. - The runtime interleaves pipelines so communication and computation - overlap across chunks. - """ - if max_concurrent <= 0: - raise ValueError(f"max_concurrent must be > 0, got {max_concurrent}") - - have_new = True - task_index = 0 - previous_tasks: list[_Task] = [] - - try: - while have_new or previous_tasks: - running_tasks: list[_Task] = [] - - # Admit one new pipeline per iteration (staggered admission). - # Admitting one at a time ensures that while chunk N does NS - # compute on the default stream, chunk N+1's NCCL all-to-all - # runs concurrently on the NCCL stream — creating real - # communication/computation overlap on the GPU. - if have_new and len(previous_tasks) < max_concurrent: - try: - gen = next(pipelines) - task = _Task(gen, task_index) - task_index += 1 - running_tasks.append(task) - except StopIteration: - have_new = False - - # Advance every previously-yielded task by one step. - for task in previous_tasks: - if task.step(): - running_tasks.append(task) - - previous_tasks = running_tasks - except BaseException: - # Clean up all in-flight generators to release GPU resources. - for task in previous_tasks: - task.close() - raise diff --git a/build/torch28-cxx11-cu126-x86_64-linux/core.py b/build/torch28-cxx11-cu126-x86_64-linux/core.py deleted file mode 100644 index c69d515afef305ad0ed66374095fa2d2468d99cc..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/core.py +++ /dev/null @@ -1,219 +0,0 @@ -import logging -import math -from dataclasses import dataclass -from typing import List - -import torch -from torch.distributed import ProcessGroup -from torch.distributed.tensor import DTensor - -# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into -# parameter FQNs. Activation checkpointing similarly inserts -# "_checkpoint_wrapped_module". Strip these so name-based matching (skip_keys, -# expert_keys, QK layer parsing) works regardless of wrapper nesting. -_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) - -logger = logging.getLogger(__name__) - - -def normalize_fqn(name: str) -> str: - """Strip torch.compile / checkpoint wrapper components from a parameter FQN.""" - return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS) - - -@dataclass -class _muon_state: - worker_rank: int - process_group: ProcessGroup - rank_indices: dict[int, tuple] # local_rank -> per-dim indices - rank_numels: dict[int, int] # local_rank -> numel - name: str - qk_clip_state: torch.Tensor | None = None - - -def _batch_momentum( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update (no nesterov).""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - - -def _batch_momentum_nesterov( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update with nesterov correction.""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - nesterov_terms = torch._foreach_mul(momentum_bufs, momentum) - torch._foreach_add_(grads, nesterov_terms) - - -_compiled_momentum: dict[bool, callable] = {} -_use_momentum_compile = True - - -def set_momentum_compile(enabled: bool): - """Toggle torch.compile for batched momentum.""" - global _use_momentum_compile - _use_momentum_compile = enabled - - -def batch_pre_ortho( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, - nesterov: bool, -) -> None: - """Batched momentum update on lists of plain tensors. - - Mirrors dion's ``muon_update_pre_orthogonalize``. - Inputs must be plain CUDA tensors (not DTensor). - Modifies ``momentum_bufs`` and (for nesterov) ``grads`` in-place. - - When compile is enabled, uses separately compiled functions for - nesterov=True/False to avoid graph breaks from the branch. - """ - fn = _batch_momentum_nesterov if nesterov else _batch_momentum - if _use_momentum_compile: - if nesterov not in _compiled_momentum: - _compiled_momentum[nesterov] = torch.compile(fn) - fn = _compiled_momentum[nesterov] - fn(grads, momentum_bufs, momentum) - - -def _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay): - """Weight-decay + update on plain tensors. - - Not compiled: per-param @torch.compile caused ~0.25ms TorchDynamo cache - lookup per call × 256+ params = massive overhead. The pipeline path uses - batched _foreach_* ops instead; this function remains for base() and - distributed_muon(). - """ - p_data.mul_(1 - lr * weight_decay) - p_data.add_(u_data, alpha=-adjusted_lr) - - -def update_p(p, u, lr, adjusted_lr, weight_decay): - """Apply weight decay and orthogonalized update to parameter. - - Args: - p: Parameter (torch.nn.Parameter or DTensor). - u: Orthogonalized update tensor. - lr: Base learning rate. - adjusted_lr: Size-adjusted learning rate. - weight_decay: Weight decay coefficient. - """ - # Unwrap Parameter -> underlying data tensor. - p_data = p.data if isinstance(p, torch.nn.Parameter) else p - # Unwrap DTensor -> local CUDA tensor for compiled kernel. - if isinstance(p_data, DTensor): - p_data = p_data._local_tensor - u_data = u._local_tensor if isinstance(u, DTensor) else u - _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay) - - -def adjust_lr_for_muon(lr, param_shape): - """Scale learning rate based on parameter matrix dimensions. - - Args: - lr: Base learning rate. - param_shape: Shape of the parameter tensor. - - Returns: - Adjusted learning rate. - """ - A, B = param_shape[:2] - # We adjust the learning rate and weight decay based on the size of the parameter matrix - # as described in the paper - adjusted_ratio = 0.2 * math.sqrt(max(A, B)) - adjusted_lr = lr * adjusted_ratio - return adjusted_lr - - -def _match_key(parts, key): - """Check if key matches as contiguous components in parts. - - Single-component keys (e.g. "experts") match any single component. - Multi-component keys (e.g. "experts.w1") match as a contiguous subsequence. - """ - key_parts = key.split(".") - key_len = len(key_parts) - if key_len == 1: - return key in parts - return any(parts[i:i + key_len] == key_parts - for i in range(len(parts) - key_len + 1)) - - -def is_expert_param(name, expert_keys): - """Check if a parameter name matches any expert key (component-level).""" - if not expert_keys: - return False - parts = normalize_fqn(name).split(".") - return any(_match_key(parts, key) for key in expert_keys) - - -def default_is_muon(name, x, expert_keys=None): - normalized = normalize_fqn(name) - parts = normalized.split(".") - skip_keys = [ - "embed_tokens", - "lm_head", - "tok_embeddings", - "output", - "mhc_attn", - "mhc_ffn", - "lambda_proj", - ] - if any(key in parts for key in skip_keys): - logger.info( - "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d", - normalized, name, x.ndim) - return False - effective_ndim = x.ndim - is_expert = is_expert_param(name, expert_keys) - if is_expert: - effective_ndim -= 1 - result = effective_ndim >= 2 - logger.info( - "[is_muon] %s (orig: %s): ndim=%d, expert=%s, effective_ndim=%d → %s", - normalized, name, x.ndim, is_expert, effective_ndim, - "Muon" if result else "AdamW") - return result - - -def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None): - if is_muon_func is None: - is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys) - - muon_params, muon_names = [], [] - non_muon_params, non_muon_names = [], [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - non_muon_names.append(n) - - logger.info("[param_groups] expert_keys=%s, Muon=%d, AdamW=%d", - expert_keys, len(muon_names), len(non_muon_names)) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] diff --git a/build/torch28-cxx11-cu126-x86_64-linux/cpu_offload.py b/build/torch28-cxx11-cu126-x86_64-linux/cpu_offload.py deleted file mode 100644 index fb5e69154a1d4a6c884491413a37a9acf0f66c80..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/cpu_offload.py +++ /dev/null @@ -1,206 +0,0 @@ -"""CPU offloading for optimizer states. - -Manages a pinned CPU memory pool and async CUDA streams to offload -optimizer state tensors (momentum buffers, Adam moments) to CPU between -optimizer steps, freeing GPU memory. - -All tracked tensors are packed into a single flat pinned CPU buffer -(per dtype). D2H and H2D copies are performed per-tensor directly -between individual GPU tensors and their slice of the CPU flat buffer -— no GPU staging buffer is allocated, so there is **no temporary GPU -memory spike** during offload or reload. - -Individual tensor storages are freed after offload via -``untyped_storage().resize_(0)``, preserving tensor identity so -downstream caches remain valid. -""" - -import logging -from collections import defaultdict - -import torch -from torch.distributed.tensor import DTensor - -logger = logging.getLogger(__name__) - - -class CPUOffloadPool: - """Pinned CPU memory pool for async optimizer state offloading. - - Tracked tensors are grouped by dtype. Each group gets a single flat - pinned CPU buffer. D2H / H2D copies are per-tensor (into slices of - the flat buffer) to avoid allocating a GPU staging buffer. - """ - - def __init__(self): - self._managed: list[torch.Tensor] = [] - self._storage_nbytes: dict[int, int] = {} # id(t) → bytes - - # Per-dtype group: populated on first offload. - # dtype → dict with keys: - # "indices" : list[int] managed-list indices - # "offsets" : list[tuple[int,int]] (start, numel) in flat buf - # "total" : int total numel - # "cpu_flat" : Tensor pinned CPU buffer - self._groups: dict[torch.dtype, dict] = {} - - self._offload_stream: torch.cuda.Stream | None = None - self._device: torch.device | None = None - self._initialized: bool = False - self._logged: bool = False - - # ------------------------------------------------------------------ - @staticmethod - def _local(t: torch.Tensor) -> torch.Tensor: - """Unwrap DTensor to its local CUDA tensor.""" - return t._local_tensor if isinstance(t, DTensor) else t - - def _ensure_stream(self): - if self._offload_stream is None: - self._offload_stream = torch.cuda.Stream(device=self._device) - - # ------------------------------------------------------------------ - def track(self, tensor: torch.Tensor): - """Register a GPU tensor for CPU offloading. Idempotent.""" - tid = id(tensor) - if tid in self._storage_nbytes: - return - local = self._local(tensor) - if self._device is None: - self._device = local.device - storage = local.untyped_storage() - # Skip tensors with empty storage (e.g. empty FSDP shards) - if storage.size() == 0: - return - self._storage_nbytes[tid] = storage.size() - self._managed.append(tensor) - - # ------------------------------------------------------------------ - def _init_buffers(self): - """Build per-dtype flat buffers on first offload.""" - # Group managed tensors by dtype. - dtype_map: dict[torch.dtype, list[tuple[int, int]]] = defaultdict(list) - for idx, t in enumerate(self._managed): - local = self._local(t) - dtype_map[local.dtype].append((idx, local.numel())) - - total_cpu_bytes = 0 - for dtype, entries in dtype_map.items(): - offsets: list[tuple[int, int]] = [] - indices: list[int] = [] - off = 0 - for idx, n in entries: - indices.append(idx) - offsets.append((off, n)) - off += n - cpu_flat = torch.empty(off, dtype=dtype, device="cpu", pin_memory=True) - self._groups[dtype] = { - "indices": indices, - "offsets": offsets, - "total": off, - "cpu_flat": cpu_flat, - } - total_cpu_bytes += off * cpu_flat.element_size() - - self._initialized = True - logger.info( - "[CPUOffload] Pool initialized: %d tensors, %d dtype group(s), " - "%.2f MB pinned CPU memory", - len(self._managed), - len(self._groups), - total_cpu_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def offload(self): - """Per-tensor async D2H into CPU flat buffer, then free GPU storage.""" - if not self._managed: - return - if not self._initialized: - self._init_buffers() - self._ensure_stream() - - # Offload stream waits for compute to finish. - compute_event = torch.cuda.current_stream(self._device).record_event() - self._offload_stream.wait_event(compute_event) - - offloaded_bytes = 0 - - # Per-tensor D2H copies directly into CPU flat buffer slices. - # No GPU staging buffer → no temporary GPU memory spike. - with torch.cuda.stream(self._offload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - cpu_flat[off : off + n].copy_(local.reshape(-1), non_blocking=True) - - offloaded_bytes += grp["total"] * cpu_flat.element_size() - - # Wait for all D2H copies to land, then free GPU storage. - self._offload_stream.synchronize() - for t in self._managed: - storage = self._local(t).untyped_storage() - if storage.size() != 0: - storage.resize_(0) - else: - raise RuntimeError( - f"Tensor storage is already freed (size=0) before offload. " - f"This indicates a double-free or external interference. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - - if not self._logged: - logger.info( - "[CPUOffload] Offloaded %.2f MB (GPU → CPU)", - offloaded_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def reload(self): - """Per-tensor H2D from CPU flat buffer on the default stream. - - Runs on the current (default) CUDA stream to avoid stream - interaction issues with the parallel Muon pipeline. Since - pinned CPU memory is the source, the copies overlap with - GPU idle time between steps. - """ - if not self._managed or not self._initialized: - return - - reloaded_bytes = 0 - - # Re-allocate all GPU storages first. - for t in self._managed: - local = self._local(t) - storage = local.untyped_storage() - if storage.size() != 0: - raise RuntimeError( - f"Storage should have been freed (size=0) before reload, " - f"but got size={storage.size()}. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - storage.resize_(self._storage_nbytes[id(t)]) - - # Per-tensor H2D copies from CPU flat buffer slices. - # non_blocking=True with pinned source allows DMA overlap. - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off : off + n], non_blocking=True) - - reloaded_bytes += grp["total"] * cpu_flat.element_size() - - if not self._logged: - logger.info( - "[CPUOffload] Reloaded %.2f MB (CPU → GPU)", reloaded_bytes / (1024**2) - ) diff --git a/build/torch28-cxx11-cu126-x86_64-linux/distributed/utils.py b/build/torch28-cxx11-cu126-x86_64-linux/distributed/utils.py deleted file mode 100644 index 683096bcd50489e7a09959753fdbc29ba0f3b6a6..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/distributed/utils.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import (Placement, Shard, - _StridedShard) - - -def _is_shard(placement: Placement) -> bool: - """Check if a placement is a shard type (Shard or _StridedShard). - - In PyTorch 2.10+, _StridedShard no longer inherits from Shard, so - ``placement.is_shard()`` returns False for _StridedShard. This helper - handles both old and new hierarchies. - """ - return isinstance(placement, (Shard, _StridedShard)) - - -def get_slices_of_dtensor( - target: DTensor | torch.Tensor, - local_rank: int, - shard_mesh: DeviceMesh, - shard_placements: tuple[Placement], -) -> tuple[slice | torch.Tensor, ...]: - """ - Get per-dimension indices for a given rank's shard of the target tensor. - - Uses ``Shard.local_shard_size_and_offset`` and - ``_StridedShard.local_shard_size_and_offset`` for correct handling of - both contiguous and strided (non-contiguous) sharding. - - Args: - target (DTensor | torch.Tensor): The target tensor (for its shape). - local_rank (int): The local rank within the shard group. - shard_mesh (DeviceMesh): The shard mesh (only shard dimensions). - shard_placements (tuple[Placement]): The shard placements. - - Returns: - A tuple of indices (one per tensor dim). Each element is either: - - A ``slice`` (for contiguous or unsharded dims) - - A 1-D ``torch.LongTensor`` of indices (for strided sharding) - """ - - # find the global rank of the local rank in the shard mesh - rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank] - - rank_coords = (shard_mesh.mesh == rank).nonzero() - - assert len(rank_coords) == 1 - rank_coords = tuple(rank_coords[0].tolist()) - - assert len(rank_coords) == len(shard_placements) - - # Track per-shard-dim indices. - # None means "not yet sharded on this dim". - dim_indices: dict[int, torch.Tensor] = {} - - # Caution: Assuming replicate-to-shard of the shard mesh goes with - # left-to-right sharding. This is ensured by the sorting logic of - # construct_shard_mesh function. - for mesh_dim_idx, (rank_coord, placement) in enumerate( - zip(rank_coords, shard_placements)): - assert _is_shard(placement) - - num_chunks = shard_mesh.mesh.shape[mesh_dim_idx] - shard_dim = placement.dim - - # Current effective size on this dim (may already be sub-sharded) - if shard_dim in dim_indices: - curr_size = len(dim_indices[shard_dim]) - else: - curr_size = target.size()[shard_dim] - - # Compute indices for this level of sharding - if isinstance(placement, _StridedShard): - _shard_size, offsets = _StridedShard.local_shard_size_and_offset( - placement, - curr_size, - num_chunks, - rank_coord, - return_first_offset=False) - new_indices = torch.tensor(offsets, dtype=torch.long) - else: - shard_size, offset = Shard.local_shard_size_and_offset( - curr_size, num_chunks, rank_coord) - new_indices = torch.arange(offset, - offset + shard_size, - dtype=torch.long) - - # Compose with previous indices on this dim - if shard_dim in dim_indices: - dim_indices[shard_dim] = dim_indices[shard_dim][new_indices] - else: - dim_indices[shard_dim] = new_indices - - # Build result tuple - result: list[slice | torch.Tensor] = [] - for d in range(len(target.size())): - if d not in dim_indices: - result.append(slice(None)) - else: - indices = dim_indices[d] - # Convert contiguous indices to slice for efficiency - if len(indices) > 0: - start = indices[0].item() - expected = torch.arange(start, - start + len(indices), - dtype=torch.long) - if torch.equal(indices, expected): - result.append(slice(start, start + len(indices))) - else: - result.append(indices) - else: - result.append(slice(0, 0)) - - return tuple(result) - - -_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, - ProcessGroup]] = dict() - - -def construct_shard_mesh( - placements: tuple[Placement], - mesh: DeviceMesh, -) -> tuple[DeviceMesh, ProcessGroup, tuple[Placement, ...]]: - """Construct shard sub-mesh and ProcessGroup for all-to-all communication. - - Given a DTensor's placements and device mesh, extracts the "shard group" - — the set of ranks that together hold all shards of the same replica — - and creates a ProcessGroup for all-to-all among them. - - Steps: - 1. Sort placements: Replicate first, then Shard by (dim, granularity). - 2. Permute the mesh tensor to match the sorted order. - 3. Collapse Replicate dims → list of shard sub-meshes (one per replica). - 4. Create/retrieve a cached ProcessGroup for the current rank's sub-mesh. - - Example — 8 GPUs, mesh shape (2, 2, 2), - placements ``[Shard(0), Replicate, _StridedShard(0)]``:: - - Step 1 — Sort: [Replicate, _StridedShard(0), Shard(0)] - Permutation: [1, 2, 0] - - Step 2 — Permute mesh dims by [1, 2, 0]: - Original: Permuted: - [[[0,1],[2,3]], [[[0,2],[1,3]], - [[4,5],[6,7]]] [[4,6],[5,7]]] - - Step 3 — Unbind replicate dim (dim 0), giving 2 shard sub-meshes: - sub-mesh 0 = [[0,2],[1,3]] (replica group 0) - sub-mesh 1 = [[4,6],[5,7]] (replica group 1) - shard_placements = (_StridedShard(0), Shard(0)) - - Step 4 — Rank 0 → ProcessGroup([0,1,4,5]) - Rank 2 → ProcessGroup([2,3,6,7]) - - Returns: - ``(shard_mesh, process_group, shard_placements)`` - """ - my_rank = dist.get_rank() - assert mesh.mesh.device.type == 'cpu' - - # -- Fast path: 1D all-shard mesh → reuse existing PG. ---------------- - # Reuses the mesh's existing ProcessGroup directly, avoiding the - # overhead of dist.new_group(). The standard path below also handles - # subset calls safely via use_local_synchronization=True, but this - # fast path is still beneficial for the common 1D shard case. - if mesh.ndim == 1 and len(placements) == 1 and _is_shard(placements[0]): - key = (*mesh.mesh.shape, *mesh.mesh.flatten().tolist()) - if key not in _ranks_to_dist_cache: - _ranks_to_dist_cache[key] = (mesh, mesh.get_group()) - return (*_ranks_to_dist_cache[key], tuple(placements)) - - mesh_tensor = mesh.mesh.clone() - - # -- Step 1: Sort placements (Replicate first, then Shard by dim). ------ - # _StridedShard comes BEFORE regular Shard on the same dim so that - # get_slices_of_dtensor applies the outer sharding first, matching - # DTensor's left-to-right (outer-to-inner) composition order. - def _sort_key(item): - index, placement = item - assert not placement.is_partial(), "Partial placement not supported" - if placement.is_replicate(): - return (-1, 0, index) - assert _is_shard(placement), f"Unsupported: {type(placement)}" - split = (-1 / placement.split_factor if isinstance( - placement, _StridedShard) else 0) - return (placement.dim, split, index) - - indexed = sorted(enumerate(placements), key=_sort_key) - perm, sorted_placements = zip(*indexed) - - # -- Step 2: Permute mesh to match sorted placement order. -------------- - sorted_mesh = mesh_tensor.permute(perm) - - # -- Step 3: Collapse replicate dims → list of shard sub-meshes. -------- - # E.g. mesh (2, 3, 4, 4) with [R, R, S(0), S(1)] → 6 sub-meshes of (4, 4) - num_rep = sum(1 for p in sorted_placements if p.is_replicate()) - if num_rep > 0: - if num_rep > 1: - sorted_mesh = sorted_mesh.flatten(0, num_rep - 1) - shard_meshes = list(torch.unbind(sorted_mesh, dim=0)) - else: - shard_meshes = [sorted_mesh] - shard_placements = sorted_placements[num_rep:] - assert len(shard_placements) == len(set(shard_placements)) - - # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. -- - # Each rank only creates the group it belongs to, using - # use_local_synchronization=True so that only group members need to - # coordinate. This avoids deadlocks when different PP stages call - # construct_shard_mesh for different parameters. - def _cache_key(t: torch.Tensor) -> tuple: - return (*t.shape, *t.flatten().tolist()) - - my_key = None - for sm in shard_meshes: - if (my_rank == sm).any().item(): - key = _cache_key(sm) - assert my_key is None, "Rank appears in multiple shard groups" - my_key = key - if key not in _ranks_to_dist_cache: - pg = dist.new_group(sm.flatten().tolist(), - use_local_synchronization=True) - _ranks_to_dist_cache[key] = ( - DeviceMesh(device_type="cuda", mesh=sm), - pg, - ) - - return (*_ranks_to_dist_cache[my_key], shard_placements) diff --git a/build/torch28-cxx11-cu126-x86_64-linux/matmul_transpose_triton.py b/build/torch28-cxx11-cu126-x86_64-linux/matmul_transpose_triton.py deleted file mode 100644 index 792de23d82c3fb45fe33d397ab9b76a0787259d0..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/matmul_transpose_triton.py +++ /dev/null @@ -1,122 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], - restore_value=['y'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -@torch.library.custom_op("muon::matmul_transpose_assign", - mutates_args=("d_out", )) -def matmul_transpose_assign(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """Compute d_out = d_in @ d_in.T using an optimized Triton kernel.""" - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -@matmul_transpose_assign.register_fake -def _(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """FakeTensor impl: d_out is already allocated, mutation is declared.""" - pass diff --git a/build/torch28-cxx11-cu126-x86_64-linux/metadata.json b/build/torch28-cxx11-cu126-x86_64-linux/metadata.json deleted file mode 100644 index c55a35717622f1dd5c8ba376ea3a814cbcc10d78..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/metadata.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch28-cxx11-cu126-x86_64-linux/muon.py b/build/torch28-cxx11-cu126-x86_64-linux/muon.py deleted file mode 100644 index 14c0e22471fa6d47a51ed95e0e0c341dc18d5194..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/muon.py +++ /dev/null @@ -1,1068 +0,0 @@ -import logging -import types -from collections import defaultdict -from typing import Any - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Replicate, Shard -from torch.profiler import record_function - -from .adamw import _placement_cache, _tensor_cache, step_adamw -from .async_utils import run_pipeline -from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho, - get_default_muon_param_groups, is_expert_param, update_p) -from .cpu_offload import CPUOffloadPool -from .distributed.utils import (_is_shard, construct_shard_mesh, - get_slices_of_dtensor) -from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO, - _zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5_batched) -from .pipeline import muon_chunk_pipeline, prelaunch_first_gather -from .qk_clip import compute_scales, get_qk_clip_info, qk_clip - -logger = logging.getLogger(__name__) - - -def _expand_expert_params(names, params, expert_keys): - """Expand expert params by splitting on dim 0 (expert dimension). - - Params whose name matches any key in ``expert_keys`` are treated as - expert-parallel tensors. Their outermost dimension is the expert - dimension: an ``(E, out, in)`` tensor becomes ``E`` separate 2D - ``nn.Parameter`` views so that in-place updates propagate back to - the original storage. - - Non-expert params with ``ndim > 2`` trigger an ``AssertionError`` — - if they are expert params, their key must be added to ``expert_keys``. - - The grad must already be set on each expert param (e.g. after momentum). - - For DTensor expert params, placements that shard on dim 0 (expert dim) - are consumed by the split. Non-dim-0 shard placements (e.g. TP) are - preserved: each 2D slice is wrapped as a DTensor on the corresponding - submesh so the parallel pipeline handles the TP communication. - """ - expanded_names = [] - expanded_params = [] - - for n, p in zip(names, params): - is_expert = is_expert_param(n, expert_keys) - is_dtensor = isinstance(p.data, DTensor) - - if is_expert: - if is_dtensor: - logger.debug( - "[expand_expert] %s: expert DTensor, shape=%s, " - "placements=%s, mesh=%s, local_shape=%s", n, p.shape, - p.placements, p.device_mesh.mesh_dim_names, - p.to_local().shape) - else: - logger.debug( - "[expand_expert] %s: expert plain tensor, shape=%s", n, - p.data.shape) - - if not is_expert: - assert p.data.ndim <= 2, ( - f"Param {n} has ndim={p.data.ndim} but does not match " - f"expert_keys={expert_keys}. If this is an expert param, " - f"add its key to expert_keys.") - expanded_names.append(n) - expanded_params.append(p) - continue - - g = p.grad - assert g is not None, ( - f"Expert param {n} must have grad set before expansion") - - tp_mesh = None - tp_placements_2d = None - - if is_dtensor: - local_data = p.to_local() - local_grad = g.to_local() if isinstance(g, DTensor) else g - - # Find non-dim-0 shard placements (e.g. TP sharding). - # After splitting on dim 0, Shard(k) becomes Shard(k-1). - tp_dim_indices = [] - tp_placements_2d = [] - for i, pl in enumerate(p.placements): - if _is_shard(pl) and pl.dim != 0: - tp_dim_indices.append(i) - tp_placements_2d.append(Shard(pl.dim - 1)) - - if tp_dim_indices: - tp_dim_names = tuple(p.device_mesh.mesh_dim_names[i] - for i in tp_dim_indices) - if len(tp_dim_names) == 1: - tp_mesh = p.device_mesh[tp_dim_names[0]] - else: - tp_mesh = p.device_mesh[tp_dim_names] - else: - local_data = p.data - local_grad = g - - # Expand: split dim 0, reshape each slice to 2D. - num_local_experts = local_data.shape[0] - for i in range(num_local_experts): - slice_data = local_data[i] - slice_grad = local_grad[i] - - if tp_mesh is not None: - # Wrap as DTensor on TP submesh so the pipeline handles - # TP communication (gather/scatter across TP ranks). - dt_data = DTensor.from_local(slice_data, - device_mesh=tp_mesh, - placements=tp_placements_2d) - dt_grad = DTensor.from_local(slice_grad, - device_mesh=tp_mesh, - placements=tp_placements_2d) - expert_param = torch.nn.Parameter(dt_data, requires_grad=False) - expert_param.grad = dt_grad - else: - expert_param = torch.nn.Parameter(slice_data, - requires_grad=False) - expert_param.grad = slice_grad - - expanded_names.append(f"{n}[{i}]") - expanded_params.append(expert_param) - - p.grad = None # allow expert grad storage to be freed after pipeline - - return expanded_names, expanded_params - - -class Muon(torch.optim.Optimizer): - """ - Muon - MomentUm Orthogonalized by Newton-schulz - - Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- - processing step, in which each 2D parameter's update is replaced with the nearest orthogonal - matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has - the advantage that it can be stably run in bfloat16 on the GPU. - - Some warnings: - - We believe this optimizer is unlikely to work well for training with small batch size. - - We believe it may not work well for finetuning pretrained models, but we haven't tested this. - - Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. - lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) - momentum: The momentum used by the internal SGD. (0.95 is a good default) - nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) - ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. - Parameters that are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW instead. - adamw_lr: The learning rate for the internal AdamW. - adamw_betas: The betas for the internal AdamW. - adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - warmup_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher warmup_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. - chunk_size : Batch size of parameters to process in each - all2all gather/compute/scatter step. - Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified. - use_distributed_muon: Use distributed muon by Liu et al. (2024). - For testing purpose only. - expert_keys: List of strings to identify expert-parallel parameters. - If any key appears in a parameter's name, its outermost - dimension is treated as the expert dimension and expanded - into per-expert 2D params for Muon. For example, - ``expert_keys=["experts"]`` matches any param whose name - contains "experts". 3D+ params not matched by any key - will raise an error. - """ - - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config=None, - warmup_step=5, - chunk_size=-1, - use_distributed_muon=False, - expert_keys=None): - defaults = dict( - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - nesterov=nesterov, - ns_steps=ns_steps, - adamw_betas=adamw_betas, - adamw_eps=adamw_eps, - none_grad=none_grad, - use_muon=True, - ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) - super().__init__(params, defaults) - - self.debug = debug - self.clip_config = clip_config if clip_config is not None else { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100, - } - self.warmup_step = warmup_step - self.chunk_size = chunk_size - self.use_distributed_muon = use_distributed_muon - self.expert_keys = expert_keys - self.cpu_offload = False - self._cpu_offload_pool: CPUOffloadPool | None = None - self._offload_initialized = False - self._parallel_cache: dict[tuple[str, ...], dict] = {} - self._expert_expand_cache: dict[tuple[int, ...], dict] = {} - - def _calc_flops(self, G, steps): - assert len(G.shape) == 2 - M, N = G.shape - if M > N: - M, N = N, M - - return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) - - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - p.placements, p.device_mesh) - - return shard_mesh, shard_pg, shard_placements - - def init_state_and_assign_params(self, names, params, group, qk_logits): - param_to_state = {} - param_to_flops = {} - - total_flops = 0 - for p in params: - g = p.grad - if g is None: - continue - assert g.ndim == 2, "Muon only supports 2D parameters." - - flops = self._calc_flops(g, group["ns_steps"]) - param_to_flops[id(p)] = flops - total_flops += flops - - if self.debug: - logger.debug("Total TFLOPs for Muon: %.2f TFLOPs", - total_flops / 1e12) - - paired = list(zip(names, params)) - - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) - - round_robin = 0 - mesh = ordered_params[0].device_mesh - placements = ordered_params[0].placements - - shard_mesh, shard_pg, shard_placements = self.get_shard_mesh( - ordered_params[0]) - shard_mesh_flattened = shard_mesh.mesh.flatten() - num_ranks = dist.get_world_size(group=shard_pg) - - for n, p in zip(ordered_names, ordered_params): - if mesh != p.device_mesh: - raise ValueError("All parameters must be on the same mesh.") - if placements != p.placements: - raise ValueError("All parameters must have same placements.") - - worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks - round_robin = (round_robin + 1) % len(shard_mesh_flattened) - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - # Precompute per-rank indices and numels for all-to-all. - rank_indices: dict[int, tuple] = {} - rank_numels: dict[int, int] = {} - for r in range(num_ranks): - indices = get_slices_of_dtensor(p, r, shard_mesh, - shard_placements) - rank_indices[r] = indices - numel = 1 - for idx, dim_size in zip(indices, p.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) - rank_numels[r] = numel - - param_to_state[id(p)] = _muon_state( - worker_rank=worker_rank, - process_group=shard_pg, - rank_indices=rank_indices, - rank_numels=rank_numels, - name=n, - qk_clip_state=qk_clip_state, - ) - - return param_to_state, ordered_params - - def base(self, names, params, group, lr, weight_decay, qk_logits): - # Momentum is already applied by _step_muon before this method. - for n, p in zip(names, params): - g = p.grad - if g is None: - continue - - u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) - - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - def distributed_muon( - self, - names: list[str], - params: list[torch.nn.Parameter], - group: dict[str, Any], - lr: float, - weight_decay: float, - qk_logits: list[torch.Tensor | DTensor] | None, - ): - """Batched Distributed Muon — for testing/correctness verification only. - - Uses all-gather to reconstruct full tensors, computes Newton-Schulz on - the full grad, then slices back to local shards. This is simpler but - slower than the parallel pipeline (all2all) path, so it serves as a - reference implementation for verifying correctness. - """ - with record_function("distributed_muon"): - # Momentum is already applied by _step_muon before this method. - ns_steps = group["ns_steps"] - - # Separate plain tensors (no communication) from DTensors. - plain_names, plain_params = [], [] - dtensor_names, dtensor_params = [], [] - for n, p in zip(names, params): - if p.grad is None: - continue - if isinstance(p.data, DTensor): - dtensor_names.append(n) - dtensor_params.append(p) - else: - plain_names.append(n) - plain_params.append(p) - - # Process plain tensors per-param (no communication). - for n, p in zip(plain_names, plain_params): - u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE), - steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, - qk_logits) - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - if not dtensor_params: - return - - # Group DTensors by (placements, mesh) for batched all-gather. - placement_groups: dict[tuple, - tuple[list, - list]] = defaultdict(lambda: ([], [])) - for n, p in zip(dtensor_names, dtensor_params): - key = (p.placements, p.device_mesh) - placement_groups[key][0].append(n) - placement_groups[key][1].append(p) - - logger.info( - "distributed_muon: %d placement groups, %d total dtensors", - len(placement_groups), len(dtensor_params)) - - for (placements, mesh), (grp_names, - grp_params) in placement_groups.items(): - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - placements, mesh) - rank = dist.get_rank(shard_pg) - world_size = dist.get_world_size(shard_pg) - - logger.info(" group: %d params, placements=%s, world_size=%d", - len(grp_params), placements, world_size) - - # Separate params that can be batched (all shard dims evenly - # divisible) from those needing per-param full_tensor - # (e.g. MoE gate weights with fewer rows than shard ranks). - # all_gather_into_tensor requires equal buffer sizes across - # ranks, so uneven splits must use DTensor full_tensor(). - batch_names, batch_params = [], [] - single_names, single_params = [], [] - for n, p in zip(grp_names, grp_params): - even = all(p.shape[pl.dim] % - shard_mesh.mesh.shape[dim_idx] == 0 - for dim_idx, pl in enumerate(shard_placements)) - if even: - batch_names.append(n) - batch_params.append(p) - else: - single_names.append(n) - single_params.append(p) - - # Process uneven-split params per-param via full_tensor(). - for n, p in zip(single_names, single_params): - with record_function("distributed_muon::newton_schulz"): - g_full = p.grad.full_tensor().to(COMM_DTYPE) - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - if not batch_params: - continue - - logger.info(" batched=%d, single=%d", len(batch_params), - len(single_params)) - - # Concat all local grad shards into a single flat buffer. - with record_function("distributed_muon::gather"): - grad_locals = [ - p.grad.to_local().to(COMM_DTYPE).flatten() - for p in batch_params - ] - numels = [g.numel() for g in grad_locals] - grad_concat = torch.cat(grad_locals) - del grad_locals - - # Single all-gather (replaces N separate full_tensor). - grad_gathered = torch.empty( - grad_concat.numel() * world_size, - dtype=COMM_DTYPE, - device="cuda", - ) - dist.all_gather_into_tensor(grad_gathered, - grad_concat, - group=shard_pg) - - total_numel = grad_concat.numel() - del grad_concat - - # Precompute per-param offsets within the concat buffer. - offsets = [] - off = 0 - for ne in numels: - offsets.append(off) - off += ne - - # Per-param: reconstruct full grad → NS → local update. - for i, (n, p) in enumerate(zip(batch_names, batch_params)): - with record_function("distributed_muon::newton_schulz"): - g_full = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - for r in range(world_size): - r_start = r * total_numel + offsets[i] - shard = grad_gathered[r_start:r_start + numels[i]] - indices = get_slices_of_dtensor( - p, r, shard_mesh, shard_placements) - g_full[indices] = shard.reshape( - g_full[indices].shape) - - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - def _setup_parallel(self, names, params, group, qk_logits): - """Compute (or retrieve cached) parallel pipeline metadata. - - Returns: - (ordered_params, param_to_state, rank, chunk_size) - """ - cache_key = tuple(names) - - if cache_key not in self._parallel_cache: - # First call: compute metadata and populate cache. - param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - shard_pg = param_to_state[id(ordered_params[0])].process_group - rank = dist.get_rank(group=shard_pg) - - if self.chunk_size == -1: - shard_ranks = dist.get_world_size(shard_pg) - chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO - elif self.chunk_size > 0: - chunk_size = self.chunk_size - else: - raise ValueError( - "chunk_size must be -1 or a positive integer.") - - ordered_names = [ - param_to_state[id(p)].name for p in ordered_params - ] - name_to_state = { - param_to_state[id(p)].name: param_to_state[id(p)] - for p in ordered_params - } - self._parallel_cache[cache_key] = { - 'ordered_names': ordered_names, - 'name_to_state': name_to_state, - 'rank': rank, - 'chunk_size': chunk_size, - } - else: - # Cached path: rebuild param_to_state with current id(p) keys. - cache = self._parallel_cache[cache_key] - rank = cache['rank'] - chunk_size = cache['chunk_size'] - - name_to_param = dict(zip(names, params)) - ordered_params = [name_to_param[n] for n in cache['ordered_names']] - - param_to_state = {} - for p, n in zip(ordered_params, cache['ordered_names']): - cached_state = cache['name_to_state'][n] - param_to_state[id(p)] = _muon_state( - worker_rank=cached_state.worker_rank, - process_group=cached_state.process_group, - rank_indices=cached_state.rank_indices, - rank_numels=cached_state.rank_numels, - name=n, - qk_clip_state=get_qk_clip_info(self.clip_config, n, - qk_logits), - ) - - return ordered_params, param_to_state, rank, chunk_size - - def parallel(self, - names, - params, - group, - lr, - weight_decay, - qk_logits, - prelaunch_gather=None): - """ - Perform a parallel optimization step using Muon. - - Parameters are chunked and each chunk is processed by a - :func:`muon_chunk_pipeline` generator. :func:`run_pipeline` - interleaves multiple chunks so that communication and computation - overlap across chunks (the same overlap previously achieved by the - warmup + main-loop index scheduling). - - If ``prelaunch_gather`` is provided, it is passed to the first - chunk's generator to skip re-launching the already in-flight - A2A gather. - """ - - # Momentum is already applied by _step_muon before this method. - - ordered_params, param_to_state, rank, chunk_size = ( - self._setup_parallel(names, params, group, qk_logits)) - - def pipelines(): - first = True - for start in range(0, len(ordered_params), chunk_size): - chunk = ordered_params[start:start + chunk_size] - if chunk: - kwargs = dict( - params=chunk, - param_to_state=param_to_state, - rank=rank, - ns_steps=group["ns_steps"], - lr=lr, - weight_decay=weight_decay, - none_grad=group["none_grad"], - ) - if first and prelaunch_gather is not None: - kwargs['prelaunch_gather'] = prelaunch_gather - first = False - yield muon_chunk_pipeline(**kwargs) - - with record_function("muon::pipeline"): - run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1) - - def _step_muon(self, group, qk_logits=None): - params = group["params"] - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - # Apply momentum to all params before routing/expansion. - # Batched using _foreach_* ops (compiled, fullgraph=True). - with record_function("muon::momentum"): - active_params = [p for p in params if p.grad is not None] - if active_params: - # Ensure momentum buffers exist (avoid zeros_like when already present). - for p in active_params: - if "momentum_buffer" not in self.state[p]: - self.state[p]["momentum_buffer"] = torch.zeros_like( - p.grad) - - # Extract local tensors for compiled batch function. - local_grads = [ - p.grad._local_tensor - if isinstance(p.grad, DTensor) else p.grad - for p in active_params - ] - local_bufs = [ - self.state[p]["momentum_buffer"]._local_tensor - if isinstance(self.state[p]["momentum_buffer"], DTensor) - else self.state[p]["momentum_buffer"] - for p in active_params - ] - - # Wrap momentum as tensor for torch.compile. - batch_pre_ortho(local_grads, local_bufs, - torch.tensor(momentum), group["nesterov"]) - - # For non-nesterov, the result is the momentum buffer. - if not group["nesterov"]: - for p in active_params: - p.grad = self.state[p]["momentum_buffer"] - - # Identify batched experts for deferred NS. - # Detection is cheap (condition checks only); actual NS compute is - # deferred so it can overlap with the first chunk's A2A gather. - deferred_expert_work = [] - if self.expert_keys: - batched_expert_indices = [] - for i, (n, p) in enumerate(zip(names, params)): - if not (is_expert_param(n, self.expert_keys) - and p.grad is not None): - continue - # Eligible: plain tensor, or DTensor with no non-dim-0 shards. - if isinstance(p.data, DTensor): - has_tp = any( - _is_shard(pl) and pl.dim != 0 for pl in p.placements) - if has_tp: - continue - batched_expert_indices.append(i) - - if batched_expert_indices: - # Save refs for deferred NS; free grads from param list. - for i in batched_expert_indices: - p = params[i] - g = p.grad - local_g = (g._local_tensor - if isinstance(g, DTensor) else g) - local_data = (p.data._local_tensor if isinstance( - p.data, DTensor) else p.data) - deferred_expert_work.append((local_data, local_g)) - p.grad = None - - # Remove batched experts from lists before expansion. - keep = sorted( - set(range(len(params))) - set(batched_expert_indices)) - names = [names[i] for i in keep] - params = [params[i] for i in keep] - - def _run_deferred_expert_ns(): - """Execute deferred batched expert NS.""" - if not deferred_expert_work: - return - with record_function("muon::batched_expert_ns"): - ns_steps = group["ns_steps"] - for local_data, local_g in deferred_expert_work: - u = zeropower_via_newtonschulz5_batched( - local_g.to(COMM_DTYPE), steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:]) - local_data.mul_(1 - lr * weight_decay) - local_data.add_(u, alpha=-adjusted_lr) - - # Expand expert params by splitting on dim 0. - logger.debug("[_step_muon] before expand: %d params, expert_keys=%s", - len(params), self.expert_keys) - if self.expert_keys: - cache_key = tuple(id(p) for p in params) - cache = self._expert_expand_cache.get(cache_key) - - if cache is None: - # Cold path: full expansion + build cache metadata. - exp_names, exp_params = _expand_expert_params( - names, params, self.expert_keys) - - # Build per-expert-group info for hot-path grad updates. - grad_info = [] - exp_idx = 0 - for orig_idx, (n, p) in enumerate(zip(names, params)): - if not is_expert_param(n, self.expert_keys): - exp_idx += 1 - continue - - is_dt = isinstance(p.data, DTensor) - num_experts = (p.to_local() if is_dt else p.data).shape[0] - - # Detect TP mesh from the first expanded expert param. - tp_mesh = None - tp_pls = None - sample = exp_params[exp_idx] - if isinstance(sample.data, DTensor): - tp_mesh = sample.data.device_mesh - tp_pls = list(sample.data.placements) - - grad_info.append((orig_idx, num_experts, exp_idx, is_dt, - tp_mesh, tp_pls)) - exp_idx += num_experts - - self._expert_expand_cache[cache_key] = { - 'names': exp_names, - 'params': exp_params, - 'grad_info': grad_info, - } - names, params = exp_names, exp_params - else: - # Hot path: reuse cached params, only update expert grads. - for (orig_idx, num_experts, exp_start, is_dt, tp_mesh, - tp_pls) in cache['grad_info']: - p = params[orig_idx] - g = p.grad - local_grad = (g.to_local() - if is_dt and isinstance(g, DTensor) else g) - for i in range(num_experts): - expert_p = cache['params'][exp_start + i] - sg = local_grad[i] - if tp_mesh is not None: - expert_p.grad = DTensor.from_local( - sg, device_mesh=tp_mesh, placements=tp_pls) - else: - expert_p.grad = sg - p.grad = None - - names = cache['names'] - params = cache['params'] - else: - names, params = _expand_expert_params(names, params, - self.expert_keys) - logger.debug("[_step_muon] after expand: %d params", len(params)) - - param_dtensors = [] - name_dtensors = [] - - param_tensors = [] - name_tensors = [] - - # distributed_muon is a reference implementation for testing only. - # The parallel pipeline (all2all) path below is the production path. - if self.use_distributed_muon: - _run_deferred_expert_ns() - self.distributed_muon(names=names, - params=params, - group=group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits) - return - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - logger.debug( - "[route] %s → base (DTensor all-Replicate), " - "shape=%s, placements=%s", n, p.shape, p.placements) - param_tensors.append(p) - name_tensors.append(n) - else: - logger.debug( - "[route] %s → parallel (DTensor), shape=%s, " - "placements=%s, mesh=%s", n, p.shape, p.placements, - p.device_mesh.mesh_dim_names) - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - logger.debug("[route] %s → base (plain tensor), shape=%s", n, - p.data.shape) - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError(f"Unsupported parameter type: {type(p.data)}") - - logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, " - f"{len(param_tensors)} Tensors → base") - - def group_dtensors(dtensors, names): - # To support different placements, we group parameters by placements - # and run parallel Muon on each group. - - placement_to_params = defaultdict(lambda: ([], [])) - - assert len(dtensors) == len(names) - for p, n in zip(dtensors, names): - placement_to_params[tuple([p.placements, - p.device_mesh])][0].append(n) - placement_to_params[tuple([p.placements, - p.device_mesh])][1].append(p) - return placement_to_params - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - dtensor_group = group_dtensors(param_dtensors, name_dtensors) - - # Pre-launch the first chunk's A2A gather so that the NCCL - # communication overlaps with the (deferred) batched expert NS - # compute on the default CUDA stream. - prelaunch = None - if deferred_expert_work: - first_names, first_params = next(iter(dtensor_group.values())) - ordered, pts, rnk, csz = self._setup_parallel( - first_names, first_params, group, qk_logits) - first_chunk = ordered[:csz] - if first_chunk: - prelaunch = prelaunch_first_gather(first_chunk, pts, rnk, - group["none_grad"]) - - _run_deferred_expert_ns() - - first_group = True - for _, (names, params) in dtensor_group.items(): - pg = prelaunch if first_group else None - first_group = False - self.parallel( - names, - params, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - prelaunch_gather=pg, - ) - else: - _run_deferred_expert_ns() - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - ) - - def _register_states_for_offload(self): - """Register all optimizer state tensors with the CPU offload pool. - - Called once after the first step when states have been lazily created. - Offloads all param states (momentum buffers for Muon, moment1/moment2 - for AdamW) to free GPU memory between steps. - """ - pool = self._cpu_offload_pool - tracked = 0 - for group in self.param_groups: - for p in group["params"]: - if p not in self.state: - continue - state = self.state[p] - if group.get("use_muon", False): - if "momentum_buffer" in state: - pool.track(state["momentum_buffer"]) - tracked += 1 - else: - if "moment1" in state: - pool.track(state["moment1"]) - if "moment2" in state: - pool.track(state["moment2"]) - tracked += 1 - logger.info("[CPUOffload] Registered %d param states for offload", - tracked) - - @torch.no_grad - def step(self, closure=None, qk_logits=None): - """Perform a single optimization step. - - Args: - closure (Callable, optional): A closure that reevaluates the model - and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). - """ - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # H2D: reload optimizer states from CPU before computation. - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload() - - logger.debug("[Muon.step] expert_keys=%s, %d param groups", - self.expert_keys, len(self.param_groups)) - - for i, group in enumerate(self.param_groups): - if group["use_muon"]: - logger.debug("[Muon.step] group %d: use_muon=True, %d params", - i, len(group["params"])) - self._step_muon(group, qk_logits=qk_logits) - else: - logger.debug( - "[Muon.step] group %d: use_muon=False (AdamW), %d params", - i, len(group["params"])) - step_adamw(self.state, group) - - # D2H: offload optimizer states to CPU after computation. - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - return loss - - # ------------------------------------------------------------------ - # CPU offload public helpers - # ------------------------------------------------------------------ - - def turn_on_cpu_offload(self): - """Enable CPU offload for optimizer states.""" - if self.cpu_offload: - return - logger.info("[Muon] turn_on_cpu_offload") - self.cpu_offload = True - if not self.state: - return - self._cpu_offload_pool = CPUOffloadPool() - self._offload_initialized = False - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_off_cpu_offload(self): - """Disable CPU offload and keep optimizer states resident on GPU.""" - if not self.cpu_offload: - return - logger.info("[Muon] turn_off_cpu_offload") - if self._offload_initialized: - self._cpu_offload_pool.reload() - torch.cuda.current_stream().synchronize() - self._cpu_offload_pool = None - self._offload_initialized = False - self.cpu_offload = False - - # ------------------------------------------------------------------ - # Checkpoint support for cpu_offload - # ------------------------------------------------------------------ - - def state_dict(self) -> dict: - if self.cpu_offload: - raise RuntimeError( - "Muon.state_dict() requires turn_off_cpu_offload() before checkpoint save." - ) - return super().state_dict() - - def load_state_dict(self, state_dict: dict) -> None: - if self.cpu_offload: - raise RuntimeError( - "Muon.load_state_dict() requires turn_off_cpu_offload() before checkpoint load." - ) - super().load_state_dict(state_dict) - - # Invalidate adamw.py's module-level tensor caches so that - # the next step rebuilds them with the newly loaded state tensors. - _placement_cache.clear() - _tensor_cache.clear() diff --git a/build/torch28-cxx11-cu126-x86_64-linux/newton_schulz.py b/build/torch28-cxx11-cu126-x86_64-linux/newton_schulz.py deleted file mode 100644 index d939264b69a34e7a3fa78859f34dc265a1159d59..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/newton_schulz.py +++ /dev/null @@ -1,240 +0,0 @@ -from itertools import repeat -from math import inf, sqrt - -import numpy as np -import torch - -from .matmul_transpose_triton import matmul_transpose_assign - -COMM_DTYPE = torch.bfloat16 -DEFAULT_CHUNK_SIZE_RATIO = 4 - - -def _optimal_quintic(l, u, max_iter=1000): - """ - Use the simplified Remez algorithm to find the optimal odd quintic approximant - to the constant function x -> 1 over the interval [l, u]. - - Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum - approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the - two interior equioscillation nodes q, r until convergence. Returns the - closed-form equioscillating solution when l ≈ u. - - Raises ValueError if any intermediate value (a, b, c, E, q, r) is non-finite - (NaN or inf). Raises RuntimeError if convergence is not reached within - max_iter iterations. - """ - assert 0 <= l <= u - if 1 - 5e-6 <= l / u: - return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5) - q = (3 * l + u) / 4 - r = (l + 3 * u) / 4 - E = inf - for _ in range(max_iter): - old_E = E - LHS = np.array( - [ - [l, l**3, l**5, 1], - [q, q**3, q**5, -1], - [r, r**3, r**5, 1], - [u, u**3, u**5, -1], - ] - ) - a, b, c, E = np.linalg.solve(LHS, np.ones(4)) - if not np.all(np.isfinite([a, b, c, E])): - raise ValueError( - f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}" - ) - q, r = np.sqrt( - (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / (10 * c) - ) - if not np.all(np.isfinite([q, r])): - raise ValueError(f"_optimal_quintic: non-finite node update q={q}, r={r}") - if abs(old_E - E) <= 1e-15: - break - else: - raise RuntimeError( - f"_optimal_quintic: did not converge after {max_iter} iterations" - ) - return float(a), float(b), float(c) - - -def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0): - """ - Compute the Polar Express coefficient series for `num_iters` quintic iterations. - - Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that - compose to map singular values from [l, 1] toward 1. At each step: - 1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion` - prevents near-zero singular values from stalling by raising the effective - lower bound; if it is active (cushion*u > l), the coefficients are - rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u]. - 2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the - last iteration, providing numerical headroom at the cost of a slightly slower - final convergence step. - 3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1). - - Returns a list of (a, b, c) tuples, one per iteration. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - u = 1 - assert 0 <= l <= u - safety_factor = 1 + safety_factor_eps - coefficients = [] - for iter in range(num_iters): - a, b, c = _optimal_quintic(max(l, cushion * u), u) - if cushion * u > l: - pl = a * l + b * l**3 + c * l**5 - pu = a * u + b * u**3 + c * u**5 - rescaler = 2 / (pl + pu) - a *= rescaler - b *= rescaler - c *= rescaler - if iter < num_iters - 1: - a /= safety_factor - b /= safety_factor**3 - c /= safety_factor**5 - coefficients.append((a, b, c)) - l = a * l + b * l**3 + c * l**5 - u = 2 - l - return coefficients - - -# Precomputed Polar Express coefficients (a, b, c) for 10 quintic Newton-Schulz -# iterations. Each tuple is the minimax-optimal (Remez/equioscillation) odd quintic -# approximant to x->1 over the current singular-value interval, computed once at -# import time and reused across all optimizer steps. -# -# Contrast with the former hardcoded NS coefficients (5 fixed tuples): -# - Former: empirically tuned to maximize slope at zero; did not converge -# singular values to 1, yielding US'V^T with S' ~ Uniform(0.5, 1.5) instead -# of the true polar factor UV^T. -# - Polar Express: analytically optimal per step, adapting to the shrinking -# singular-value interval [l, u] as iterations progress; converges all -# singular values to 1, producing the exact polar factor UV^T. -_coeffs_list = _optimal_composition( - l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02 -) - - -# This code is adapted from: -# KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py) -# NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress) -# matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon) -@torch.no_grad() -def _zeropower_via_newtonschulz5(G, steps): - """ - Compute the polar factor of G via the Polar Express method. - - Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c) - are the Polar Express coefficients from `_coeffs_list`. Each step is the - optimal odd quintic approximant to x -> 1 over the current singular-value - interval, minimizing the maximum approximation error (Remez / minimax criterion). - The composition maps singular values from [l, 1] to near 1, producing the - polar factor (orthogonal factor in the polar decomposition G = UP). - - `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2, - cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE - X = G # no manual typecast - - if G.size(0) > G.size(1): - X = X.T - - X = X / (X.norm() + 1e-7) - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - # Perform the NS iterations - for a, b, c in hs: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(0) > G.size(1): - X = X.T - - return X - - -@torch.no_grad() -def _zeropower_via_newtonschulz5_batched(G, steps): - """Batched polar factor computation for 3D (E, out, in) tensors. - - Same algorithm as ``_zeropower_via_newtonschulz5`` but uses - ``torch.bmm`` / ``torch.baddbmm`` instead of the 2D Triton kernel, - processing all E expert matrices in a single batched call. - """ - assert len(G.shape) == 3 - assert G.dtype == COMM_DTYPE - X = G - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - # Per-expert Frobenius norm. - X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) - - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - for a, b, c in hs: - buf1 = torch.bmm(X, X.transpose(-2, -1)) - buf2 = torch.bmm(buf1, buf1.transpose(-2, -1)) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.baddbmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - return X - - -_ns_per_shape: dict[tuple[int, ...], callable] = {} -_use_compile = True - - -def set_ns_compile(enabled: bool): - """Toggle torch.compile for Newton-Schulz iteration.""" - global _use_compile - _use_compile = enabled - - -def zeropower_via_newtonschulz5(G, steps=5): - if not _use_compile: - return _zeropower_via_newtonschulz5(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() - - -def zeropower_via_newtonschulz5_batched(G, steps=5): - """Compile-cached batched Newton-Schulz for 3D expert tensors.""" - if not _use_compile: - return _zeropower_via_newtonschulz5_batched(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile( - _zeropower_via_newtonschulz5_batched, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() diff --git a/build/torch28-cxx11-cu126-x86_64-linux/optimizer/__init__.py b/build/torch28-cxx11-cu126-x86_64-linux/optimizer/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/optimizer/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch28-cxx11-cu126-x86_64-linux/pipeline.py b/build/torch28-cxx11-cu126-x86_64-linux/pipeline.py deleted file mode 100644 index c0c2d515856182d8d15ad27dd4e4e093b29397d6..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/pipeline.py +++ /dev/null @@ -1,468 +0,0 @@ -import logging -from typing import Generator - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -from .core import _muon_state, adjust_lr_for_muon -from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5 -from .qk_clip import compute_scales - -logger = logging.getLogger(__name__) - -# ====================================================================== -# Stage helpers -# ====================================================================== - - -def _launch_gather( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Allocate gather buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_gather``). - gathered_grads: ``{id(p): empty_tensor}`` for owned params, - ``None`` for non-owned. - recv_counts: Per-source-rank element counts. - """ - # Allocate gathered-grad buffers - gathered_grads: dict[int, torch.Tensor | None] = {} - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - gathered_grads[id(p)] = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - else: - gathered_grads[id(p)] = None - - # Build send buffer – batch grad copies via torch.cat - # (1-2 fused kernels vs N individual narrow().copy_() calls). - send_counts = [0] * num_ranks - for p in params: - state = param_to_state[id(p)] - send_counts[state.worker_rank] += state.rank_numels[rank] - - total_send = sum(send_counts) - if total_send > 0: - # Group grad slices by destination rank in a single pass. - dst_to_grads = [[] for _ in range(num_ranks)] - for p in params: - state = param_to_state[id(p)] - n = state.rank_numels[rank] - if n > 0: - g = p.grad.to_local() - dst_to_grads[state.worker_rank].append(g.reshape(-1)) - - # Flatten in dst order and cat once. - all_slices = [] - for dst in range(num_ranks): - all_slices.extend(dst_to_grads[dst]) - send_buf = torch.cat(all_slices) - if send_buf.dtype != COMM_DTYPE: - send_buf = send_buf.to(COMM_DTYPE) - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += state.rank_numels[src] - recv_counts[src] = total - - recv_buf = torch.empty(sum(recv_counts), dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - logger.debug(f"send_buf size: {send_buf.numel()}, " - f"recv_buf size: {recv_buf.numel()}, " - f"recv_counts: {recv_counts}, " - f"send_counts: {send_counts}, " - f"process_group: {str(process_group)}") - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, gathered_grads, recv_counts - - -def _complete_gather( - recv_buf: torch.Tensor, - recv_counts: list[int], - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - param_to_state: dict[int, _muon_state], - rank: int, -) -> None: - """Reconstruct gathered grads from the recv buffer (in-place).""" - off = 0 - for src in range(len(recv_counts)): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - - indices = state.rank_indices[src] - - shard_view = gathered_grads[id(p)][indices] - n = shard_view.numel() - if n == 0: - continue - - sg = recv_buf.narrow(0, off + inner_off, n) - sg = sg.reshape(shard_view.shape) - gathered_grads[id(p)][indices] = sg - - inner_off += n - assert inner_off == block - off += block - - -def _compute_ns( - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - ns_steps: int, -) -> dict[int, torch.Tensor | None]: - """Run Newton-Schulz orthogonalization on owned parameters. - - Returns: - computed_us: ``{id(p): orthogonalized_update}`` for owned params. - """ - computed_us: dict[int, torch.Tensor | None] = {} - for p in owned_params: - u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps) - gathered_grads[id(p)] = None # free gathered grad - computed_us[id(p)] = u - return computed_us - - -def _launch_scatter( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, - computed_us: dict[int, torch.Tensor | None], -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor], list[int]]: - """Allocate scatter buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_scatter``). - scattered_us: Empty dict, populated by ``_complete_scatter`` with - zero-copy views into ``recv_buf``. - recv_counts: Per-source-rank element counts. - """ - # scattered_us is populated by _complete_scatter with zero-copy views - # into recv_buf, avoiding N empty_like allocations + N copy_ calls. - # Pre-seed entries for params whose local shard is empty (rank_numels == 0) - # so _update_params can iterate all params without KeyError. - scattered_us: dict[int, torch.Tensor] = {} - for p in params: - if param_to_state[id(p)].rank_numels[rank] == 0: - scattered_us[id(p)] = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - # Build send buffer – batch via torch.cat - # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls). - send_counts = [0] * num_ranks - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - for dst_rank in range(num_ranks): - send_counts[dst_rank] += state.rank_numels[dst_rank] - - total_send = sum(send_counts) - if total_send > 0: - # Cache u_full conversions to avoid redundant .to() per dst_rank. - u_fulls = {} - for p in owned_params: - u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous() - - # Collect slices in dst order (matches all-to-all send layout). - all_slices = [] - for dst_rank in range(num_ranks): - for p in owned_params: - state = param_to_state[id(p)] - su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten() - if su.numel() > 0: - all_slices.append(su) - - send_buf = torch.cat(all_slices) if all_slices else torch.empty( - 0, dtype=COMM_DTYPE, device="cuda") - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += state.rank_numels[rank] - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, scattered_us, recv_counts - - -def _complete_scatter( - recv_buf: torch.Tensor, - recv_counts: list[int], - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], -) -> None: - """Populate scattered_us with zero-copy views into recv_buf. - - Instead of pre-allocating tensors and copying, we assign views directly - from ``recv_buf``. This eliminates N ``empty_like`` + N ``copy_`` calls. - The underlying storage of ``recv_buf`` is kept alive through the views - until ``scattered_us`` is cleared after ``_update_params``. - """ - off = 0 - for src in range(len(recv_counts)): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = state.rank_numels[rank] - if n == 0: - continue - - scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - - inner_off += n - - assert inner_off == block - off += block - - -def _update_params( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], - lr: float, - weight_decay: float, -) -> None: - """Apply weight decay, Muon update, and optional QK clipping. - - Uses batched ``_foreach_mul_`` for weight decay and batched - ``_foreach_add_`` for the Muon update, grouping parameters by - adjusted_lr to minimize kernel launches while preserving float32 - precision for the alpha scaling. - """ - if not params: - return - - # Batched weight decay: p *= (1 - lr * wd) — single fused kernel. - p_locals = [p._local_tensor for p in params] - torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay) - - # Group params by adjusted_lr so _foreach_add_ can use a single - # alpha per group (preserves float32 precision for alpha scaling). - lr_groups: dict[float, tuple[list, list]] = {} - for p in params: - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - if adjusted_lr not in lr_groups: - lr_groups[adjusted_lr] = ([], []) - lr_groups[adjusted_lr][0].append(p._local_tensor) - lr_groups[adjusted_lr][1].append(scattered_us[id(p)]) - - for adjusted_lr, (p_group, u_group) in lr_groups.items(): - torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr) - - # QK clipping – applied directly on the local tensor to - # avoid DTensor sharding-propagation issues with _StridedShard. - for p in params: - state = param_to_state[id(p)] - if state.qk_clip_state is None: - continue - scales_full = compute_scales(p, state.qk_clip_state) - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = state.rank_indices[rank][0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - -# ====================================================================== -# Pre-launch helper for overlapping first chunk's gather with other work. -# ====================================================================== - - -@torch.no_grad() -def prelaunch_first_gather( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - none_grad: bool, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Launch the first chunk's A2A gather early for overlap with other compute. - - Call this *before* expensive GPU work (e.g. batched expert NS) so that - the NCCL all-to-all runs concurrently on the NCCL stream while the - default stream executes compute. - - Returns the same 4-tuple that ``_launch_gather`` produces, which should - be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - with record_function("muon::prelaunch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - return work, recv_buf, gathered_grads, recv_counts - - -# ====================================================================== -# Main generator – thin orchestrator that wires stages together. -# ====================================================================== - - -@torch.no_grad() -def muon_chunk_pipeline( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - ns_steps: int, - lr: float, - weight_decay: float, - none_grad: bool, - prelaunch_gather: tuple | None = None, -) -> Generator[None, None, None]: - """Process one chunk of parameters through the full Muon pipeline. - - Stages: gather -> compute (Newton-Schulz) -> scatter -> update. - - Each ``yield`` lets :func:`run_pipeline` interleave other chunks so - that communication and computation overlap across chunks. Async - communication is launched via ``async_op=True`` and completed after - the yield with ``work.wait()``. - - Overlap happens because :func:`run_pipeline` admits one new chunk - per iteration (staggered admission). While chunk *N* does NS - compute on the default CUDA stream, chunk *N+1*'s async all-to-all - runs concurrently on the NCCL stream — no separate ``comm_stream`` - is required. - - If ``prelaunch_gather`` is provided, the gather was already launched - by :func:`prelaunch_first_gather` and we skip launching it again. - - Yields exactly **2** times: - - 1. After launching async all-to-all gather (or immediately if pre-launched). - 2. After launching async all-to-all scatter. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - if prelaunch_gather is not None: - # Gather was pre-launched; none_grad already handled by caller. - work, recv_buf, gathered_grads, recv_counts = prelaunch_gather - else: - # Normal path: launch async gather. - with record_function("muon::launch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - yield # --- YIELD 1: other chunks can launch their gather --- - - with record_function("muon::wait_gather"): - work.wait() - _complete_gather(recv_buf, recv_counts, owned_params, gathered_grads, - param_to_state, rank) - del recv_buf - - # Stage 3: Newton-Schulz orthogonalization. - with record_function("muon::newton_schulz"): - computed_us = _compute_ns(owned_params, gathered_grads, ns_steps) - gathered_grads.clear() - - # Stages 4-5: launch async scatter. - with record_function("muon::launch_scatter"): - work, recv_buf, scattered_us, recv_counts = _launch_scatter( - params, owned_params, param_to_state, rank, num_ranks, - process_group, computed_us) - computed_us.clear() - - yield # --- YIELD 2: other chunks can launch their scatter --- - - with record_function("muon::wait_scatter"): - work.wait() - _complete_scatter(recv_buf, recv_counts, params, param_to_state, rank, - scattered_us) - del recv_buf - - # Stage 6: apply parameter updates. - with record_function("muon::update_params"): - _update_params(params, param_to_state, rank, scattered_us, lr, - weight_decay) - scattered_us.clear() diff --git a/build/torch28-cxx11-cu126-x86_64-linux/qk_clip.py b/build/torch28-cxx11-cu126-x86_64-linux/qk_clip.py deleted file mode 100644 index 2aba711b3004b7f09e7141da7ef834bd61cc2430..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/qk_clip.py +++ /dev/null @@ -1,198 +0,0 @@ -import logging -import math -from dataclasses import dataclass - -import torch -from torch.distributed.tensor import DTensor - -from .core import normalize_fqn - -logger = logging.getLogger(__name__) - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - and return (kind, layer_index). - - Supported kinds: - MHA/GQA: 'wq', 'wk', 'q_proj', 'k_proj' - MLA: 'wq_b' (Q up-proj), 'wkv_b' (KV up-proj) - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.1.attn.wq_b.weight' -> ('wq_b', 1) - 'model.0.attn.wkv_b.weight' -> ('wkv_b', 0) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = normalize_fqn(name).split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: str | None # 'wq'/'q_proj'/'wq_b' or 'wk'/'k_proj'/'wkv_b' or None - indices: list[int] # which heads to consider for clipping - head_dim: int # from config (qk_head_dim for MLA wq_b) - threshold: float # from config - logit: torch.Tensor | None - - # MLA-specific fields - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 - - -def get_qk_clip_info(clip_config, n, qk_logits): - """Extract QK clipping info for a named parameter. - - Args: - clip_config: QK clipping configuration dict (or None). - MHA/GQA keys: head_dim, threshold, q_indices, k_indices - MLA extra keys: is_mla=True, qk_nope_head_dim, qk_rope_head_dim, v_head_dim - n: Parameter name string. - qk_logits: Dict mapping layer indices to logit tensors (or None). - - Returns: - QKClipInfo instance with clipping configuration for this parameter. - """ - if clip_config is None: - return None - - head_dim = clip_config.get('head_dim') - threshold = clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - # In TP settings, qk_logits may be DTensor - # We convert it to full tensor here for simplicity - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - -def compute_scales(p, qk_clip_state): - """Compute per-head scaling factors for QK clipping. - - Returns scales tensor (√γ per head) if any head exceeds threshold, else None. - For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim. - """ - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - # Check if any head exceeds threshold before allocating. - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - - if not head_scales: - return None - - # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows - if qk_clip_state.is_mla and kind == 'wkv_b': - effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim - else: - effective_head_dim = head_dim - - H_global = p.shape[0] // effective_head_dim - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale - return scales_full - - -def qk_clip(p, scales, info): - """Apply per-head scaling to a Q/K projection weight matrix. - - Args: - p: Parameter (nn.Parameter or raw tensor). - scales: [n_heads] tensor, each element = √γ_h. - info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions. - - MLA sub-region scaling per Algorithm 1 (MuonClip): - wq_b: q_nope rows → √γ, q_pe rows → γ - wkv_b: k_nope rows → √γ, v rows → unchanged - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not info.is_mla: - # MHA/GQA: uniform √γ applied to all rows in each head - W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: vectorized sub-region scaling within each head - if info.kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_head_dim = qk_nope + info.qk_rope_head_dim - W_3d = W.view(-1, qk_head_dim, W.shape[1]) # [H, qk_head_dim, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # q_nope → √γ - W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1, - 1)) # q_pe → γ - - elif info.kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - W_3d = W.view(-1, kv_stride, W.shape[1]) # [H, kv_stride, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # k_nope → √γ - # v rows: not touched (k_R shared rotary unchanged) diff --git a/build/torch28-cxx11-cu128-x86_64-linux/_ops.py b/build/torch28-cxx11-cu128-x86_64-linux/_ops.py deleted file mode 100644 index 034bff088659b5df6f6d401feb18c89dc5f33b29..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _optimizer_8d53b78_dirty -ops = torch.ops._optimizer_8d53b78_dirty - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_optimizer_8d53b78_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-cu128-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so b/build/torch28-cxx11-cu128-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so deleted file mode 100755 index 99e4220345748e9633a27b083af5e5ac2605be0b..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:090f5a44cdfa4554147159cc36bb7e8ee9dba1ffb1fea4825aa838461fdaddf9 -size 1999872 diff --git a/build/torch28-cxx11-cu128-x86_64-linux/adamw.py b/build/torch28-cxx11-cu128-x86_64-linux/adamw.py deleted file mode 100644 index b5a95816a9f5b9e1889eaadae65373bfbced809a..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/adamw.py +++ /dev/null @@ -1,271 +0,0 @@ -import logging -from collections import defaultdict -from typing import cast - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -def fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: float | torch.Tensor, - weight_decay: float, - eps: float, - maximize: bool, -) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: dict | None = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) - - -def _to_local(t): - """Unwrap DTensor to local tensor for fused ops.""" - return t._local_tensor if isinstance(t, DTensor) else t - - -# --------------------------------------------------------------------------- -# Caches for eliminating per-step Python overhead. -# -# Placement grouping and tensor list assembly are identical every step -# (params don't change placement, moment/step tensors are the same objects -# after initialisation). We cache them keyed by id() of the param list -# stored in param_groups (stable across steps). -# -# Only gradients change each step and must be collected fresh. -# --------------------------------------------------------------------------- - -# id(group["params"]) → dict[placement_key, list[param]] -_placement_cache: dict[int, dict[tuple, list]] = {} - -# id(placement_group_list) → (params_local, moment1, moment2, state_steps) -_tensor_cache: dict[int, tuple[list, list, list, list]] = {} - - -def _step_adamw_params_slow(optimizer_state, params, group): - """Uncached fallback for the rare case where some params lack grads.""" - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - g = p.grad - if g is None: - continue - state = optimizer_state[p] - params_with_grads.append(_to_local(p)) - grads.append(_to_local(g)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - if not params_with_grads: - return - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw_params(optimizer_state, params, group): - """Run fused AdamW on a list of parameters sharing the same placement. - - After the first call, cached tensor lists (params_local, moment1, - moment2, state_steps) are reused — only gradients are collected fresh. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - params: List of parameters to update. - group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay. - """ - # Collect grads — the only thing that changes each step. - with record_function("adamw::collect_grads"): - grads = [] - for p in params: - g = p.grad - if g is None: - # Rare: fall back to slow path that filters per-param. - _step_adamw_params_slow(optimizer_state, params, group) - return - grads.append(_to_local(g)) - - tensor_key = id(params) - if tensor_key not in _tensor_cache: - with record_function("adamw::init_tensor_cache"): - params_local = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - state = optimizer_state[p] - params_local.append(_to_local(p)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(p.grad) - state["moment2"] = torch.zeros_like(p.grad) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - _tensor_cache[tensor_key] = (params_local, moment1, moment2, - state_steps) - - params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key] - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - with record_function("adamw::fused_adamw"): - fused_adamw( - params_local, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw(optimizer_state, group): - """Dispatch AdamW step, grouping parameters by type and placement. - - Placement grouping is cached after the first call since params never - change their placement between steps. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - group: Parameter group dict. - """ - params = group["params"] - placement_key = id(params) - - if placement_key not in _placement_cache: - with record_function("adamw::group_by_placement"): - placement_to_params: dict[tuple, - list[torch.Tensor]] = defaultdict(list) - for p in params: - match p: - case DTensor(): - logger.debug( - "[AdamW] DTensor param: shape=%s, placements=%s, " - "mesh=%s, grad=%s", p.shape, p.placements, - p.device_mesh.mesh_dim_names, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple( - [p.placements, p.device_mesh])].append(p) - case torch.Tensor(): - logger.debug( - "[AdamW] plain param: shape=%s, grad=%s", p.shape, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple([torch.Tensor, - None])].append(p) - - logger.debug("[AdamW] %d placement groups, %d total params", - len(placement_to_params), len(params)) - - _placement_cache[placement_key] = dict(placement_to_params) - - for group_params in _placement_cache[placement_key].values(): - step_adamw_params(optimizer_state, group_params, group) diff --git a/build/torch28-cxx11-cu128-x86_64-linux/async_utils.py b/build/torch28-cxx11-cu128-x86_64-linux/async_utils.py deleted file mode 100644 index a45c530ac9cad88e3555ec1047a6aa59f225347e..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/async_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import logging -from typing import Generator - -logger = logging.getLogger(__name__) - - -class _Task: - """Internal: wraps a generator, advances one yield at a time.""" - - def __init__(self, generator: Generator[None, None, None], index: int): - self._generator = generator - self._index = index - self._steps_completed = 0 - self.step() # run to first yield - - def step(self) -> bool: - try: - next(self._generator) - self._steps_completed += 1 - logger.debug("pipeline[%d] completed stage %d", self._index, - self._steps_completed) - return True - except StopIteration: - logger.debug("pipeline[%d] finished after %d stages", self._index, - self._steps_completed) - return False - - def close(self): - self._generator.close() - - -def run_pipeline( - pipelines: Generator[Generator[None, None, None], None, None], - max_concurrent: int, -) -> None: - """Run generator-based pipelines with bounded concurrency. - - Each pipeline is a generator that yields at stage boundaries. - The runtime interleaves pipelines so communication and computation - overlap across chunks. - """ - if max_concurrent <= 0: - raise ValueError(f"max_concurrent must be > 0, got {max_concurrent}") - - have_new = True - task_index = 0 - previous_tasks: list[_Task] = [] - - try: - while have_new or previous_tasks: - running_tasks: list[_Task] = [] - - # Admit one new pipeline per iteration (staggered admission). - # Admitting one at a time ensures that while chunk N does NS - # compute on the default stream, chunk N+1's NCCL all-to-all - # runs concurrently on the NCCL stream — creating real - # communication/computation overlap on the GPU. - if have_new and len(previous_tasks) < max_concurrent: - try: - gen = next(pipelines) - task = _Task(gen, task_index) - task_index += 1 - running_tasks.append(task) - except StopIteration: - have_new = False - - # Advance every previously-yielded task by one step. - for task in previous_tasks: - if task.step(): - running_tasks.append(task) - - previous_tasks = running_tasks - except BaseException: - # Clean up all in-flight generators to release GPU resources. - for task in previous_tasks: - task.close() - raise diff --git a/build/torch28-cxx11-cu128-x86_64-linux/core.py b/build/torch28-cxx11-cu128-x86_64-linux/core.py deleted file mode 100644 index c69d515afef305ad0ed66374095fa2d2468d99cc..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/core.py +++ /dev/null @@ -1,219 +0,0 @@ -import logging -import math -from dataclasses import dataclass -from typing import List - -import torch -from torch.distributed import ProcessGroup -from torch.distributed.tensor import DTensor - -# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into -# parameter FQNs. Activation checkpointing similarly inserts -# "_checkpoint_wrapped_module". Strip these so name-based matching (skip_keys, -# expert_keys, QK layer parsing) works regardless of wrapper nesting. -_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) - -logger = logging.getLogger(__name__) - - -def normalize_fqn(name: str) -> str: - """Strip torch.compile / checkpoint wrapper components from a parameter FQN.""" - return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS) - - -@dataclass -class _muon_state: - worker_rank: int - process_group: ProcessGroup - rank_indices: dict[int, tuple] # local_rank -> per-dim indices - rank_numels: dict[int, int] # local_rank -> numel - name: str - qk_clip_state: torch.Tensor | None = None - - -def _batch_momentum( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update (no nesterov).""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - - -def _batch_momentum_nesterov( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update with nesterov correction.""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - nesterov_terms = torch._foreach_mul(momentum_bufs, momentum) - torch._foreach_add_(grads, nesterov_terms) - - -_compiled_momentum: dict[bool, callable] = {} -_use_momentum_compile = True - - -def set_momentum_compile(enabled: bool): - """Toggle torch.compile for batched momentum.""" - global _use_momentum_compile - _use_momentum_compile = enabled - - -def batch_pre_ortho( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, - nesterov: bool, -) -> None: - """Batched momentum update on lists of plain tensors. - - Mirrors dion's ``muon_update_pre_orthogonalize``. - Inputs must be plain CUDA tensors (not DTensor). - Modifies ``momentum_bufs`` and (for nesterov) ``grads`` in-place. - - When compile is enabled, uses separately compiled functions for - nesterov=True/False to avoid graph breaks from the branch. - """ - fn = _batch_momentum_nesterov if nesterov else _batch_momentum - if _use_momentum_compile: - if nesterov not in _compiled_momentum: - _compiled_momentum[nesterov] = torch.compile(fn) - fn = _compiled_momentum[nesterov] - fn(grads, momentum_bufs, momentum) - - -def _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay): - """Weight-decay + update on plain tensors. - - Not compiled: per-param @torch.compile caused ~0.25ms TorchDynamo cache - lookup per call × 256+ params = massive overhead. The pipeline path uses - batched _foreach_* ops instead; this function remains for base() and - distributed_muon(). - """ - p_data.mul_(1 - lr * weight_decay) - p_data.add_(u_data, alpha=-adjusted_lr) - - -def update_p(p, u, lr, adjusted_lr, weight_decay): - """Apply weight decay and orthogonalized update to parameter. - - Args: - p: Parameter (torch.nn.Parameter or DTensor). - u: Orthogonalized update tensor. - lr: Base learning rate. - adjusted_lr: Size-adjusted learning rate. - weight_decay: Weight decay coefficient. - """ - # Unwrap Parameter -> underlying data tensor. - p_data = p.data if isinstance(p, torch.nn.Parameter) else p - # Unwrap DTensor -> local CUDA tensor for compiled kernel. - if isinstance(p_data, DTensor): - p_data = p_data._local_tensor - u_data = u._local_tensor if isinstance(u, DTensor) else u - _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay) - - -def adjust_lr_for_muon(lr, param_shape): - """Scale learning rate based on parameter matrix dimensions. - - Args: - lr: Base learning rate. - param_shape: Shape of the parameter tensor. - - Returns: - Adjusted learning rate. - """ - A, B = param_shape[:2] - # We adjust the learning rate and weight decay based on the size of the parameter matrix - # as described in the paper - adjusted_ratio = 0.2 * math.sqrt(max(A, B)) - adjusted_lr = lr * adjusted_ratio - return adjusted_lr - - -def _match_key(parts, key): - """Check if key matches as contiguous components in parts. - - Single-component keys (e.g. "experts") match any single component. - Multi-component keys (e.g. "experts.w1") match as a contiguous subsequence. - """ - key_parts = key.split(".") - key_len = len(key_parts) - if key_len == 1: - return key in parts - return any(parts[i:i + key_len] == key_parts - for i in range(len(parts) - key_len + 1)) - - -def is_expert_param(name, expert_keys): - """Check if a parameter name matches any expert key (component-level).""" - if not expert_keys: - return False - parts = normalize_fqn(name).split(".") - return any(_match_key(parts, key) for key in expert_keys) - - -def default_is_muon(name, x, expert_keys=None): - normalized = normalize_fqn(name) - parts = normalized.split(".") - skip_keys = [ - "embed_tokens", - "lm_head", - "tok_embeddings", - "output", - "mhc_attn", - "mhc_ffn", - "lambda_proj", - ] - if any(key in parts for key in skip_keys): - logger.info( - "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d", - normalized, name, x.ndim) - return False - effective_ndim = x.ndim - is_expert = is_expert_param(name, expert_keys) - if is_expert: - effective_ndim -= 1 - result = effective_ndim >= 2 - logger.info( - "[is_muon] %s (orig: %s): ndim=%d, expert=%s, effective_ndim=%d → %s", - normalized, name, x.ndim, is_expert, effective_ndim, - "Muon" if result else "AdamW") - return result - - -def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None): - if is_muon_func is None: - is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys) - - muon_params, muon_names = [], [] - non_muon_params, non_muon_names = [], [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - non_muon_names.append(n) - - logger.info("[param_groups] expert_keys=%s, Muon=%d, AdamW=%d", - expert_keys, len(muon_names), len(non_muon_names)) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] diff --git a/build/torch28-cxx11-cu128-x86_64-linux/cpu_offload.py b/build/torch28-cxx11-cu128-x86_64-linux/cpu_offload.py deleted file mode 100644 index fb5e69154a1d4a6c884491413a37a9acf0f66c80..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/cpu_offload.py +++ /dev/null @@ -1,206 +0,0 @@ -"""CPU offloading for optimizer states. - -Manages a pinned CPU memory pool and async CUDA streams to offload -optimizer state tensors (momentum buffers, Adam moments) to CPU between -optimizer steps, freeing GPU memory. - -All tracked tensors are packed into a single flat pinned CPU buffer -(per dtype). D2H and H2D copies are performed per-tensor directly -between individual GPU tensors and their slice of the CPU flat buffer -— no GPU staging buffer is allocated, so there is **no temporary GPU -memory spike** during offload or reload. - -Individual tensor storages are freed after offload via -``untyped_storage().resize_(0)``, preserving tensor identity so -downstream caches remain valid. -""" - -import logging -from collections import defaultdict - -import torch -from torch.distributed.tensor import DTensor - -logger = logging.getLogger(__name__) - - -class CPUOffloadPool: - """Pinned CPU memory pool for async optimizer state offloading. - - Tracked tensors are grouped by dtype. Each group gets a single flat - pinned CPU buffer. D2H / H2D copies are per-tensor (into slices of - the flat buffer) to avoid allocating a GPU staging buffer. - """ - - def __init__(self): - self._managed: list[torch.Tensor] = [] - self._storage_nbytes: dict[int, int] = {} # id(t) → bytes - - # Per-dtype group: populated on first offload. - # dtype → dict with keys: - # "indices" : list[int] managed-list indices - # "offsets" : list[tuple[int,int]] (start, numel) in flat buf - # "total" : int total numel - # "cpu_flat" : Tensor pinned CPU buffer - self._groups: dict[torch.dtype, dict] = {} - - self._offload_stream: torch.cuda.Stream | None = None - self._device: torch.device | None = None - self._initialized: bool = False - self._logged: bool = False - - # ------------------------------------------------------------------ - @staticmethod - def _local(t: torch.Tensor) -> torch.Tensor: - """Unwrap DTensor to its local CUDA tensor.""" - return t._local_tensor if isinstance(t, DTensor) else t - - def _ensure_stream(self): - if self._offload_stream is None: - self._offload_stream = torch.cuda.Stream(device=self._device) - - # ------------------------------------------------------------------ - def track(self, tensor: torch.Tensor): - """Register a GPU tensor for CPU offloading. Idempotent.""" - tid = id(tensor) - if tid in self._storage_nbytes: - return - local = self._local(tensor) - if self._device is None: - self._device = local.device - storage = local.untyped_storage() - # Skip tensors with empty storage (e.g. empty FSDP shards) - if storage.size() == 0: - return - self._storage_nbytes[tid] = storage.size() - self._managed.append(tensor) - - # ------------------------------------------------------------------ - def _init_buffers(self): - """Build per-dtype flat buffers on first offload.""" - # Group managed tensors by dtype. - dtype_map: dict[torch.dtype, list[tuple[int, int]]] = defaultdict(list) - for idx, t in enumerate(self._managed): - local = self._local(t) - dtype_map[local.dtype].append((idx, local.numel())) - - total_cpu_bytes = 0 - for dtype, entries in dtype_map.items(): - offsets: list[tuple[int, int]] = [] - indices: list[int] = [] - off = 0 - for idx, n in entries: - indices.append(idx) - offsets.append((off, n)) - off += n - cpu_flat = torch.empty(off, dtype=dtype, device="cpu", pin_memory=True) - self._groups[dtype] = { - "indices": indices, - "offsets": offsets, - "total": off, - "cpu_flat": cpu_flat, - } - total_cpu_bytes += off * cpu_flat.element_size() - - self._initialized = True - logger.info( - "[CPUOffload] Pool initialized: %d tensors, %d dtype group(s), " - "%.2f MB pinned CPU memory", - len(self._managed), - len(self._groups), - total_cpu_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def offload(self): - """Per-tensor async D2H into CPU flat buffer, then free GPU storage.""" - if not self._managed: - return - if not self._initialized: - self._init_buffers() - self._ensure_stream() - - # Offload stream waits for compute to finish. - compute_event = torch.cuda.current_stream(self._device).record_event() - self._offload_stream.wait_event(compute_event) - - offloaded_bytes = 0 - - # Per-tensor D2H copies directly into CPU flat buffer slices. - # No GPU staging buffer → no temporary GPU memory spike. - with torch.cuda.stream(self._offload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - cpu_flat[off : off + n].copy_(local.reshape(-1), non_blocking=True) - - offloaded_bytes += grp["total"] * cpu_flat.element_size() - - # Wait for all D2H copies to land, then free GPU storage. - self._offload_stream.synchronize() - for t in self._managed: - storage = self._local(t).untyped_storage() - if storage.size() != 0: - storage.resize_(0) - else: - raise RuntimeError( - f"Tensor storage is already freed (size=0) before offload. " - f"This indicates a double-free or external interference. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - - if not self._logged: - logger.info( - "[CPUOffload] Offloaded %.2f MB (GPU → CPU)", - offloaded_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def reload(self): - """Per-tensor H2D from CPU flat buffer on the default stream. - - Runs on the current (default) CUDA stream to avoid stream - interaction issues with the parallel Muon pipeline. Since - pinned CPU memory is the source, the copies overlap with - GPU idle time between steps. - """ - if not self._managed or not self._initialized: - return - - reloaded_bytes = 0 - - # Re-allocate all GPU storages first. - for t in self._managed: - local = self._local(t) - storage = local.untyped_storage() - if storage.size() != 0: - raise RuntimeError( - f"Storage should have been freed (size=0) before reload, " - f"but got size={storage.size()}. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - storage.resize_(self._storage_nbytes[id(t)]) - - # Per-tensor H2D copies from CPU flat buffer slices. - # non_blocking=True with pinned source allows DMA overlap. - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off : off + n], non_blocking=True) - - reloaded_bytes += grp["total"] * cpu_flat.element_size() - - if not self._logged: - logger.info( - "[CPUOffload] Reloaded %.2f MB (CPU → GPU)", reloaded_bytes / (1024**2) - ) diff --git a/build/torch28-cxx11-cu128-x86_64-linux/distributed/utils.py b/build/torch28-cxx11-cu128-x86_64-linux/distributed/utils.py deleted file mode 100644 index 683096bcd50489e7a09959753fdbc29ba0f3b6a6..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/distributed/utils.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import (Placement, Shard, - _StridedShard) - - -def _is_shard(placement: Placement) -> bool: - """Check if a placement is a shard type (Shard or _StridedShard). - - In PyTorch 2.10+, _StridedShard no longer inherits from Shard, so - ``placement.is_shard()`` returns False for _StridedShard. This helper - handles both old and new hierarchies. - """ - return isinstance(placement, (Shard, _StridedShard)) - - -def get_slices_of_dtensor( - target: DTensor | torch.Tensor, - local_rank: int, - shard_mesh: DeviceMesh, - shard_placements: tuple[Placement], -) -> tuple[slice | torch.Tensor, ...]: - """ - Get per-dimension indices for a given rank's shard of the target tensor. - - Uses ``Shard.local_shard_size_and_offset`` and - ``_StridedShard.local_shard_size_and_offset`` for correct handling of - both contiguous and strided (non-contiguous) sharding. - - Args: - target (DTensor | torch.Tensor): The target tensor (for its shape). - local_rank (int): The local rank within the shard group. - shard_mesh (DeviceMesh): The shard mesh (only shard dimensions). - shard_placements (tuple[Placement]): The shard placements. - - Returns: - A tuple of indices (one per tensor dim). Each element is either: - - A ``slice`` (for contiguous or unsharded dims) - - A 1-D ``torch.LongTensor`` of indices (for strided sharding) - """ - - # find the global rank of the local rank in the shard mesh - rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank] - - rank_coords = (shard_mesh.mesh == rank).nonzero() - - assert len(rank_coords) == 1 - rank_coords = tuple(rank_coords[0].tolist()) - - assert len(rank_coords) == len(shard_placements) - - # Track per-shard-dim indices. - # None means "not yet sharded on this dim". - dim_indices: dict[int, torch.Tensor] = {} - - # Caution: Assuming replicate-to-shard of the shard mesh goes with - # left-to-right sharding. This is ensured by the sorting logic of - # construct_shard_mesh function. - for mesh_dim_idx, (rank_coord, placement) in enumerate( - zip(rank_coords, shard_placements)): - assert _is_shard(placement) - - num_chunks = shard_mesh.mesh.shape[mesh_dim_idx] - shard_dim = placement.dim - - # Current effective size on this dim (may already be sub-sharded) - if shard_dim in dim_indices: - curr_size = len(dim_indices[shard_dim]) - else: - curr_size = target.size()[shard_dim] - - # Compute indices for this level of sharding - if isinstance(placement, _StridedShard): - _shard_size, offsets = _StridedShard.local_shard_size_and_offset( - placement, - curr_size, - num_chunks, - rank_coord, - return_first_offset=False) - new_indices = torch.tensor(offsets, dtype=torch.long) - else: - shard_size, offset = Shard.local_shard_size_and_offset( - curr_size, num_chunks, rank_coord) - new_indices = torch.arange(offset, - offset + shard_size, - dtype=torch.long) - - # Compose with previous indices on this dim - if shard_dim in dim_indices: - dim_indices[shard_dim] = dim_indices[shard_dim][new_indices] - else: - dim_indices[shard_dim] = new_indices - - # Build result tuple - result: list[slice | torch.Tensor] = [] - for d in range(len(target.size())): - if d not in dim_indices: - result.append(slice(None)) - else: - indices = dim_indices[d] - # Convert contiguous indices to slice for efficiency - if len(indices) > 0: - start = indices[0].item() - expected = torch.arange(start, - start + len(indices), - dtype=torch.long) - if torch.equal(indices, expected): - result.append(slice(start, start + len(indices))) - else: - result.append(indices) - else: - result.append(slice(0, 0)) - - return tuple(result) - - -_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, - ProcessGroup]] = dict() - - -def construct_shard_mesh( - placements: tuple[Placement], - mesh: DeviceMesh, -) -> tuple[DeviceMesh, ProcessGroup, tuple[Placement, ...]]: - """Construct shard sub-mesh and ProcessGroup for all-to-all communication. - - Given a DTensor's placements and device mesh, extracts the "shard group" - — the set of ranks that together hold all shards of the same replica — - and creates a ProcessGroup for all-to-all among them. - - Steps: - 1. Sort placements: Replicate first, then Shard by (dim, granularity). - 2. Permute the mesh tensor to match the sorted order. - 3. Collapse Replicate dims → list of shard sub-meshes (one per replica). - 4. Create/retrieve a cached ProcessGroup for the current rank's sub-mesh. - - Example — 8 GPUs, mesh shape (2, 2, 2), - placements ``[Shard(0), Replicate, _StridedShard(0)]``:: - - Step 1 — Sort: [Replicate, _StridedShard(0), Shard(0)] - Permutation: [1, 2, 0] - - Step 2 — Permute mesh dims by [1, 2, 0]: - Original: Permuted: - [[[0,1],[2,3]], [[[0,2],[1,3]], - [[4,5],[6,7]]] [[4,6],[5,7]]] - - Step 3 — Unbind replicate dim (dim 0), giving 2 shard sub-meshes: - sub-mesh 0 = [[0,2],[1,3]] (replica group 0) - sub-mesh 1 = [[4,6],[5,7]] (replica group 1) - shard_placements = (_StridedShard(0), Shard(0)) - - Step 4 — Rank 0 → ProcessGroup([0,1,4,5]) - Rank 2 → ProcessGroup([2,3,6,7]) - - Returns: - ``(shard_mesh, process_group, shard_placements)`` - """ - my_rank = dist.get_rank() - assert mesh.mesh.device.type == 'cpu' - - # -- Fast path: 1D all-shard mesh → reuse existing PG. ---------------- - # Reuses the mesh's existing ProcessGroup directly, avoiding the - # overhead of dist.new_group(). The standard path below also handles - # subset calls safely via use_local_synchronization=True, but this - # fast path is still beneficial for the common 1D shard case. - if mesh.ndim == 1 and len(placements) == 1 and _is_shard(placements[0]): - key = (*mesh.mesh.shape, *mesh.mesh.flatten().tolist()) - if key not in _ranks_to_dist_cache: - _ranks_to_dist_cache[key] = (mesh, mesh.get_group()) - return (*_ranks_to_dist_cache[key], tuple(placements)) - - mesh_tensor = mesh.mesh.clone() - - # -- Step 1: Sort placements (Replicate first, then Shard by dim). ------ - # _StridedShard comes BEFORE regular Shard on the same dim so that - # get_slices_of_dtensor applies the outer sharding first, matching - # DTensor's left-to-right (outer-to-inner) composition order. - def _sort_key(item): - index, placement = item - assert not placement.is_partial(), "Partial placement not supported" - if placement.is_replicate(): - return (-1, 0, index) - assert _is_shard(placement), f"Unsupported: {type(placement)}" - split = (-1 / placement.split_factor if isinstance( - placement, _StridedShard) else 0) - return (placement.dim, split, index) - - indexed = sorted(enumerate(placements), key=_sort_key) - perm, sorted_placements = zip(*indexed) - - # -- Step 2: Permute mesh to match sorted placement order. -------------- - sorted_mesh = mesh_tensor.permute(perm) - - # -- Step 3: Collapse replicate dims → list of shard sub-meshes. -------- - # E.g. mesh (2, 3, 4, 4) with [R, R, S(0), S(1)] → 6 sub-meshes of (4, 4) - num_rep = sum(1 for p in sorted_placements if p.is_replicate()) - if num_rep > 0: - if num_rep > 1: - sorted_mesh = sorted_mesh.flatten(0, num_rep - 1) - shard_meshes = list(torch.unbind(sorted_mesh, dim=0)) - else: - shard_meshes = [sorted_mesh] - shard_placements = sorted_placements[num_rep:] - assert len(shard_placements) == len(set(shard_placements)) - - # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. -- - # Each rank only creates the group it belongs to, using - # use_local_synchronization=True so that only group members need to - # coordinate. This avoids deadlocks when different PP stages call - # construct_shard_mesh for different parameters. - def _cache_key(t: torch.Tensor) -> tuple: - return (*t.shape, *t.flatten().tolist()) - - my_key = None - for sm in shard_meshes: - if (my_rank == sm).any().item(): - key = _cache_key(sm) - assert my_key is None, "Rank appears in multiple shard groups" - my_key = key - if key not in _ranks_to_dist_cache: - pg = dist.new_group(sm.flatten().tolist(), - use_local_synchronization=True) - _ranks_to_dist_cache[key] = ( - DeviceMesh(device_type="cuda", mesh=sm), - pg, - ) - - return (*_ranks_to_dist_cache[my_key], shard_placements) diff --git a/build/torch28-cxx11-cu128-x86_64-linux/matmul_transpose_triton.py b/build/torch28-cxx11-cu128-x86_64-linux/matmul_transpose_triton.py deleted file mode 100644 index 792de23d82c3fb45fe33d397ab9b76a0787259d0..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/matmul_transpose_triton.py +++ /dev/null @@ -1,122 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], - restore_value=['y'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -@torch.library.custom_op("muon::matmul_transpose_assign", - mutates_args=("d_out", )) -def matmul_transpose_assign(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """Compute d_out = d_in @ d_in.T using an optimized Triton kernel.""" - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -@matmul_transpose_assign.register_fake -def _(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """FakeTensor impl: d_out is already allocated, mutation is declared.""" - pass diff --git a/build/torch28-cxx11-cu128-x86_64-linux/metadata.json b/build/torch28-cxx11-cu128-x86_64-linux/metadata.json deleted file mode 100644 index c55a35717622f1dd5c8ba376ea3a814cbcc10d78..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/metadata.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch28-cxx11-cu128-x86_64-linux/muon.py b/build/torch28-cxx11-cu128-x86_64-linux/muon.py deleted file mode 100644 index 14c0e22471fa6d47a51ed95e0e0c341dc18d5194..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/muon.py +++ /dev/null @@ -1,1068 +0,0 @@ -import logging -import types -from collections import defaultdict -from typing import Any - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Replicate, Shard -from torch.profiler import record_function - -from .adamw import _placement_cache, _tensor_cache, step_adamw -from .async_utils import run_pipeline -from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho, - get_default_muon_param_groups, is_expert_param, update_p) -from .cpu_offload import CPUOffloadPool -from .distributed.utils import (_is_shard, construct_shard_mesh, - get_slices_of_dtensor) -from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO, - _zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5_batched) -from .pipeline import muon_chunk_pipeline, prelaunch_first_gather -from .qk_clip import compute_scales, get_qk_clip_info, qk_clip - -logger = logging.getLogger(__name__) - - -def _expand_expert_params(names, params, expert_keys): - """Expand expert params by splitting on dim 0 (expert dimension). - - Params whose name matches any key in ``expert_keys`` are treated as - expert-parallel tensors. Their outermost dimension is the expert - dimension: an ``(E, out, in)`` tensor becomes ``E`` separate 2D - ``nn.Parameter`` views so that in-place updates propagate back to - the original storage. - - Non-expert params with ``ndim > 2`` trigger an ``AssertionError`` — - if they are expert params, their key must be added to ``expert_keys``. - - The grad must already be set on each expert param (e.g. after momentum). - - For DTensor expert params, placements that shard on dim 0 (expert dim) - are consumed by the split. Non-dim-0 shard placements (e.g. TP) are - preserved: each 2D slice is wrapped as a DTensor on the corresponding - submesh so the parallel pipeline handles the TP communication. - """ - expanded_names = [] - expanded_params = [] - - for n, p in zip(names, params): - is_expert = is_expert_param(n, expert_keys) - is_dtensor = isinstance(p.data, DTensor) - - if is_expert: - if is_dtensor: - logger.debug( - "[expand_expert] %s: expert DTensor, shape=%s, " - "placements=%s, mesh=%s, local_shape=%s", n, p.shape, - p.placements, p.device_mesh.mesh_dim_names, - p.to_local().shape) - else: - logger.debug( - "[expand_expert] %s: expert plain tensor, shape=%s", n, - p.data.shape) - - if not is_expert: - assert p.data.ndim <= 2, ( - f"Param {n} has ndim={p.data.ndim} but does not match " - f"expert_keys={expert_keys}. If this is an expert param, " - f"add its key to expert_keys.") - expanded_names.append(n) - expanded_params.append(p) - continue - - g = p.grad - assert g is not None, ( - f"Expert param {n} must have grad set before expansion") - - tp_mesh = None - tp_placements_2d = None - - if is_dtensor: - local_data = p.to_local() - local_grad = g.to_local() if isinstance(g, DTensor) else g - - # Find non-dim-0 shard placements (e.g. TP sharding). - # After splitting on dim 0, Shard(k) becomes Shard(k-1). - tp_dim_indices = [] - tp_placements_2d = [] - for i, pl in enumerate(p.placements): - if _is_shard(pl) and pl.dim != 0: - tp_dim_indices.append(i) - tp_placements_2d.append(Shard(pl.dim - 1)) - - if tp_dim_indices: - tp_dim_names = tuple(p.device_mesh.mesh_dim_names[i] - for i in tp_dim_indices) - if len(tp_dim_names) == 1: - tp_mesh = p.device_mesh[tp_dim_names[0]] - else: - tp_mesh = p.device_mesh[tp_dim_names] - else: - local_data = p.data - local_grad = g - - # Expand: split dim 0, reshape each slice to 2D. - num_local_experts = local_data.shape[0] - for i in range(num_local_experts): - slice_data = local_data[i] - slice_grad = local_grad[i] - - if tp_mesh is not None: - # Wrap as DTensor on TP submesh so the pipeline handles - # TP communication (gather/scatter across TP ranks). - dt_data = DTensor.from_local(slice_data, - device_mesh=tp_mesh, - placements=tp_placements_2d) - dt_grad = DTensor.from_local(slice_grad, - device_mesh=tp_mesh, - placements=tp_placements_2d) - expert_param = torch.nn.Parameter(dt_data, requires_grad=False) - expert_param.grad = dt_grad - else: - expert_param = torch.nn.Parameter(slice_data, - requires_grad=False) - expert_param.grad = slice_grad - - expanded_names.append(f"{n}[{i}]") - expanded_params.append(expert_param) - - p.grad = None # allow expert grad storage to be freed after pipeline - - return expanded_names, expanded_params - - -class Muon(torch.optim.Optimizer): - """ - Muon - MomentUm Orthogonalized by Newton-schulz - - Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- - processing step, in which each 2D parameter's update is replaced with the nearest orthogonal - matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has - the advantage that it can be stably run in bfloat16 on the GPU. - - Some warnings: - - We believe this optimizer is unlikely to work well for training with small batch size. - - We believe it may not work well for finetuning pretrained models, but we haven't tested this. - - Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. - lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) - momentum: The momentum used by the internal SGD. (0.95 is a good default) - nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) - ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. - Parameters that are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW instead. - adamw_lr: The learning rate for the internal AdamW. - adamw_betas: The betas for the internal AdamW. - adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - warmup_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher warmup_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. - chunk_size : Batch size of parameters to process in each - all2all gather/compute/scatter step. - Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified. - use_distributed_muon: Use distributed muon by Liu et al. (2024). - For testing purpose only. - expert_keys: List of strings to identify expert-parallel parameters. - If any key appears in a parameter's name, its outermost - dimension is treated as the expert dimension and expanded - into per-expert 2D params for Muon. For example, - ``expert_keys=["experts"]`` matches any param whose name - contains "experts". 3D+ params not matched by any key - will raise an error. - """ - - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config=None, - warmup_step=5, - chunk_size=-1, - use_distributed_muon=False, - expert_keys=None): - defaults = dict( - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - nesterov=nesterov, - ns_steps=ns_steps, - adamw_betas=adamw_betas, - adamw_eps=adamw_eps, - none_grad=none_grad, - use_muon=True, - ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) - super().__init__(params, defaults) - - self.debug = debug - self.clip_config = clip_config if clip_config is not None else { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100, - } - self.warmup_step = warmup_step - self.chunk_size = chunk_size - self.use_distributed_muon = use_distributed_muon - self.expert_keys = expert_keys - self.cpu_offload = False - self._cpu_offload_pool: CPUOffloadPool | None = None - self._offload_initialized = False - self._parallel_cache: dict[tuple[str, ...], dict] = {} - self._expert_expand_cache: dict[tuple[int, ...], dict] = {} - - def _calc_flops(self, G, steps): - assert len(G.shape) == 2 - M, N = G.shape - if M > N: - M, N = N, M - - return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) - - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - p.placements, p.device_mesh) - - return shard_mesh, shard_pg, shard_placements - - def init_state_and_assign_params(self, names, params, group, qk_logits): - param_to_state = {} - param_to_flops = {} - - total_flops = 0 - for p in params: - g = p.grad - if g is None: - continue - assert g.ndim == 2, "Muon only supports 2D parameters." - - flops = self._calc_flops(g, group["ns_steps"]) - param_to_flops[id(p)] = flops - total_flops += flops - - if self.debug: - logger.debug("Total TFLOPs for Muon: %.2f TFLOPs", - total_flops / 1e12) - - paired = list(zip(names, params)) - - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) - - round_robin = 0 - mesh = ordered_params[0].device_mesh - placements = ordered_params[0].placements - - shard_mesh, shard_pg, shard_placements = self.get_shard_mesh( - ordered_params[0]) - shard_mesh_flattened = shard_mesh.mesh.flatten() - num_ranks = dist.get_world_size(group=shard_pg) - - for n, p in zip(ordered_names, ordered_params): - if mesh != p.device_mesh: - raise ValueError("All parameters must be on the same mesh.") - if placements != p.placements: - raise ValueError("All parameters must have same placements.") - - worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks - round_robin = (round_robin + 1) % len(shard_mesh_flattened) - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - # Precompute per-rank indices and numels for all-to-all. - rank_indices: dict[int, tuple] = {} - rank_numels: dict[int, int] = {} - for r in range(num_ranks): - indices = get_slices_of_dtensor(p, r, shard_mesh, - shard_placements) - rank_indices[r] = indices - numel = 1 - for idx, dim_size in zip(indices, p.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) - rank_numels[r] = numel - - param_to_state[id(p)] = _muon_state( - worker_rank=worker_rank, - process_group=shard_pg, - rank_indices=rank_indices, - rank_numels=rank_numels, - name=n, - qk_clip_state=qk_clip_state, - ) - - return param_to_state, ordered_params - - def base(self, names, params, group, lr, weight_decay, qk_logits): - # Momentum is already applied by _step_muon before this method. - for n, p in zip(names, params): - g = p.grad - if g is None: - continue - - u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) - - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - def distributed_muon( - self, - names: list[str], - params: list[torch.nn.Parameter], - group: dict[str, Any], - lr: float, - weight_decay: float, - qk_logits: list[torch.Tensor | DTensor] | None, - ): - """Batched Distributed Muon — for testing/correctness verification only. - - Uses all-gather to reconstruct full tensors, computes Newton-Schulz on - the full grad, then slices back to local shards. This is simpler but - slower than the parallel pipeline (all2all) path, so it serves as a - reference implementation for verifying correctness. - """ - with record_function("distributed_muon"): - # Momentum is already applied by _step_muon before this method. - ns_steps = group["ns_steps"] - - # Separate plain tensors (no communication) from DTensors. - plain_names, plain_params = [], [] - dtensor_names, dtensor_params = [], [] - for n, p in zip(names, params): - if p.grad is None: - continue - if isinstance(p.data, DTensor): - dtensor_names.append(n) - dtensor_params.append(p) - else: - plain_names.append(n) - plain_params.append(p) - - # Process plain tensors per-param (no communication). - for n, p in zip(plain_names, plain_params): - u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE), - steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, - qk_logits) - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - if not dtensor_params: - return - - # Group DTensors by (placements, mesh) for batched all-gather. - placement_groups: dict[tuple, - tuple[list, - list]] = defaultdict(lambda: ([], [])) - for n, p in zip(dtensor_names, dtensor_params): - key = (p.placements, p.device_mesh) - placement_groups[key][0].append(n) - placement_groups[key][1].append(p) - - logger.info( - "distributed_muon: %d placement groups, %d total dtensors", - len(placement_groups), len(dtensor_params)) - - for (placements, mesh), (grp_names, - grp_params) in placement_groups.items(): - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - placements, mesh) - rank = dist.get_rank(shard_pg) - world_size = dist.get_world_size(shard_pg) - - logger.info(" group: %d params, placements=%s, world_size=%d", - len(grp_params), placements, world_size) - - # Separate params that can be batched (all shard dims evenly - # divisible) from those needing per-param full_tensor - # (e.g. MoE gate weights with fewer rows than shard ranks). - # all_gather_into_tensor requires equal buffer sizes across - # ranks, so uneven splits must use DTensor full_tensor(). - batch_names, batch_params = [], [] - single_names, single_params = [], [] - for n, p in zip(grp_names, grp_params): - even = all(p.shape[pl.dim] % - shard_mesh.mesh.shape[dim_idx] == 0 - for dim_idx, pl in enumerate(shard_placements)) - if even: - batch_names.append(n) - batch_params.append(p) - else: - single_names.append(n) - single_params.append(p) - - # Process uneven-split params per-param via full_tensor(). - for n, p in zip(single_names, single_params): - with record_function("distributed_muon::newton_schulz"): - g_full = p.grad.full_tensor().to(COMM_DTYPE) - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - if not batch_params: - continue - - logger.info(" batched=%d, single=%d", len(batch_params), - len(single_params)) - - # Concat all local grad shards into a single flat buffer. - with record_function("distributed_muon::gather"): - grad_locals = [ - p.grad.to_local().to(COMM_DTYPE).flatten() - for p in batch_params - ] - numels = [g.numel() for g in grad_locals] - grad_concat = torch.cat(grad_locals) - del grad_locals - - # Single all-gather (replaces N separate full_tensor). - grad_gathered = torch.empty( - grad_concat.numel() * world_size, - dtype=COMM_DTYPE, - device="cuda", - ) - dist.all_gather_into_tensor(grad_gathered, - grad_concat, - group=shard_pg) - - total_numel = grad_concat.numel() - del grad_concat - - # Precompute per-param offsets within the concat buffer. - offsets = [] - off = 0 - for ne in numels: - offsets.append(off) - off += ne - - # Per-param: reconstruct full grad → NS → local update. - for i, (n, p) in enumerate(zip(batch_names, batch_params)): - with record_function("distributed_muon::newton_schulz"): - g_full = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - for r in range(world_size): - r_start = r * total_numel + offsets[i] - shard = grad_gathered[r_start:r_start + numels[i]] - indices = get_slices_of_dtensor( - p, r, shard_mesh, shard_placements) - g_full[indices] = shard.reshape( - g_full[indices].shape) - - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - def _setup_parallel(self, names, params, group, qk_logits): - """Compute (or retrieve cached) parallel pipeline metadata. - - Returns: - (ordered_params, param_to_state, rank, chunk_size) - """ - cache_key = tuple(names) - - if cache_key not in self._parallel_cache: - # First call: compute metadata and populate cache. - param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - shard_pg = param_to_state[id(ordered_params[0])].process_group - rank = dist.get_rank(group=shard_pg) - - if self.chunk_size == -1: - shard_ranks = dist.get_world_size(shard_pg) - chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO - elif self.chunk_size > 0: - chunk_size = self.chunk_size - else: - raise ValueError( - "chunk_size must be -1 or a positive integer.") - - ordered_names = [ - param_to_state[id(p)].name for p in ordered_params - ] - name_to_state = { - param_to_state[id(p)].name: param_to_state[id(p)] - for p in ordered_params - } - self._parallel_cache[cache_key] = { - 'ordered_names': ordered_names, - 'name_to_state': name_to_state, - 'rank': rank, - 'chunk_size': chunk_size, - } - else: - # Cached path: rebuild param_to_state with current id(p) keys. - cache = self._parallel_cache[cache_key] - rank = cache['rank'] - chunk_size = cache['chunk_size'] - - name_to_param = dict(zip(names, params)) - ordered_params = [name_to_param[n] for n in cache['ordered_names']] - - param_to_state = {} - for p, n in zip(ordered_params, cache['ordered_names']): - cached_state = cache['name_to_state'][n] - param_to_state[id(p)] = _muon_state( - worker_rank=cached_state.worker_rank, - process_group=cached_state.process_group, - rank_indices=cached_state.rank_indices, - rank_numels=cached_state.rank_numels, - name=n, - qk_clip_state=get_qk_clip_info(self.clip_config, n, - qk_logits), - ) - - return ordered_params, param_to_state, rank, chunk_size - - def parallel(self, - names, - params, - group, - lr, - weight_decay, - qk_logits, - prelaunch_gather=None): - """ - Perform a parallel optimization step using Muon. - - Parameters are chunked and each chunk is processed by a - :func:`muon_chunk_pipeline` generator. :func:`run_pipeline` - interleaves multiple chunks so that communication and computation - overlap across chunks (the same overlap previously achieved by the - warmup + main-loop index scheduling). - - If ``prelaunch_gather`` is provided, it is passed to the first - chunk's generator to skip re-launching the already in-flight - A2A gather. - """ - - # Momentum is already applied by _step_muon before this method. - - ordered_params, param_to_state, rank, chunk_size = ( - self._setup_parallel(names, params, group, qk_logits)) - - def pipelines(): - first = True - for start in range(0, len(ordered_params), chunk_size): - chunk = ordered_params[start:start + chunk_size] - if chunk: - kwargs = dict( - params=chunk, - param_to_state=param_to_state, - rank=rank, - ns_steps=group["ns_steps"], - lr=lr, - weight_decay=weight_decay, - none_grad=group["none_grad"], - ) - if first and prelaunch_gather is not None: - kwargs['prelaunch_gather'] = prelaunch_gather - first = False - yield muon_chunk_pipeline(**kwargs) - - with record_function("muon::pipeline"): - run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1) - - def _step_muon(self, group, qk_logits=None): - params = group["params"] - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - # Apply momentum to all params before routing/expansion. - # Batched using _foreach_* ops (compiled, fullgraph=True). - with record_function("muon::momentum"): - active_params = [p for p in params if p.grad is not None] - if active_params: - # Ensure momentum buffers exist (avoid zeros_like when already present). - for p in active_params: - if "momentum_buffer" not in self.state[p]: - self.state[p]["momentum_buffer"] = torch.zeros_like( - p.grad) - - # Extract local tensors for compiled batch function. - local_grads = [ - p.grad._local_tensor - if isinstance(p.grad, DTensor) else p.grad - for p in active_params - ] - local_bufs = [ - self.state[p]["momentum_buffer"]._local_tensor - if isinstance(self.state[p]["momentum_buffer"], DTensor) - else self.state[p]["momentum_buffer"] - for p in active_params - ] - - # Wrap momentum as tensor for torch.compile. - batch_pre_ortho(local_grads, local_bufs, - torch.tensor(momentum), group["nesterov"]) - - # For non-nesterov, the result is the momentum buffer. - if not group["nesterov"]: - for p in active_params: - p.grad = self.state[p]["momentum_buffer"] - - # Identify batched experts for deferred NS. - # Detection is cheap (condition checks only); actual NS compute is - # deferred so it can overlap with the first chunk's A2A gather. - deferred_expert_work = [] - if self.expert_keys: - batched_expert_indices = [] - for i, (n, p) in enumerate(zip(names, params)): - if not (is_expert_param(n, self.expert_keys) - and p.grad is not None): - continue - # Eligible: plain tensor, or DTensor with no non-dim-0 shards. - if isinstance(p.data, DTensor): - has_tp = any( - _is_shard(pl) and pl.dim != 0 for pl in p.placements) - if has_tp: - continue - batched_expert_indices.append(i) - - if batched_expert_indices: - # Save refs for deferred NS; free grads from param list. - for i in batched_expert_indices: - p = params[i] - g = p.grad - local_g = (g._local_tensor - if isinstance(g, DTensor) else g) - local_data = (p.data._local_tensor if isinstance( - p.data, DTensor) else p.data) - deferred_expert_work.append((local_data, local_g)) - p.grad = None - - # Remove batched experts from lists before expansion. - keep = sorted( - set(range(len(params))) - set(batched_expert_indices)) - names = [names[i] for i in keep] - params = [params[i] for i in keep] - - def _run_deferred_expert_ns(): - """Execute deferred batched expert NS.""" - if not deferred_expert_work: - return - with record_function("muon::batched_expert_ns"): - ns_steps = group["ns_steps"] - for local_data, local_g in deferred_expert_work: - u = zeropower_via_newtonschulz5_batched( - local_g.to(COMM_DTYPE), steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:]) - local_data.mul_(1 - lr * weight_decay) - local_data.add_(u, alpha=-adjusted_lr) - - # Expand expert params by splitting on dim 0. - logger.debug("[_step_muon] before expand: %d params, expert_keys=%s", - len(params), self.expert_keys) - if self.expert_keys: - cache_key = tuple(id(p) for p in params) - cache = self._expert_expand_cache.get(cache_key) - - if cache is None: - # Cold path: full expansion + build cache metadata. - exp_names, exp_params = _expand_expert_params( - names, params, self.expert_keys) - - # Build per-expert-group info for hot-path grad updates. - grad_info = [] - exp_idx = 0 - for orig_idx, (n, p) in enumerate(zip(names, params)): - if not is_expert_param(n, self.expert_keys): - exp_idx += 1 - continue - - is_dt = isinstance(p.data, DTensor) - num_experts = (p.to_local() if is_dt else p.data).shape[0] - - # Detect TP mesh from the first expanded expert param. - tp_mesh = None - tp_pls = None - sample = exp_params[exp_idx] - if isinstance(sample.data, DTensor): - tp_mesh = sample.data.device_mesh - tp_pls = list(sample.data.placements) - - grad_info.append((orig_idx, num_experts, exp_idx, is_dt, - tp_mesh, tp_pls)) - exp_idx += num_experts - - self._expert_expand_cache[cache_key] = { - 'names': exp_names, - 'params': exp_params, - 'grad_info': grad_info, - } - names, params = exp_names, exp_params - else: - # Hot path: reuse cached params, only update expert grads. - for (orig_idx, num_experts, exp_start, is_dt, tp_mesh, - tp_pls) in cache['grad_info']: - p = params[orig_idx] - g = p.grad - local_grad = (g.to_local() - if is_dt and isinstance(g, DTensor) else g) - for i in range(num_experts): - expert_p = cache['params'][exp_start + i] - sg = local_grad[i] - if tp_mesh is not None: - expert_p.grad = DTensor.from_local( - sg, device_mesh=tp_mesh, placements=tp_pls) - else: - expert_p.grad = sg - p.grad = None - - names = cache['names'] - params = cache['params'] - else: - names, params = _expand_expert_params(names, params, - self.expert_keys) - logger.debug("[_step_muon] after expand: %d params", len(params)) - - param_dtensors = [] - name_dtensors = [] - - param_tensors = [] - name_tensors = [] - - # distributed_muon is a reference implementation for testing only. - # The parallel pipeline (all2all) path below is the production path. - if self.use_distributed_muon: - _run_deferred_expert_ns() - self.distributed_muon(names=names, - params=params, - group=group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits) - return - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - logger.debug( - "[route] %s → base (DTensor all-Replicate), " - "shape=%s, placements=%s", n, p.shape, p.placements) - param_tensors.append(p) - name_tensors.append(n) - else: - logger.debug( - "[route] %s → parallel (DTensor), shape=%s, " - "placements=%s, mesh=%s", n, p.shape, p.placements, - p.device_mesh.mesh_dim_names) - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - logger.debug("[route] %s → base (plain tensor), shape=%s", n, - p.data.shape) - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError(f"Unsupported parameter type: {type(p.data)}") - - logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, " - f"{len(param_tensors)} Tensors → base") - - def group_dtensors(dtensors, names): - # To support different placements, we group parameters by placements - # and run parallel Muon on each group. - - placement_to_params = defaultdict(lambda: ([], [])) - - assert len(dtensors) == len(names) - for p, n in zip(dtensors, names): - placement_to_params[tuple([p.placements, - p.device_mesh])][0].append(n) - placement_to_params[tuple([p.placements, - p.device_mesh])][1].append(p) - return placement_to_params - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - dtensor_group = group_dtensors(param_dtensors, name_dtensors) - - # Pre-launch the first chunk's A2A gather so that the NCCL - # communication overlaps with the (deferred) batched expert NS - # compute on the default CUDA stream. - prelaunch = None - if deferred_expert_work: - first_names, first_params = next(iter(dtensor_group.values())) - ordered, pts, rnk, csz = self._setup_parallel( - first_names, first_params, group, qk_logits) - first_chunk = ordered[:csz] - if first_chunk: - prelaunch = prelaunch_first_gather(first_chunk, pts, rnk, - group["none_grad"]) - - _run_deferred_expert_ns() - - first_group = True - for _, (names, params) in dtensor_group.items(): - pg = prelaunch if first_group else None - first_group = False - self.parallel( - names, - params, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - prelaunch_gather=pg, - ) - else: - _run_deferred_expert_ns() - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - ) - - def _register_states_for_offload(self): - """Register all optimizer state tensors with the CPU offload pool. - - Called once after the first step when states have been lazily created. - Offloads all param states (momentum buffers for Muon, moment1/moment2 - for AdamW) to free GPU memory between steps. - """ - pool = self._cpu_offload_pool - tracked = 0 - for group in self.param_groups: - for p in group["params"]: - if p not in self.state: - continue - state = self.state[p] - if group.get("use_muon", False): - if "momentum_buffer" in state: - pool.track(state["momentum_buffer"]) - tracked += 1 - else: - if "moment1" in state: - pool.track(state["moment1"]) - if "moment2" in state: - pool.track(state["moment2"]) - tracked += 1 - logger.info("[CPUOffload] Registered %d param states for offload", - tracked) - - @torch.no_grad - def step(self, closure=None, qk_logits=None): - """Perform a single optimization step. - - Args: - closure (Callable, optional): A closure that reevaluates the model - and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). - """ - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # H2D: reload optimizer states from CPU before computation. - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload() - - logger.debug("[Muon.step] expert_keys=%s, %d param groups", - self.expert_keys, len(self.param_groups)) - - for i, group in enumerate(self.param_groups): - if group["use_muon"]: - logger.debug("[Muon.step] group %d: use_muon=True, %d params", - i, len(group["params"])) - self._step_muon(group, qk_logits=qk_logits) - else: - logger.debug( - "[Muon.step] group %d: use_muon=False (AdamW), %d params", - i, len(group["params"])) - step_adamw(self.state, group) - - # D2H: offload optimizer states to CPU after computation. - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - return loss - - # ------------------------------------------------------------------ - # CPU offload public helpers - # ------------------------------------------------------------------ - - def turn_on_cpu_offload(self): - """Enable CPU offload for optimizer states.""" - if self.cpu_offload: - return - logger.info("[Muon] turn_on_cpu_offload") - self.cpu_offload = True - if not self.state: - return - self._cpu_offload_pool = CPUOffloadPool() - self._offload_initialized = False - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_off_cpu_offload(self): - """Disable CPU offload and keep optimizer states resident on GPU.""" - if not self.cpu_offload: - return - logger.info("[Muon] turn_off_cpu_offload") - if self._offload_initialized: - self._cpu_offload_pool.reload() - torch.cuda.current_stream().synchronize() - self._cpu_offload_pool = None - self._offload_initialized = False - self.cpu_offload = False - - # ------------------------------------------------------------------ - # Checkpoint support for cpu_offload - # ------------------------------------------------------------------ - - def state_dict(self) -> dict: - if self.cpu_offload: - raise RuntimeError( - "Muon.state_dict() requires turn_off_cpu_offload() before checkpoint save." - ) - return super().state_dict() - - def load_state_dict(self, state_dict: dict) -> None: - if self.cpu_offload: - raise RuntimeError( - "Muon.load_state_dict() requires turn_off_cpu_offload() before checkpoint load." - ) - super().load_state_dict(state_dict) - - # Invalidate adamw.py's module-level tensor caches so that - # the next step rebuilds them with the newly loaded state tensors. - _placement_cache.clear() - _tensor_cache.clear() diff --git a/build/torch28-cxx11-cu128-x86_64-linux/newton_schulz.py b/build/torch28-cxx11-cu128-x86_64-linux/newton_schulz.py deleted file mode 100644 index d939264b69a34e7a3fa78859f34dc265a1159d59..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/newton_schulz.py +++ /dev/null @@ -1,240 +0,0 @@ -from itertools import repeat -from math import inf, sqrt - -import numpy as np -import torch - -from .matmul_transpose_triton import matmul_transpose_assign - -COMM_DTYPE = torch.bfloat16 -DEFAULT_CHUNK_SIZE_RATIO = 4 - - -def _optimal_quintic(l, u, max_iter=1000): - """ - Use the simplified Remez algorithm to find the optimal odd quintic approximant - to the constant function x -> 1 over the interval [l, u]. - - Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum - approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the - two interior equioscillation nodes q, r until convergence. Returns the - closed-form equioscillating solution when l ≈ u. - - Raises ValueError if any intermediate value (a, b, c, E, q, r) is non-finite - (NaN or inf). Raises RuntimeError if convergence is not reached within - max_iter iterations. - """ - assert 0 <= l <= u - if 1 - 5e-6 <= l / u: - return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5) - q = (3 * l + u) / 4 - r = (l + 3 * u) / 4 - E = inf - for _ in range(max_iter): - old_E = E - LHS = np.array( - [ - [l, l**3, l**5, 1], - [q, q**3, q**5, -1], - [r, r**3, r**5, 1], - [u, u**3, u**5, -1], - ] - ) - a, b, c, E = np.linalg.solve(LHS, np.ones(4)) - if not np.all(np.isfinite([a, b, c, E])): - raise ValueError( - f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}" - ) - q, r = np.sqrt( - (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / (10 * c) - ) - if not np.all(np.isfinite([q, r])): - raise ValueError(f"_optimal_quintic: non-finite node update q={q}, r={r}") - if abs(old_E - E) <= 1e-15: - break - else: - raise RuntimeError( - f"_optimal_quintic: did not converge after {max_iter} iterations" - ) - return float(a), float(b), float(c) - - -def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0): - """ - Compute the Polar Express coefficient series for `num_iters` quintic iterations. - - Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that - compose to map singular values from [l, 1] toward 1. At each step: - 1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion` - prevents near-zero singular values from stalling by raising the effective - lower bound; if it is active (cushion*u > l), the coefficients are - rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u]. - 2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the - last iteration, providing numerical headroom at the cost of a slightly slower - final convergence step. - 3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1). - - Returns a list of (a, b, c) tuples, one per iteration. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - u = 1 - assert 0 <= l <= u - safety_factor = 1 + safety_factor_eps - coefficients = [] - for iter in range(num_iters): - a, b, c = _optimal_quintic(max(l, cushion * u), u) - if cushion * u > l: - pl = a * l + b * l**3 + c * l**5 - pu = a * u + b * u**3 + c * u**5 - rescaler = 2 / (pl + pu) - a *= rescaler - b *= rescaler - c *= rescaler - if iter < num_iters - 1: - a /= safety_factor - b /= safety_factor**3 - c /= safety_factor**5 - coefficients.append((a, b, c)) - l = a * l + b * l**3 + c * l**5 - u = 2 - l - return coefficients - - -# Precomputed Polar Express coefficients (a, b, c) for 10 quintic Newton-Schulz -# iterations. Each tuple is the minimax-optimal (Remez/equioscillation) odd quintic -# approximant to x->1 over the current singular-value interval, computed once at -# import time and reused across all optimizer steps. -# -# Contrast with the former hardcoded NS coefficients (5 fixed tuples): -# - Former: empirically tuned to maximize slope at zero; did not converge -# singular values to 1, yielding US'V^T with S' ~ Uniform(0.5, 1.5) instead -# of the true polar factor UV^T. -# - Polar Express: analytically optimal per step, adapting to the shrinking -# singular-value interval [l, u] as iterations progress; converges all -# singular values to 1, producing the exact polar factor UV^T. -_coeffs_list = _optimal_composition( - l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02 -) - - -# This code is adapted from: -# KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py) -# NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress) -# matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon) -@torch.no_grad() -def _zeropower_via_newtonschulz5(G, steps): - """ - Compute the polar factor of G via the Polar Express method. - - Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c) - are the Polar Express coefficients from `_coeffs_list`. Each step is the - optimal odd quintic approximant to x -> 1 over the current singular-value - interval, minimizing the maximum approximation error (Remez / minimax criterion). - The composition maps singular values from [l, 1] to near 1, producing the - polar factor (orthogonal factor in the polar decomposition G = UP). - - `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2, - cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE - X = G # no manual typecast - - if G.size(0) > G.size(1): - X = X.T - - X = X / (X.norm() + 1e-7) - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - # Perform the NS iterations - for a, b, c in hs: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(0) > G.size(1): - X = X.T - - return X - - -@torch.no_grad() -def _zeropower_via_newtonschulz5_batched(G, steps): - """Batched polar factor computation for 3D (E, out, in) tensors. - - Same algorithm as ``_zeropower_via_newtonschulz5`` but uses - ``torch.bmm`` / ``torch.baddbmm`` instead of the 2D Triton kernel, - processing all E expert matrices in a single batched call. - """ - assert len(G.shape) == 3 - assert G.dtype == COMM_DTYPE - X = G - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - # Per-expert Frobenius norm. - X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) - - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - for a, b, c in hs: - buf1 = torch.bmm(X, X.transpose(-2, -1)) - buf2 = torch.bmm(buf1, buf1.transpose(-2, -1)) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.baddbmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - return X - - -_ns_per_shape: dict[tuple[int, ...], callable] = {} -_use_compile = True - - -def set_ns_compile(enabled: bool): - """Toggle torch.compile for Newton-Schulz iteration.""" - global _use_compile - _use_compile = enabled - - -def zeropower_via_newtonschulz5(G, steps=5): - if not _use_compile: - return _zeropower_via_newtonschulz5(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() - - -def zeropower_via_newtonschulz5_batched(G, steps=5): - """Compile-cached batched Newton-Schulz for 3D expert tensors.""" - if not _use_compile: - return _zeropower_via_newtonschulz5_batched(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile( - _zeropower_via_newtonschulz5_batched, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() diff --git a/build/torch28-cxx11-cu128-x86_64-linux/optimizer/__init__.py b/build/torch28-cxx11-cu128-x86_64-linux/optimizer/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/optimizer/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch28-cxx11-cu128-x86_64-linux/pipeline.py b/build/torch28-cxx11-cu128-x86_64-linux/pipeline.py deleted file mode 100644 index c0c2d515856182d8d15ad27dd4e4e093b29397d6..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/pipeline.py +++ /dev/null @@ -1,468 +0,0 @@ -import logging -from typing import Generator - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -from .core import _muon_state, adjust_lr_for_muon -from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5 -from .qk_clip import compute_scales - -logger = logging.getLogger(__name__) - -# ====================================================================== -# Stage helpers -# ====================================================================== - - -def _launch_gather( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Allocate gather buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_gather``). - gathered_grads: ``{id(p): empty_tensor}`` for owned params, - ``None`` for non-owned. - recv_counts: Per-source-rank element counts. - """ - # Allocate gathered-grad buffers - gathered_grads: dict[int, torch.Tensor | None] = {} - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - gathered_grads[id(p)] = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - else: - gathered_grads[id(p)] = None - - # Build send buffer – batch grad copies via torch.cat - # (1-2 fused kernels vs N individual narrow().copy_() calls). - send_counts = [0] * num_ranks - for p in params: - state = param_to_state[id(p)] - send_counts[state.worker_rank] += state.rank_numels[rank] - - total_send = sum(send_counts) - if total_send > 0: - # Group grad slices by destination rank in a single pass. - dst_to_grads = [[] for _ in range(num_ranks)] - for p in params: - state = param_to_state[id(p)] - n = state.rank_numels[rank] - if n > 0: - g = p.grad.to_local() - dst_to_grads[state.worker_rank].append(g.reshape(-1)) - - # Flatten in dst order and cat once. - all_slices = [] - for dst in range(num_ranks): - all_slices.extend(dst_to_grads[dst]) - send_buf = torch.cat(all_slices) - if send_buf.dtype != COMM_DTYPE: - send_buf = send_buf.to(COMM_DTYPE) - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += state.rank_numels[src] - recv_counts[src] = total - - recv_buf = torch.empty(sum(recv_counts), dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - logger.debug(f"send_buf size: {send_buf.numel()}, " - f"recv_buf size: {recv_buf.numel()}, " - f"recv_counts: {recv_counts}, " - f"send_counts: {send_counts}, " - f"process_group: {str(process_group)}") - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, gathered_grads, recv_counts - - -def _complete_gather( - recv_buf: torch.Tensor, - recv_counts: list[int], - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - param_to_state: dict[int, _muon_state], - rank: int, -) -> None: - """Reconstruct gathered grads from the recv buffer (in-place).""" - off = 0 - for src in range(len(recv_counts)): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - - indices = state.rank_indices[src] - - shard_view = gathered_grads[id(p)][indices] - n = shard_view.numel() - if n == 0: - continue - - sg = recv_buf.narrow(0, off + inner_off, n) - sg = sg.reshape(shard_view.shape) - gathered_grads[id(p)][indices] = sg - - inner_off += n - assert inner_off == block - off += block - - -def _compute_ns( - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - ns_steps: int, -) -> dict[int, torch.Tensor | None]: - """Run Newton-Schulz orthogonalization on owned parameters. - - Returns: - computed_us: ``{id(p): orthogonalized_update}`` for owned params. - """ - computed_us: dict[int, torch.Tensor | None] = {} - for p in owned_params: - u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps) - gathered_grads[id(p)] = None # free gathered grad - computed_us[id(p)] = u - return computed_us - - -def _launch_scatter( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, - computed_us: dict[int, torch.Tensor | None], -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor], list[int]]: - """Allocate scatter buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_scatter``). - scattered_us: Empty dict, populated by ``_complete_scatter`` with - zero-copy views into ``recv_buf``. - recv_counts: Per-source-rank element counts. - """ - # scattered_us is populated by _complete_scatter with zero-copy views - # into recv_buf, avoiding N empty_like allocations + N copy_ calls. - # Pre-seed entries for params whose local shard is empty (rank_numels == 0) - # so _update_params can iterate all params without KeyError. - scattered_us: dict[int, torch.Tensor] = {} - for p in params: - if param_to_state[id(p)].rank_numels[rank] == 0: - scattered_us[id(p)] = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - # Build send buffer – batch via torch.cat - # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls). - send_counts = [0] * num_ranks - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - for dst_rank in range(num_ranks): - send_counts[dst_rank] += state.rank_numels[dst_rank] - - total_send = sum(send_counts) - if total_send > 0: - # Cache u_full conversions to avoid redundant .to() per dst_rank. - u_fulls = {} - for p in owned_params: - u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous() - - # Collect slices in dst order (matches all-to-all send layout). - all_slices = [] - for dst_rank in range(num_ranks): - for p in owned_params: - state = param_to_state[id(p)] - su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten() - if su.numel() > 0: - all_slices.append(su) - - send_buf = torch.cat(all_slices) if all_slices else torch.empty( - 0, dtype=COMM_DTYPE, device="cuda") - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += state.rank_numels[rank] - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, scattered_us, recv_counts - - -def _complete_scatter( - recv_buf: torch.Tensor, - recv_counts: list[int], - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], -) -> None: - """Populate scattered_us with zero-copy views into recv_buf. - - Instead of pre-allocating tensors and copying, we assign views directly - from ``recv_buf``. This eliminates N ``empty_like`` + N ``copy_`` calls. - The underlying storage of ``recv_buf`` is kept alive through the views - until ``scattered_us`` is cleared after ``_update_params``. - """ - off = 0 - for src in range(len(recv_counts)): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = state.rank_numels[rank] - if n == 0: - continue - - scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - - inner_off += n - - assert inner_off == block - off += block - - -def _update_params( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], - lr: float, - weight_decay: float, -) -> None: - """Apply weight decay, Muon update, and optional QK clipping. - - Uses batched ``_foreach_mul_`` for weight decay and batched - ``_foreach_add_`` for the Muon update, grouping parameters by - adjusted_lr to minimize kernel launches while preserving float32 - precision for the alpha scaling. - """ - if not params: - return - - # Batched weight decay: p *= (1 - lr * wd) — single fused kernel. - p_locals = [p._local_tensor for p in params] - torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay) - - # Group params by adjusted_lr so _foreach_add_ can use a single - # alpha per group (preserves float32 precision for alpha scaling). - lr_groups: dict[float, tuple[list, list]] = {} - for p in params: - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - if adjusted_lr not in lr_groups: - lr_groups[adjusted_lr] = ([], []) - lr_groups[adjusted_lr][0].append(p._local_tensor) - lr_groups[adjusted_lr][1].append(scattered_us[id(p)]) - - for adjusted_lr, (p_group, u_group) in lr_groups.items(): - torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr) - - # QK clipping – applied directly on the local tensor to - # avoid DTensor sharding-propagation issues with _StridedShard. - for p in params: - state = param_to_state[id(p)] - if state.qk_clip_state is None: - continue - scales_full = compute_scales(p, state.qk_clip_state) - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = state.rank_indices[rank][0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - -# ====================================================================== -# Pre-launch helper for overlapping first chunk's gather with other work. -# ====================================================================== - - -@torch.no_grad() -def prelaunch_first_gather( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - none_grad: bool, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Launch the first chunk's A2A gather early for overlap with other compute. - - Call this *before* expensive GPU work (e.g. batched expert NS) so that - the NCCL all-to-all runs concurrently on the NCCL stream while the - default stream executes compute. - - Returns the same 4-tuple that ``_launch_gather`` produces, which should - be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - with record_function("muon::prelaunch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - return work, recv_buf, gathered_grads, recv_counts - - -# ====================================================================== -# Main generator – thin orchestrator that wires stages together. -# ====================================================================== - - -@torch.no_grad() -def muon_chunk_pipeline( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - ns_steps: int, - lr: float, - weight_decay: float, - none_grad: bool, - prelaunch_gather: tuple | None = None, -) -> Generator[None, None, None]: - """Process one chunk of parameters through the full Muon pipeline. - - Stages: gather -> compute (Newton-Schulz) -> scatter -> update. - - Each ``yield`` lets :func:`run_pipeline` interleave other chunks so - that communication and computation overlap across chunks. Async - communication is launched via ``async_op=True`` and completed after - the yield with ``work.wait()``. - - Overlap happens because :func:`run_pipeline` admits one new chunk - per iteration (staggered admission). While chunk *N* does NS - compute on the default CUDA stream, chunk *N+1*'s async all-to-all - runs concurrently on the NCCL stream — no separate ``comm_stream`` - is required. - - If ``prelaunch_gather`` is provided, the gather was already launched - by :func:`prelaunch_first_gather` and we skip launching it again. - - Yields exactly **2** times: - - 1. After launching async all-to-all gather (or immediately if pre-launched). - 2. After launching async all-to-all scatter. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - if prelaunch_gather is not None: - # Gather was pre-launched; none_grad already handled by caller. - work, recv_buf, gathered_grads, recv_counts = prelaunch_gather - else: - # Normal path: launch async gather. - with record_function("muon::launch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - yield # --- YIELD 1: other chunks can launch their gather --- - - with record_function("muon::wait_gather"): - work.wait() - _complete_gather(recv_buf, recv_counts, owned_params, gathered_grads, - param_to_state, rank) - del recv_buf - - # Stage 3: Newton-Schulz orthogonalization. - with record_function("muon::newton_schulz"): - computed_us = _compute_ns(owned_params, gathered_grads, ns_steps) - gathered_grads.clear() - - # Stages 4-5: launch async scatter. - with record_function("muon::launch_scatter"): - work, recv_buf, scattered_us, recv_counts = _launch_scatter( - params, owned_params, param_to_state, rank, num_ranks, - process_group, computed_us) - computed_us.clear() - - yield # --- YIELD 2: other chunks can launch their scatter --- - - with record_function("muon::wait_scatter"): - work.wait() - _complete_scatter(recv_buf, recv_counts, params, param_to_state, rank, - scattered_us) - del recv_buf - - # Stage 6: apply parameter updates. - with record_function("muon::update_params"): - _update_params(params, param_to_state, rank, scattered_us, lr, - weight_decay) - scattered_us.clear() diff --git a/build/torch28-cxx11-cu128-x86_64-linux/qk_clip.py b/build/torch28-cxx11-cu128-x86_64-linux/qk_clip.py deleted file mode 100644 index 2aba711b3004b7f09e7141da7ef834bd61cc2430..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/qk_clip.py +++ /dev/null @@ -1,198 +0,0 @@ -import logging -import math -from dataclasses import dataclass - -import torch -from torch.distributed.tensor import DTensor - -from .core import normalize_fqn - -logger = logging.getLogger(__name__) - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - and return (kind, layer_index). - - Supported kinds: - MHA/GQA: 'wq', 'wk', 'q_proj', 'k_proj' - MLA: 'wq_b' (Q up-proj), 'wkv_b' (KV up-proj) - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.1.attn.wq_b.weight' -> ('wq_b', 1) - 'model.0.attn.wkv_b.weight' -> ('wkv_b', 0) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = normalize_fqn(name).split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: str | None # 'wq'/'q_proj'/'wq_b' or 'wk'/'k_proj'/'wkv_b' or None - indices: list[int] # which heads to consider for clipping - head_dim: int # from config (qk_head_dim for MLA wq_b) - threshold: float # from config - logit: torch.Tensor | None - - # MLA-specific fields - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 - - -def get_qk_clip_info(clip_config, n, qk_logits): - """Extract QK clipping info for a named parameter. - - Args: - clip_config: QK clipping configuration dict (or None). - MHA/GQA keys: head_dim, threshold, q_indices, k_indices - MLA extra keys: is_mla=True, qk_nope_head_dim, qk_rope_head_dim, v_head_dim - n: Parameter name string. - qk_logits: Dict mapping layer indices to logit tensors (or None). - - Returns: - QKClipInfo instance with clipping configuration for this parameter. - """ - if clip_config is None: - return None - - head_dim = clip_config.get('head_dim') - threshold = clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - # In TP settings, qk_logits may be DTensor - # We convert it to full tensor here for simplicity - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - -def compute_scales(p, qk_clip_state): - """Compute per-head scaling factors for QK clipping. - - Returns scales tensor (√γ per head) if any head exceeds threshold, else None. - For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim. - """ - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - # Check if any head exceeds threshold before allocating. - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - - if not head_scales: - return None - - # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows - if qk_clip_state.is_mla and kind == 'wkv_b': - effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim - else: - effective_head_dim = head_dim - - H_global = p.shape[0] // effective_head_dim - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale - return scales_full - - -def qk_clip(p, scales, info): - """Apply per-head scaling to a Q/K projection weight matrix. - - Args: - p: Parameter (nn.Parameter or raw tensor). - scales: [n_heads] tensor, each element = √γ_h. - info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions. - - MLA sub-region scaling per Algorithm 1 (MuonClip): - wq_b: q_nope rows → √γ, q_pe rows → γ - wkv_b: k_nope rows → √γ, v rows → unchanged - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not info.is_mla: - # MHA/GQA: uniform √γ applied to all rows in each head - W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: vectorized sub-region scaling within each head - if info.kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_head_dim = qk_nope + info.qk_rope_head_dim - W_3d = W.view(-1, qk_head_dim, W.shape[1]) # [H, qk_head_dim, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # q_nope → √γ - W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1, - 1)) # q_pe → γ - - elif info.kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - W_3d = W.view(-1, kv_stride, W.shape[1]) # [H, kv_stride, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # k_nope → √γ - # v rows: not touched (k_R shared rotary unchanged) diff --git a/build/torch28-cxx11-cu129-x86_64-linux/__init__.py b/build/torch28-cxx11-cu129-x86_64-linux/__init__.py deleted file mode 100644 index 239c7a65f8293e7d0df28f05fce645af56d628c0..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .muon import Muon - -__all__ = [ - "Muon", -] diff --git a/build/torch28-cxx11-cu129-x86_64-linux/_ops.py b/build/torch28-cxx11-cu129-x86_64-linux/_ops.py deleted file mode 100644 index 034bff088659b5df6f6d401feb18c89dc5f33b29..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _optimizer_8d53b78_dirty -ops = torch.ops._optimizer_8d53b78_dirty - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_optimizer_8d53b78_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-cu129-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so b/build/torch28-cxx11-cu129-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so deleted file mode 100755 index e1e05434ff697a1f0daae635177d29e1b6f8531b..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:46baa92bf8f5ec5913df4081a01f662049fda475eb01bc7ed0f6154755fa88d5 -size 1999872 diff --git a/build/torch28-cxx11-cu129-x86_64-linux/adamw.py b/build/torch28-cxx11-cu129-x86_64-linux/adamw.py deleted file mode 100644 index b5a95816a9f5b9e1889eaadae65373bfbced809a..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/adamw.py +++ /dev/null @@ -1,271 +0,0 @@ -import logging -from collections import defaultdict -from typing import cast - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -def fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: float | torch.Tensor, - weight_decay: float, - eps: float, - maximize: bool, -) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: dict | None = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) - - -def _to_local(t): - """Unwrap DTensor to local tensor for fused ops.""" - return t._local_tensor if isinstance(t, DTensor) else t - - -# --------------------------------------------------------------------------- -# Caches for eliminating per-step Python overhead. -# -# Placement grouping and tensor list assembly are identical every step -# (params don't change placement, moment/step tensors are the same objects -# after initialisation). We cache them keyed by id() of the param list -# stored in param_groups (stable across steps). -# -# Only gradients change each step and must be collected fresh. -# --------------------------------------------------------------------------- - -# id(group["params"]) → dict[placement_key, list[param]] -_placement_cache: dict[int, dict[tuple, list]] = {} - -# id(placement_group_list) → (params_local, moment1, moment2, state_steps) -_tensor_cache: dict[int, tuple[list, list, list, list]] = {} - - -def _step_adamw_params_slow(optimizer_state, params, group): - """Uncached fallback for the rare case where some params lack grads.""" - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - g = p.grad - if g is None: - continue - state = optimizer_state[p] - params_with_grads.append(_to_local(p)) - grads.append(_to_local(g)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - if not params_with_grads: - return - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw_params(optimizer_state, params, group): - """Run fused AdamW on a list of parameters sharing the same placement. - - After the first call, cached tensor lists (params_local, moment1, - moment2, state_steps) are reused — only gradients are collected fresh. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - params: List of parameters to update. - group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay. - """ - # Collect grads — the only thing that changes each step. - with record_function("adamw::collect_grads"): - grads = [] - for p in params: - g = p.grad - if g is None: - # Rare: fall back to slow path that filters per-param. - _step_adamw_params_slow(optimizer_state, params, group) - return - grads.append(_to_local(g)) - - tensor_key = id(params) - if tensor_key not in _tensor_cache: - with record_function("adamw::init_tensor_cache"): - params_local = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - state = optimizer_state[p] - params_local.append(_to_local(p)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(p.grad) - state["moment2"] = torch.zeros_like(p.grad) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - _tensor_cache[tensor_key] = (params_local, moment1, moment2, - state_steps) - - params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key] - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - with record_function("adamw::fused_adamw"): - fused_adamw( - params_local, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw(optimizer_state, group): - """Dispatch AdamW step, grouping parameters by type and placement. - - Placement grouping is cached after the first call since params never - change their placement between steps. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - group: Parameter group dict. - """ - params = group["params"] - placement_key = id(params) - - if placement_key not in _placement_cache: - with record_function("adamw::group_by_placement"): - placement_to_params: dict[tuple, - list[torch.Tensor]] = defaultdict(list) - for p in params: - match p: - case DTensor(): - logger.debug( - "[AdamW] DTensor param: shape=%s, placements=%s, " - "mesh=%s, grad=%s", p.shape, p.placements, - p.device_mesh.mesh_dim_names, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple( - [p.placements, p.device_mesh])].append(p) - case torch.Tensor(): - logger.debug( - "[AdamW] plain param: shape=%s, grad=%s", p.shape, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple([torch.Tensor, - None])].append(p) - - logger.debug("[AdamW] %d placement groups, %d total params", - len(placement_to_params), len(params)) - - _placement_cache[placement_key] = dict(placement_to_params) - - for group_params in _placement_cache[placement_key].values(): - step_adamw_params(optimizer_state, group_params, group) diff --git a/build/torch28-cxx11-cu129-x86_64-linux/async_utils.py b/build/torch28-cxx11-cu129-x86_64-linux/async_utils.py deleted file mode 100644 index a45c530ac9cad88e3555ec1047a6aa59f225347e..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/async_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import logging -from typing import Generator - -logger = logging.getLogger(__name__) - - -class _Task: - """Internal: wraps a generator, advances one yield at a time.""" - - def __init__(self, generator: Generator[None, None, None], index: int): - self._generator = generator - self._index = index - self._steps_completed = 0 - self.step() # run to first yield - - def step(self) -> bool: - try: - next(self._generator) - self._steps_completed += 1 - logger.debug("pipeline[%d] completed stage %d", self._index, - self._steps_completed) - return True - except StopIteration: - logger.debug("pipeline[%d] finished after %d stages", self._index, - self._steps_completed) - return False - - def close(self): - self._generator.close() - - -def run_pipeline( - pipelines: Generator[Generator[None, None, None], None, None], - max_concurrent: int, -) -> None: - """Run generator-based pipelines with bounded concurrency. - - Each pipeline is a generator that yields at stage boundaries. - The runtime interleaves pipelines so communication and computation - overlap across chunks. - """ - if max_concurrent <= 0: - raise ValueError(f"max_concurrent must be > 0, got {max_concurrent}") - - have_new = True - task_index = 0 - previous_tasks: list[_Task] = [] - - try: - while have_new or previous_tasks: - running_tasks: list[_Task] = [] - - # Admit one new pipeline per iteration (staggered admission). - # Admitting one at a time ensures that while chunk N does NS - # compute on the default stream, chunk N+1's NCCL all-to-all - # runs concurrently on the NCCL stream — creating real - # communication/computation overlap on the GPU. - if have_new and len(previous_tasks) < max_concurrent: - try: - gen = next(pipelines) - task = _Task(gen, task_index) - task_index += 1 - running_tasks.append(task) - except StopIteration: - have_new = False - - # Advance every previously-yielded task by one step. - for task in previous_tasks: - if task.step(): - running_tasks.append(task) - - previous_tasks = running_tasks - except BaseException: - # Clean up all in-flight generators to release GPU resources. - for task in previous_tasks: - task.close() - raise diff --git a/build/torch28-cxx11-cu129-x86_64-linux/core.py b/build/torch28-cxx11-cu129-x86_64-linux/core.py deleted file mode 100644 index c69d515afef305ad0ed66374095fa2d2468d99cc..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/core.py +++ /dev/null @@ -1,219 +0,0 @@ -import logging -import math -from dataclasses import dataclass -from typing import List - -import torch -from torch.distributed import ProcessGroup -from torch.distributed.tensor import DTensor - -# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into -# parameter FQNs. Activation checkpointing similarly inserts -# "_checkpoint_wrapped_module". Strip these so name-based matching (skip_keys, -# expert_keys, QK layer parsing) works regardless of wrapper nesting. -_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) - -logger = logging.getLogger(__name__) - - -def normalize_fqn(name: str) -> str: - """Strip torch.compile / checkpoint wrapper components from a parameter FQN.""" - return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS) - - -@dataclass -class _muon_state: - worker_rank: int - process_group: ProcessGroup - rank_indices: dict[int, tuple] # local_rank -> per-dim indices - rank_numels: dict[int, int] # local_rank -> numel - name: str - qk_clip_state: torch.Tensor | None = None - - -def _batch_momentum( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update (no nesterov).""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - - -def _batch_momentum_nesterov( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update with nesterov correction.""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - nesterov_terms = torch._foreach_mul(momentum_bufs, momentum) - torch._foreach_add_(grads, nesterov_terms) - - -_compiled_momentum: dict[bool, callable] = {} -_use_momentum_compile = True - - -def set_momentum_compile(enabled: bool): - """Toggle torch.compile for batched momentum.""" - global _use_momentum_compile - _use_momentum_compile = enabled - - -def batch_pre_ortho( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, - nesterov: bool, -) -> None: - """Batched momentum update on lists of plain tensors. - - Mirrors dion's ``muon_update_pre_orthogonalize``. - Inputs must be plain CUDA tensors (not DTensor). - Modifies ``momentum_bufs`` and (for nesterov) ``grads`` in-place. - - When compile is enabled, uses separately compiled functions for - nesterov=True/False to avoid graph breaks from the branch. - """ - fn = _batch_momentum_nesterov if nesterov else _batch_momentum - if _use_momentum_compile: - if nesterov not in _compiled_momentum: - _compiled_momentum[nesterov] = torch.compile(fn) - fn = _compiled_momentum[nesterov] - fn(grads, momentum_bufs, momentum) - - -def _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay): - """Weight-decay + update on plain tensors. - - Not compiled: per-param @torch.compile caused ~0.25ms TorchDynamo cache - lookup per call × 256+ params = massive overhead. The pipeline path uses - batched _foreach_* ops instead; this function remains for base() and - distributed_muon(). - """ - p_data.mul_(1 - lr * weight_decay) - p_data.add_(u_data, alpha=-adjusted_lr) - - -def update_p(p, u, lr, adjusted_lr, weight_decay): - """Apply weight decay and orthogonalized update to parameter. - - Args: - p: Parameter (torch.nn.Parameter or DTensor). - u: Orthogonalized update tensor. - lr: Base learning rate. - adjusted_lr: Size-adjusted learning rate. - weight_decay: Weight decay coefficient. - """ - # Unwrap Parameter -> underlying data tensor. - p_data = p.data if isinstance(p, torch.nn.Parameter) else p - # Unwrap DTensor -> local CUDA tensor for compiled kernel. - if isinstance(p_data, DTensor): - p_data = p_data._local_tensor - u_data = u._local_tensor if isinstance(u, DTensor) else u - _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay) - - -def adjust_lr_for_muon(lr, param_shape): - """Scale learning rate based on parameter matrix dimensions. - - Args: - lr: Base learning rate. - param_shape: Shape of the parameter tensor. - - Returns: - Adjusted learning rate. - """ - A, B = param_shape[:2] - # We adjust the learning rate and weight decay based on the size of the parameter matrix - # as described in the paper - adjusted_ratio = 0.2 * math.sqrt(max(A, B)) - adjusted_lr = lr * adjusted_ratio - return adjusted_lr - - -def _match_key(parts, key): - """Check if key matches as contiguous components in parts. - - Single-component keys (e.g. "experts") match any single component. - Multi-component keys (e.g. "experts.w1") match as a contiguous subsequence. - """ - key_parts = key.split(".") - key_len = len(key_parts) - if key_len == 1: - return key in parts - return any(parts[i:i + key_len] == key_parts - for i in range(len(parts) - key_len + 1)) - - -def is_expert_param(name, expert_keys): - """Check if a parameter name matches any expert key (component-level).""" - if not expert_keys: - return False - parts = normalize_fqn(name).split(".") - return any(_match_key(parts, key) for key in expert_keys) - - -def default_is_muon(name, x, expert_keys=None): - normalized = normalize_fqn(name) - parts = normalized.split(".") - skip_keys = [ - "embed_tokens", - "lm_head", - "tok_embeddings", - "output", - "mhc_attn", - "mhc_ffn", - "lambda_proj", - ] - if any(key in parts for key in skip_keys): - logger.info( - "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d", - normalized, name, x.ndim) - return False - effective_ndim = x.ndim - is_expert = is_expert_param(name, expert_keys) - if is_expert: - effective_ndim -= 1 - result = effective_ndim >= 2 - logger.info( - "[is_muon] %s (orig: %s): ndim=%d, expert=%s, effective_ndim=%d → %s", - normalized, name, x.ndim, is_expert, effective_ndim, - "Muon" if result else "AdamW") - return result - - -def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None): - if is_muon_func is None: - is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys) - - muon_params, muon_names = [], [] - non_muon_params, non_muon_names = [], [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - non_muon_names.append(n) - - logger.info("[param_groups] expert_keys=%s, Muon=%d, AdamW=%d", - expert_keys, len(muon_names), len(non_muon_names)) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] diff --git a/build/torch28-cxx11-cu129-x86_64-linux/cpu_offload.py b/build/torch28-cxx11-cu129-x86_64-linux/cpu_offload.py deleted file mode 100644 index fb5e69154a1d4a6c884491413a37a9acf0f66c80..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/cpu_offload.py +++ /dev/null @@ -1,206 +0,0 @@ -"""CPU offloading for optimizer states. - -Manages a pinned CPU memory pool and async CUDA streams to offload -optimizer state tensors (momentum buffers, Adam moments) to CPU between -optimizer steps, freeing GPU memory. - -All tracked tensors are packed into a single flat pinned CPU buffer -(per dtype). D2H and H2D copies are performed per-tensor directly -between individual GPU tensors and their slice of the CPU flat buffer -— no GPU staging buffer is allocated, so there is **no temporary GPU -memory spike** during offload or reload. - -Individual tensor storages are freed after offload via -``untyped_storage().resize_(0)``, preserving tensor identity so -downstream caches remain valid. -""" - -import logging -from collections import defaultdict - -import torch -from torch.distributed.tensor import DTensor - -logger = logging.getLogger(__name__) - - -class CPUOffloadPool: - """Pinned CPU memory pool for async optimizer state offloading. - - Tracked tensors are grouped by dtype. Each group gets a single flat - pinned CPU buffer. D2H / H2D copies are per-tensor (into slices of - the flat buffer) to avoid allocating a GPU staging buffer. - """ - - def __init__(self): - self._managed: list[torch.Tensor] = [] - self._storage_nbytes: dict[int, int] = {} # id(t) → bytes - - # Per-dtype group: populated on first offload. - # dtype → dict with keys: - # "indices" : list[int] managed-list indices - # "offsets" : list[tuple[int,int]] (start, numel) in flat buf - # "total" : int total numel - # "cpu_flat" : Tensor pinned CPU buffer - self._groups: dict[torch.dtype, dict] = {} - - self._offload_stream: torch.cuda.Stream | None = None - self._device: torch.device | None = None - self._initialized: bool = False - self._logged: bool = False - - # ------------------------------------------------------------------ - @staticmethod - def _local(t: torch.Tensor) -> torch.Tensor: - """Unwrap DTensor to its local CUDA tensor.""" - return t._local_tensor if isinstance(t, DTensor) else t - - def _ensure_stream(self): - if self._offload_stream is None: - self._offload_stream = torch.cuda.Stream(device=self._device) - - # ------------------------------------------------------------------ - def track(self, tensor: torch.Tensor): - """Register a GPU tensor for CPU offloading. Idempotent.""" - tid = id(tensor) - if tid in self._storage_nbytes: - return - local = self._local(tensor) - if self._device is None: - self._device = local.device - storage = local.untyped_storage() - # Skip tensors with empty storage (e.g. empty FSDP shards) - if storage.size() == 0: - return - self._storage_nbytes[tid] = storage.size() - self._managed.append(tensor) - - # ------------------------------------------------------------------ - def _init_buffers(self): - """Build per-dtype flat buffers on first offload.""" - # Group managed tensors by dtype. - dtype_map: dict[torch.dtype, list[tuple[int, int]]] = defaultdict(list) - for idx, t in enumerate(self._managed): - local = self._local(t) - dtype_map[local.dtype].append((idx, local.numel())) - - total_cpu_bytes = 0 - for dtype, entries in dtype_map.items(): - offsets: list[tuple[int, int]] = [] - indices: list[int] = [] - off = 0 - for idx, n in entries: - indices.append(idx) - offsets.append((off, n)) - off += n - cpu_flat = torch.empty(off, dtype=dtype, device="cpu", pin_memory=True) - self._groups[dtype] = { - "indices": indices, - "offsets": offsets, - "total": off, - "cpu_flat": cpu_flat, - } - total_cpu_bytes += off * cpu_flat.element_size() - - self._initialized = True - logger.info( - "[CPUOffload] Pool initialized: %d tensors, %d dtype group(s), " - "%.2f MB pinned CPU memory", - len(self._managed), - len(self._groups), - total_cpu_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def offload(self): - """Per-tensor async D2H into CPU flat buffer, then free GPU storage.""" - if not self._managed: - return - if not self._initialized: - self._init_buffers() - self._ensure_stream() - - # Offload stream waits for compute to finish. - compute_event = torch.cuda.current_stream(self._device).record_event() - self._offload_stream.wait_event(compute_event) - - offloaded_bytes = 0 - - # Per-tensor D2H copies directly into CPU flat buffer slices. - # No GPU staging buffer → no temporary GPU memory spike. - with torch.cuda.stream(self._offload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - cpu_flat[off : off + n].copy_(local.reshape(-1), non_blocking=True) - - offloaded_bytes += grp["total"] * cpu_flat.element_size() - - # Wait for all D2H copies to land, then free GPU storage. - self._offload_stream.synchronize() - for t in self._managed: - storage = self._local(t).untyped_storage() - if storage.size() != 0: - storage.resize_(0) - else: - raise RuntimeError( - f"Tensor storage is already freed (size=0) before offload. " - f"This indicates a double-free or external interference. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - - if not self._logged: - logger.info( - "[CPUOffload] Offloaded %.2f MB (GPU → CPU)", - offloaded_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def reload(self): - """Per-tensor H2D from CPU flat buffer on the default stream. - - Runs on the current (default) CUDA stream to avoid stream - interaction issues with the parallel Muon pipeline. Since - pinned CPU memory is the source, the copies overlap with - GPU idle time between steps. - """ - if not self._managed or not self._initialized: - return - - reloaded_bytes = 0 - - # Re-allocate all GPU storages first. - for t in self._managed: - local = self._local(t) - storage = local.untyped_storage() - if storage.size() != 0: - raise RuntimeError( - f"Storage should have been freed (size=0) before reload, " - f"but got size={storage.size()}. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - storage.resize_(self._storage_nbytes[id(t)]) - - # Per-tensor H2D copies from CPU flat buffer slices. - # non_blocking=True with pinned source allows DMA overlap. - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off : off + n], non_blocking=True) - - reloaded_bytes += grp["total"] * cpu_flat.element_size() - - if not self._logged: - logger.info( - "[CPUOffload] Reloaded %.2f MB (CPU → GPU)", reloaded_bytes / (1024**2) - ) diff --git a/build/torch28-cxx11-cu129-x86_64-linux/distributed/utils.py b/build/torch28-cxx11-cu129-x86_64-linux/distributed/utils.py deleted file mode 100644 index 683096bcd50489e7a09959753fdbc29ba0f3b6a6..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/distributed/utils.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import (Placement, Shard, - _StridedShard) - - -def _is_shard(placement: Placement) -> bool: - """Check if a placement is a shard type (Shard or _StridedShard). - - In PyTorch 2.10+, _StridedShard no longer inherits from Shard, so - ``placement.is_shard()`` returns False for _StridedShard. This helper - handles both old and new hierarchies. - """ - return isinstance(placement, (Shard, _StridedShard)) - - -def get_slices_of_dtensor( - target: DTensor | torch.Tensor, - local_rank: int, - shard_mesh: DeviceMesh, - shard_placements: tuple[Placement], -) -> tuple[slice | torch.Tensor, ...]: - """ - Get per-dimension indices for a given rank's shard of the target tensor. - - Uses ``Shard.local_shard_size_and_offset`` and - ``_StridedShard.local_shard_size_and_offset`` for correct handling of - both contiguous and strided (non-contiguous) sharding. - - Args: - target (DTensor | torch.Tensor): The target tensor (for its shape). - local_rank (int): The local rank within the shard group. - shard_mesh (DeviceMesh): The shard mesh (only shard dimensions). - shard_placements (tuple[Placement]): The shard placements. - - Returns: - A tuple of indices (one per tensor dim). Each element is either: - - A ``slice`` (for contiguous or unsharded dims) - - A 1-D ``torch.LongTensor`` of indices (for strided sharding) - """ - - # find the global rank of the local rank in the shard mesh - rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank] - - rank_coords = (shard_mesh.mesh == rank).nonzero() - - assert len(rank_coords) == 1 - rank_coords = tuple(rank_coords[0].tolist()) - - assert len(rank_coords) == len(shard_placements) - - # Track per-shard-dim indices. - # None means "not yet sharded on this dim". - dim_indices: dict[int, torch.Tensor] = {} - - # Caution: Assuming replicate-to-shard of the shard mesh goes with - # left-to-right sharding. This is ensured by the sorting logic of - # construct_shard_mesh function. - for mesh_dim_idx, (rank_coord, placement) in enumerate( - zip(rank_coords, shard_placements)): - assert _is_shard(placement) - - num_chunks = shard_mesh.mesh.shape[mesh_dim_idx] - shard_dim = placement.dim - - # Current effective size on this dim (may already be sub-sharded) - if shard_dim in dim_indices: - curr_size = len(dim_indices[shard_dim]) - else: - curr_size = target.size()[shard_dim] - - # Compute indices for this level of sharding - if isinstance(placement, _StridedShard): - _shard_size, offsets = _StridedShard.local_shard_size_and_offset( - placement, - curr_size, - num_chunks, - rank_coord, - return_first_offset=False) - new_indices = torch.tensor(offsets, dtype=torch.long) - else: - shard_size, offset = Shard.local_shard_size_and_offset( - curr_size, num_chunks, rank_coord) - new_indices = torch.arange(offset, - offset + shard_size, - dtype=torch.long) - - # Compose with previous indices on this dim - if shard_dim in dim_indices: - dim_indices[shard_dim] = dim_indices[shard_dim][new_indices] - else: - dim_indices[shard_dim] = new_indices - - # Build result tuple - result: list[slice | torch.Tensor] = [] - for d in range(len(target.size())): - if d not in dim_indices: - result.append(slice(None)) - else: - indices = dim_indices[d] - # Convert contiguous indices to slice for efficiency - if len(indices) > 0: - start = indices[0].item() - expected = torch.arange(start, - start + len(indices), - dtype=torch.long) - if torch.equal(indices, expected): - result.append(slice(start, start + len(indices))) - else: - result.append(indices) - else: - result.append(slice(0, 0)) - - return tuple(result) - - -_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, - ProcessGroup]] = dict() - - -def construct_shard_mesh( - placements: tuple[Placement], - mesh: DeviceMesh, -) -> tuple[DeviceMesh, ProcessGroup, tuple[Placement, ...]]: - """Construct shard sub-mesh and ProcessGroup for all-to-all communication. - - Given a DTensor's placements and device mesh, extracts the "shard group" - — the set of ranks that together hold all shards of the same replica — - and creates a ProcessGroup for all-to-all among them. - - Steps: - 1. Sort placements: Replicate first, then Shard by (dim, granularity). - 2. Permute the mesh tensor to match the sorted order. - 3. Collapse Replicate dims → list of shard sub-meshes (one per replica). - 4. Create/retrieve a cached ProcessGroup for the current rank's sub-mesh. - - Example — 8 GPUs, mesh shape (2, 2, 2), - placements ``[Shard(0), Replicate, _StridedShard(0)]``:: - - Step 1 — Sort: [Replicate, _StridedShard(0), Shard(0)] - Permutation: [1, 2, 0] - - Step 2 — Permute mesh dims by [1, 2, 0]: - Original: Permuted: - [[[0,1],[2,3]], [[[0,2],[1,3]], - [[4,5],[6,7]]] [[4,6],[5,7]]] - - Step 3 — Unbind replicate dim (dim 0), giving 2 shard sub-meshes: - sub-mesh 0 = [[0,2],[1,3]] (replica group 0) - sub-mesh 1 = [[4,6],[5,7]] (replica group 1) - shard_placements = (_StridedShard(0), Shard(0)) - - Step 4 — Rank 0 → ProcessGroup([0,1,4,5]) - Rank 2 → ProcessGroup([2,3,6,7]) - - Returns: - ``(shard_mesh, process_group, shard_placements)`` - """ - my_rank = dist.get_rank() - assert mesh.mesh.device.type == 'cpu' - - # -- Fast path: 1D all-shard mesh → reuse existing PG. ---------------- - # Reuses the mesh's existing ProcessGroup directly, avoiding the - # overhead of dist.new_group(). The standard path below also handles - # subset calls safely via use_local_synchronization=True, but this - # fast path is still beneficial for the common 1D shard case. - if mesh.ndim == 1 and len(placements) == 1 and _is_shard(placements[0]): - key = (*mesh.mesh.shape, *mesh.mesh.flatten().tolist()) - if key not in _ranks_to_dist_cache: - _ranks_to_dist_cache[key] = (mesh, mesh.get_group()) - return (*_ranks_to_dist_cache[key], tuple(placements)) - - mesh_tensor = mesh.mesh.clone() - - # -- Step 1: Sort placements (Replicate first, then Shard by dim). ------ - # _StridedShard comes BEFORE regular Shard on the same dim so that - # get_slices_of_dtensor applies the outer sharding first, matching - # DTensor's left-to-right (outer-to-inner) composition order. - def _sort_key(item): - index, placement = item - assert not placement.is_partial(), "Partial placement not supported" - if placement.is_replicate(): - return (-1, 0, index) - assert _is_shard(placement), f"Unsupported: {type(placement)}" - split = (-1 / placement.split_factor if isinstance( - placement, _StridedShard) else 0) - return (placement.dim, split, index) - - indexed = sorted(enumerate(placements), key=_sort_key) - perm, sorted_placements = zip(*indexed) - - # -- Step 2: Permute mesh to match sorted placement order. -------------- - sorted_mesh = mesh_tensor.permute(perm) - - # -- Step 3: Collapse replicate dims → list of shard sub-meshes. -------- - # E.g. mesh (2, 3, 4, 4) with [R, R, S(0), S(1)] → 6 sub-meshes of (4, 4) - num_rep = sum(1 for p in sorted_placements if p.is_replicate()) - if num_rep > 0: - if num_rep > 1: - sorted_mesh = sorted_mesh.flatten(0, num_rep - 1) - shard_meshes = list(torch.unbind(sorted_mesh, dim=0)) - else: - shard_meshes = [sorted_mesh] - shard_placements = sorted_placements[num_rep:] - assert len(shard_placements) == len(set(shard_placements)) - - # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. -- - # Each rank only creates the group it belongs to, using - # use_local_synchronization=True so that only group members need to - # coordinate. This avoids deadlocks when different PP stages call - # construct_shard_mesh for different parameters. - def _cache_key(t: torch.Tensor) -> tuple: - return (*t.shape, *t.flatten().tolist()) - - my_key = None - for sm in shard_meshes: - if (my_rank == sm).any().item(): - key = _cache_key(sm) - assert my_key is None, "Rank appears in multiple shard groups" - my_key = key - if key not in _ranks_to_dist_cache: - pg = dist.new_group(sm.flatten().tolist(), - use_local_synchronization=True) - _ranks_to_dist_cache[key] = ( - DeviceMesh(device_type="cuda", mesh=sm), - pg, - ) - - return (*_ranks_to_dist_cache[my_key], shard_placements) diff --git a/build/torch28-cxx11-cu129-x86_64-linux/matmul_transpose_triton.py b/build/torch28-cxx11-cu129-x86_64-linux/matmul_transpose_triton.py deleted file mode 100644 index 792de23d82c3fb45fe33d397ab9b76a0787259d0..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/matmul_transpose_triton.py +++ /dev/null @@ -1,122 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], - restore_value=['y'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -@torch.library.custom_op("muon::matmul_transpose_assign", - mutates_args=("d_out", )) -def matmul_transpose_assign(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """Compute d_out = d_in @ d_in.T using an optimized Triton kernel.""" - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -@matmul_transpose_assign.register_fake -def _(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """FakeTensor impl: d_out is already allocated, mutation is declared.""" - pass diff --git a/build/torch28-cxx11-cu129-x86_64-linux/metadata.json b/build/torch28-cxx11-cu129-x86_64-linux/metadata.json deleted file mode 100644 index c55a35717622f1dd5c8ba376ea3a814cbcc10d78..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/metadata.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch28-cxx11-cu129-x86_64-linux/muon.py b/build/torch28-cxx11-cu129-x86_64-linux/muon.py deleted file mode 100644 index 14c0e22471fa6d47a51ed95e0e0c341dc18d5194..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/muon.py +++ /dev/null @@ -1,1068 +0,0 @@ -import logging -import types -from collections import defaultdict -from typing import Any - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Replicate, Shard -from torch.profiler import record_function - -from .adamw import _placement_cache, _tensor_cache, step_adamw -from .async_utils import run_pipeline -from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho, - get_default_muon_param_groups, is_expert_param, update_p) -from .cpu_offload import CPUOffloadPool -from .distributed.utils import (_is_shard, construct_shard_mesh, - get_slices_of_dtensor) -from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO, - _zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5_batched) -from .pipeline import muon_chunk_pipeline, prelaunch_first_gather -from .qk_clip import compute_scales, get_qk_clip_info, qk_clip - -logger = logging.getLogger(__name__) - - -def _expand_expert_params(names, params, expert_keys): - """Expand expert params by splitting on dim 0 (expert dimension). - - Params whose name matches any key in ``expert_keys`` are treated as - expert-parallel tensors. Their outermost dimension is the expert - dimension: an ``(E, out, in)`` tensor becomes ``E`` separate 2D - ``nn.Parameter`` views so that in-place updates propagate back to - the original storage. - - Non-expert params with ``ndim > 2`` trigger an ``AssertionError`` — - if they are expert params, their key must be added to ``expert_keys``. - - The grad must already be set on each expert param (e.g. after momentum). - - For DTensor expert params, placements that shard on dim 0 (expert dim) - are consumed by the split. Non-dim-0 shard placements (e.g. TP) are - preserved: each 2D slice is wrapped as a DTensor on the corresponding - submesh so the parallel pipeline handles the TP communication. - """ - expanded_names = [] - expanded_params = [] - - for n, p in zip(names, params): - is_expert = is_expert_param(n, expert_keys) - is_dtensor = isinstance(p.data, DTensor) - - if is_expert: - if is_dtensor: - logger.debug( - "[expand_expert] %s: expert DTensor, shape=%s, " - "placements=%s, mesh=%s, local_shape=%s", n, p.shape, - p.placements, p.device_mesh.mesh_dim_names, - p.to_local().shape) - else: - logger.debug( - "[expand_expert] %s: expert plain tensor, shape=%s", n, - p.data.shape) - - if not is_expert: - assert p.data.ndim <= 2, ( - f"Param {n} has ndim={p.data.ndim} but does not match " - f"expert_keys={expert_keys}. If this is an expert param, " - f"add its key to expert_keys.") - expanded_names.append(n) - expanded_params.append(p) - continue - - g = p.grad - assert g is not None, ( - f"Expert param {n} must have grad set before expansion") - - tp_mesh = None - tp_placements_2d = None - - if is_dtensor: - local_data = p.to_local() - local_grad = g.to_local() if isinstance(g, DTensor) else g - - # Find non-dim-0 shard placements (e.g. TP sharding). - # After splitting on dim 0, Shard(k) becomes Shard(k-1). - tp_dim_indices = [] - tp_placements_2d = [] - for i, pl in enumerate(p.placements): - if _is_shard(pl) and pl.dim != 0: - tp_dim_indices.append(i) - tp_placements_2d.append(Shard(pl.dim - 1)) - - if tp_dim_indices: - tp_dim_names = tuple(p.device_mesh.mesh_dim_names[i] - for i in tp_dim_indices) - if len(tp_dim_names) == 1: - tp_mesh = p.device_mesh[tp_dim_names[0]] - else: - tp_mesh = p.device_mesh[tp_dim_names] - else: - local_data = p.data - local_grad = g - - # Expand: split dim 0, reshape each slice to 2D. - num_local_experts = local_data.shape[0] - for i in range(num_local_experts): - slice_data = local_data[i] - slice_grad = local_grad[i] - - if tp_mesh is not None: - # Wrap as DTensor on TP submesh so the pipeline handles - # TP communication (gather/scatter across TP ranks). - dt_data = DTensor.from_local(slice_data, - device_mesh=tp_mesh, - placements=tp_placements_2d) - dt_grad = DTensor.from_local(slice_grad, - device_mesh=tp_mesh, - placements=tp_placements_2d) - expert_param = torch.nn.Parameter(dt_data, requires_grad=False) - expert_param.grad = dt_grad - else: - expert_param = torch.nn.Parameter(slice_data, - requires_grad=False) - expert_param.grad = slice_grad - - expanded_names.append(f"{n}[{i}]") - expanded_params.append(expert_param) - - p.grad = None # allow expert grad storage to be freed after pipeline - - return expanded_names, expanded_params - - -class Muon(torch.optim.Optimizer): - """ - Muon - MomentUm Orthogonalized by Newton-schulz - - Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- - processing step, in which each 2D parameter's update is replaced with the nearest orthogonal - matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has - the advantage that it can be stably run in bfloat16 on the GPU. - - Some warnings: - - We believe this optimizer is unlikely to work well for training with small batch size. - - We believe it may not work well for finetuning pretrained models, but we haven't tested this. - - Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. - lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) - momentum: The momentum used by the internal SGD. (0.95 is a good default) - nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) - ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. - Parameters that are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW instead. - adamw_lr: The learning rate for the internal AdamW. - adamw_betas: The betas for the internal AdamW. - adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - warmup_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher warmup_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. - chunk_size : Batch size of parameters to process in each - all2all gather/compute/scatter step. - Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified. - use_distributed_muon: Use distributed muon by Liu et al. (2024). - For testing purpose only. - expert_keys: List of strings to identify expert-parallel parameters. - If any key appears in a parameter's name, its outermost - dimension is treated as the expert dimension and expanded - into per-expert 2D params for Muon. For example, - ``expert_keys=["experts"]`` matches any param whose name - contains "experts". 3D+ params not matched by any key - will raise an error. - """ - - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config=None, - warmup_step=5, - chunk_size=-1, - use_distributed_muon=False, - expert_keys=None): - defaults = dict( - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - nesterov=nesterov, - ns_steps=ns_steps, - adamw_betas=adamw_betas, - adamw_eps=adamw_eps, - none_grad=none_grad, - use_muon=True, - ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) - super().__init__(params, defaults) - - self.debug = debug - self.clip_config = clip_config if clip_config is not None else { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100, - } - self.warmup_step = warmup_step - self.chunk_size = chunk_size - self.use_distributed_muon = use_distributed_muon - self.expert_keys = expert_keys - self.cpu_offload = False - self._cpu_offload_pool: CPUOffloadPool | None = None - self._offload_initialized = False - self._parallel_cache: dict[tuple[str, ...], dict] = {} - self._expert_expand_cache: dict[tuple[int, ...], dict] = {} - - def _calc_flops(self, G, steps): - assert len(G.shape) == 2 - M, N = G.shape - if M > N: - M, N = N, M - - return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) - - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - p.placements, p.device_mesh) - - return shard_mesh, shard_pg, shard_placements - - def init_state_and_assign_params(self, names, params, group, qk_logits): - param_to_state = {} - param_to_flops = {} - - total_flops = 0 - for p in params: - g = p.grad - if g is None: - continue - assert g.ndim == 2, "Muon only supports 2D parameters." - - flops = self._calc_flops(g, group["ns_steps"]) - param_to_flops[id(p)] = flops - total_flops += flops - - if self.debug: - logger.debug("Total TFLOPs for Muon: %.2f TFLOPs", - total_flops / 1e12) - - paired = list(zip(names, params)) - - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) - - round_robin = 0 - mesh = ordered_params[0].device_mesh - placements = ordered_params[0].placements - - shard_mesh, shard_pg, shard_placements = self.get_shard_mesh( - ordered_params[0]) - shard_mesh_flattened = shard_mesh.mesh.flatten() - num_ranks = dist.get_world_size(group=shard_pg) - - for n, p in zip(ordered_names, ordered_params): - if mesh != p.device_mesh: - raise ValueError("All parameters must be on the same mesh.") - if placements != p.placements: - raise ValueError("All parameters must have same placements.") - - worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks - round_robin = (round_robin + 1) % len(shard_mesh_flattened) - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - # Precompute per-rank indices and numels for all-to-all. - rank_indices: dict[int, tuple] = {} - rank_numels: dict[int, int] = {} - for r in range(num_ranks): - indices = get_slices_of_dtensor(p, r, shard_mesh, - shard_placements) - rank_indices[r] = indices - numel = 1 - for idx, dim_size in zip(indices, p.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) - rank_numels[r] = numel - - param_to_state[id(p)] = _muon_state( - worker_rank=worker_rank, - process_group=shard_pg, - rank_indices=rank_indices, - rank_numels=rank_numels, - name=n, - qk_clip_state=qk_clip_state, - ) - - return param_to_state, ordered_params - - def base(self, names, params, group, lr, weight_decay, qk_logits): - # Momentum is already applied by _step_muon before this method. - for n, p in zip(names, params): - g = p.grad - if g is None: - continue - - u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) - - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - def distributed_muon( - self, - names: list[str], - params: list[torch.nn.Parameter], - group: dict[str, Any], - lr: float, - weight_decay: float, - qk_logits: list[torch.Tensor | DTensor] | None, - ): - """Batched Distributed Muon — for testing/correctness verification only. - - Uses all-gather to reconstruct full tensors, computes Newton-Schulz on - the full grad, then slices back to local shards. This is simpler but - slower than the parallel pipeline (all2all) path, so it serves as a - reference implementation for verifying correctness. - """ - with record_function("distributed_muon"): - # Momentum is already applied by _step_muon before this method. - ns_steps = group["ns_steps"] - - # Separate plain tensors (no communication) from DTensors. - plain_names, plain_params = [], [] - dtensor_names, dtensor_params = [], [] - for n, p in zip(names, params): - if p.grad is None: - continue - if isinstance(p.data, DTensor): - dtensor_names.append(n) - dtensor_params.append(p) - else: - plain_names.append(n) - plain_params.append(p) - - # Process plain tensors per-param (no communication). - for n, p in zip(plain_names, plain_params): - u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE), - steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, - qk_logits) - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - if not dtensor_params: - return - - # Group DTensors by (placements, mesh) for batched all-gather. - placement_groups: dict[tuple, - tuple[list, - list]] = defaultdict(lambda: ([], [])) - for n, p in zip(dtensor_names, dtensor_params): - key = (p.placements, p.device_mesh) - placement_groups[key][0].append(n) - placement_groups[key][1].append(p) - - logger.info( - "distributed_muon: %d placement groups, %d total dtensors", - len(placement_groups), len(dtensor_params)) - - for (placements, mesh), (grp_names, - grp_params) in placement_groups.items(): - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - placements, mesh) - rank = dist.get_rank(shard_pg) - world_size = dist.get_world_size(shard_pg) - - logger.info(" group: %d params, placements=%s, world_size=%d", - len(grp_params), placements, world_size) - - # Separate params that can be batched (all shard dims evenly - # divisible) from those needing per-param full_tensor - # (e.g. MoE gate weights with fewer rows than shard ranks). - # all_gather_into_tensor requires equal buffer sizes across - # ranks, so uneven splits must use DTensor full_tensor(). - batch_names, batch_params = [], [] - single_names, single_params = [], [] - for n, p in zip(grp_names, grp_params): - even = all(p.shape[pl.dim] % - shard_mesh.mesh.shape[dim_idx] == 0 - for dim_idx, pl in enumerate(shard_placements)) - if even: - batch_names.append(n) - batch_params.append(p) - else: - single_names.append(n) - single_params.append(p) - - # Process uneven-split params per-param via full_tensor(). - for n, p in zip(single_names, single_params): - with record_function("distributed_muon::newton_schulz"): - g_full = p.grad.full_tensor().to(COMM_DTYPE) - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - if not batch_params: - continue - - logger.info(" batched=%d, single=%d", len(batch_params), - len(single_params)) - - # Concat all local grad shards into a single flat buffer. - with record_function("distributed_muon::gather"): - grad_locals = [ - p.grad.to_local().to(COMM_DTYPE).flatten() - for p in batch_params - ] - numels = [g.numel() for g in grad_locals] - grad_concat = torch.cat(grad_locals) - del grad_locals - - # Single all-gather (replaces N separate full_tensor). - grad_gathered = torch.empty( - grad_concat.numel() * world_size, - dtype=COMM_DTYPE, - device="cuda", - ) - dist.all_gather_into_tensor(grad_gathered, - grad_concat, - group=shard_pg) - - total_numel = grad_concat.numel() - del grad_concat - - # Precompute per-param offsets within the concat buffer. - offsets = [] - off = 0 - for ne in numels: - offsets.append(off) - off += ne - - # Per-param: reconstruct full grad → NS → local update. - for i, (n, p) in enumerate(zip(batch_names, batch_params)): - with record_function("distributed_muon::newton_schulz"): - g_full = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - for r in range(world_size): - r_start = r * total_numel + offsets[i] - shard = grad_gathered[r_start:r_start + numels[i]] - indices = get_slices_of_dtensor( - p, r, shard_mesh, shard_placements) - g_full[indices] = shard.reshape( - g_full[indices].shape) - - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - def _setup_parallel(self, names, params, group, qk_logits): - """Compute (or retrieve cached) parallel pipeline metadata. - - Returns: - (ordered_params, param_to_state, rank, chunk_size) - """ - cache_key = tuple(names) - - if cache_key not in self._parallel_cache: - # First call: compute metadata and populate cache. - param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - shard_pg = param_to_state[id(ordered_params[0])].process_group - rank = dist.get_rank(group=shard_pg) - - if self.chunk_size == -1: - shard_ranks = dist.get_world_size(shard_pg) - chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO - elif self.chunk_size > 0: - chunk_size = self.chunk_size - else: - raise ValueError( - "chunk_size must be -1 or a positive integer.") - - ordered_names = [ - param_to_state[id(p)].name for p in ordered_params - ] - name_to_state = { - param_to_state[id(p)].name: param_to_state[id(p)] - for p in ordered_params - } - self._parallel_cache[cache_key] = { - 'ordered_names': ordered_names, - 'name_to_state': name_to_state, - 'rank': rank, - 'chunk_size': chunk_size, - } - else: - # Cached path: rebuild param_to_state with current id(p) keys. - cache = self._parallel_cache[cache_key] - rank = cache['rank'] - chunk_size = cache['chunk_size'] - - name_to_param = dict(zip(names, params)) - ordered_params = [name_to_param[n] for n in cache['ordered_names']] - - param_to_state = {} - for p, n in zip(ordered_params, cache['ordered_names']): - cached_state = cache['name_to_state'][n] - param_to_state[id(p)] = _muon_state( - worker_rank=cached_state.worker_rank, - process_group=cached_state.process_group, - rank_indices=cached_state.rank_indices, - rank_numels=cached_state.rank_numels, - name=n, - qk_clip_state=get_qk_clip_info(self.clip_config, n, - qk_logits), - ) - - return ordered_params, param_to_state, rank, chunk_size - - def parallel(self, - names, - params, - group, - lr, - weight_decay, - qk_logits, - prelaunch_gather=None): - """ - Perform a parallel optimization step using Muon. - - Parameters are chunked and each chunk is processed by a - :func:`muon_chunk_pipeline` generator. :func:`run_pipeline` - interleaves multiple chunks so that communication and computation - overlap across chunks (the same overlap previously achieved by the - warmup + main-loop index scheduling). - - If ``prelaunch_gather`` is provided, it is passed to the first - chunk's generator to skip re-launching the already in-flight - A2A gather. - """ - - # Momentum is already applied by _step_muon before this method. - - ordered_params, param_to_state, rank, chunk_size = ( - self._setup_parallel(names, params, group, qk_logits)) - - def pipelines(): - first = True - for start in range(0, len(ordered_params), chunk_size): - chunk = ordered_params[start:start + chunk_size] - if chunk: - kwargs = dict( - params=chunk, - param_to_state=param_to_state, - rank=rank, - ns_steps=group["ns_steps"], - lr=lr, - weight_decay=weight_decay, - none_grad=group["none_grad"], - ) - if first and prelaunch_gather is not None: - kwargs['prelaunch_gather'] = prelaunch_gather - first = False - yield muon_chunk_pipeline(**kwargs) - - with record_function("muon::pipeline"): - run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1) - - def _step_muon(self, group, qk_logits=None): - params = group["params"] - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - # Apply momentum to all params before routing/expansion. - # Batched using _foreach_* ops (compiled, fullgraph=True). - with record_function("muon::momentum"): - active_params = [p for p in params if p.grad is not None] - if active_params: - # Ensure momentum buffers exist (avoid zeros_like when already present). - for p in active_params: - if "momentum_buffer" not in self.state[p]: - self.state[p]["momentum_buffer"] = torch.zeros_like( - p.grad) - - # Extract local tensors for compiled batch function. - local_grads = [ - p.grad._local_tensor - if isinstance(p.grad, DTensor) else p.grad - for p in active_params - ] - local_bufs = [ - self.state[p]["momentum_buffer"]._local_tensor - if isinstance(self.state[p]["momentum_buffer"], DTensor) - else self.state[p]["momentum_buffer"] - for p in active_params - ] - - # Wrap momentum as tensor for torch.compile. - batch_pre_ortho(local_grads, local_bufs, - torch.tensor(momentum), group["nesterov"]) - - # For non-nesterov, the result is the momentum buffer. - if not group["nesterov"]: - for p in active_params: - p.grad = self.state[p]["momentum_buffer"] - - # Identify batched experts for deferred NS. - # Detection is cheap (condition checks only); actual NS compute is - # deferred so it can overlap with the first chunk's A2A gather. - deferred_expert_work = [] - if self.expert_keys: - batched_expert_indices = [] - for i, (n, p) in enumerate(zip(names, params)): - if not (is_expert_param(n, self.expert_keys) - and p.grad is not None): - continue - # Eligible: plain tensor, or DTensor with no non-dim-0 shards. - if isinstance(p.data, DTensor): - has_tp = any( - _is_shard(pl) and pl.dim != 0 for pl in p.placements) - if has_tp: - continue - batched_expert_indices.append(i) - - if batched_expert_indices: - # Save refs for deferred NS; free grads from param list. - for i in batched_expert_indices: - p = params[i] - g = p.grad - local_g = (g._local_tensor - if isinstance(g, DTensor) else g) - local_data = (p.data._local_tensor if isinstance( - p.data, DTensor) else p.data) - deferred_expert_work.append((local_data, local_g)) - p.grad = None - - # Remove batched experts from lists before expansion. - keep = sorted( - set(range(len(params))) - set(batched_expert_indices)) - names = [names[i] for i in keep] - params = [params[i] for i in keep] - - def _run_deferred_expert_ns(): - """Execute deferred batched expert NS.""" - if not deferred_expert_work: - return - with record_function("muon::batched_expert_ns"): - ns_steps = group["ns_steps"] - for local_data, local_g in deferred_expert_work: - u = zeropower_via_newtonschulz5_batched( - local_g.to(COMM_DTYPE), steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:]) - local_data.mul_(1 - lr * weight_decay) - local_data.add_(u, alpha=-adjusted_lr) - - # Expand expert params by splitting on dim 0. - logger.debug("[_step_muon] before expand: %d params, expert_keys=%s", - len(params), self.expert_keys) - if self.expert_keys: - cache_key = tuple(id(p) for p in params) - cache = self._expert_expand_cache.get(cache_key) - - if cache is None: - # Cold path: full expansion + build cache metadata. - exp_names, exp_params = _expand_expert_params( - names, params, self.expert_keys) - - # Build per-expert-group info for hot-path grad updates. - grad_info = [] - exp_idx = 0 - for orig_idx, (n, p) in enumerate(zip(names, params)): - if not is_expert_param(n, self.expert_keys): - exp_idx += 1 - continue - - is_dt = isinstance(p.data, DTensor) - num_experts = (p.to_local() if is_dt else p.data).shape[0] - - # Detect TP mesh from the first expanded expert param. - tp_mesh = None - tp_pls = None - sample = exp_params[exp_idx] - if isinstance(sample.data, DTensor): - tp_mesh = sample.data.device_mesh - tp_pls = list(sample.data.placements) - - grad_info.append((orig_idx, num_experts, exp_idx, is_dt, - tp_mesh, tp_pls)) - exp_idx += num_experts - - self._expert_expand_cache[cache_key] = { - 'names': exp_names, - 'params': exp_params, - 'grad_info': grad_info, - } - names, params = exp_names, exp_params - else: - # Hot path: reuse cached params, only update expert grads. - for (orig_idx, num_experts, exp_start, is_dt, tp_mesh, - tp_pls) in cache['grad_info']: - p = params[orig_idx] - g = p.grad - local_grad = (g.to_local() - if is_dt and isinstance(g, DTensor) else g) - for i in range(num_experts): - expert_p = cache['params'][exp_start + i] - sg = local_grad[i] - if tp_mesh is not None: - expert_p.grad = DTensor.from_local( - sg, device_mesh=tp_mesh, placements=tp_pls) - else: - expert_p.grad = sg - p.grad = None - - names = cache['names'] - params = cache['params'] - else: - names, params = _expand_expert_params(names, params, - self.expert_keys) - logger.debug("[_step_muon] after expand: %d params", len(params)) - - param_dtensors = [] - name_dtensors = [] - - param_tensors = [] - name_tensors = [] - - # distributed_muon is a reference implementation for testing only. - # The parallel pipeline (all2all) path below is the production path. - if self.use_distributed_muon: - _run_deferred_expert_ns() - self.distributed_muon(names=names, - params=params, - group=group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits) - return - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - logger.debug( - "[route] %s → base (DTensor all-Replicate), " - "shape=%s, placements=%s", n, p.shape, p.placements) - param_tensors.append(p) - name_tensors.append(n) - else: - logger.debug( - "[route] %s → parallel (DTensor), shape=%s, " - "placements=%s, mesh=%s", n, p.shape, p.placements, - p.device_mesh.mesh_dim_names) - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - logger.debug("[route] %s → base (plain tensor), shape=%s", n, - p.data.shape) - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError(f"Unsupported parameter type: {type(p.data)}") - - logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, " - f"{len(param_tensors)} Tensors → base") - - def group_dtensors(dtensors, names): - # To support different placements, we group parameters by placements - # and run parallel Muon on each group. - - placement_to_params = defaultdict(lambda: ([], [])) - - assert len(dtensors) == len(names) - for p, n in zip(dtensors, names): - placement_to_params[tuple([p.placements, - p.device_mesh])][0].append(n) - placement_to_params[tuple([p.placements, - p.device_mesh])][1].append(p) - return placement_to_params - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - dtensor_group = group_dtensors(param_dtensors, name_dtensors) - - # Pre-launch the first chunk's A2A gather so that the NCCL - # communication overlaps with the (deferred) batched expert NS - # compute on the default CUDA stream. - prelaunch = None - if deferred_expert_work: - first_names, first_params = next(iter(dtensor_group.values())) - ordered, pts, rnk, csz = self._setup_parallel( - first_names, first_params, group, qk_logits) - first_chunk = ordered[:csz] - if first_chunk: - prelaunch = prelaunch_first_gather(first_chunk, pts, rnk, - group["none_grad"]) - - _run_deferred_expert_ns() - - first_group = True - for _, (names, params) in dtensor_group.items(): - pg = prelaunch if first_group else None - first_group = False - self.parallel( - names, - params, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - prelaunch_gather=pg, - ) - else: - _run_deferred_expert_ns() - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - ) - - def _register_states_for_offload(self): - """Register all optimizer state tensors with the CPU offload pool. - - Called once after the first step when states have been lazily created. - Offloads all param states (momentum buffers for Muon, moment1/moment2 - for AdamW) to free GPU memory between steps. - """ - pool = self._cpu_offload_pool - tracked = 0 - for group in self.param_groups: - for p in group["params"]: - if p not in self.state: - continue - state = self.state[p] - if group.get("use_muon", False): - if "momentum_buffer" in state: - pool.track(state["momentum_buffer"]) - tracked += 1 - else: - if "moment1" in state: - pool.track(state["moment1"]) - if "moment2" in state: - pool.track(state["moment2"]) - tracked += 1 - logger.info("[CPUOffload] Registered %d param states for offload", - tracked) - - @torch.no_grad - def step(self, closure=None, qk_logits=None): - """Perform a single optimization step. - - Args: - closure (Callable, optional): A closure that reevaluates the model - and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). - """ - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # H2D: reload optimizer states from CPU before computation. - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload() - - logger.debug("[Muon.step] expert_keys=%s, %d param groups", - self.expert_keys, len(self.param_groups)) - - for i, group in enumerate(self.param_groups): - if group["use_muon"]: - logger.debug("[Muon.step] group %d: use_muon=True, %d params", - i, len(group["params"])) - self._step_muon(group, qk_logits=qk_logits) - else: - logger.debug( - "[Muon.step] group %d: use_muon=False (AdamW), %d params", - i, len(group["params"])) - step_adamw(self.state, group) - - # D2H: offload optimizer states to CPU after computation. - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - return loss - - # ------------------------------------------------------------------ - # CPU offload public helpers - # ------------------------------------------------------------------ - - def turn_on_cpu_offload(self): - """Enable CPU offload for optimizer states.""" - if self.cpu_offload: - return - logger.info("[Muon] turn_on_cpu_offload") - self.cpu_offload = True - if not self.state: - return - self._cpu_offload_pool = CPUOffloadPool() - self._offload_initialized = False - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_off_cpu_offload(self): - """Disable CPU offload and keep optimizer states resident on GPU.""" - if not self.cpu_offload: - return - logger.info("[Muon] turn_off_cpu_offload") - if self._offload_initialized: - self._cpu_offload_pool.reload() - torch.cuda.current_stream().synchronize() - self._cpu_offload_pool = None - self._offload_initialized = False - self.cpu_offload = False - - # ------------------------------------------------------------------ - # Checkpoint support for cpu_offload - # ------------------------------------------------------------------ - - def state_dict(self) -> dict: - if self.cpu_offload: - raise RuntimeError( - "Muon.state_dict() requires turn_off_cpu_offload() before checkpoint save." - ) - return super().state_dict() - - def load_state_dict(self, state_dict: dict) -> None: - if self.cpu_offload: - raise RuntimeError( - "Muon.load_state_dict() requires turn_off_cpu_offload() before checkpoint load." - ) - super().load_state_dict(state_dict) - - # Invalidate adamw.py's module-level tensor caches so that - # the next step rebuilds them with the newly loaded state tensors. - _placement_cache.clear() - _tensor_cache.clear() diff --git a/build/torch28-cxx11-cu129-x86_64-linux/newton_schulz.py b/build/torch28-cxx11-cu129-x86_64-linux/newton_schulz.py deleted file mode 100644 index d939264b69a34e7a3fa78859f34dc265a1159d59..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/newton_schulz.py +++ /dev/null @@ -1,240 +0,0 @@ -from itertools import repeat -from math import inf, sqrt - -import numpy as np -import torch - -from .matmul_transpose_triton import matmul_transpose_assign - -COMM_DTYPE = torch.bfloat16 -DEFAULT_CHUNK_SIZE_RATIO = 4 - - -def _optimal_quintic(l, u, max_iter=1000): - """ - Use the simplified Remez algorithm to find the optimal odd quintic approximant - to the constant function x -> 1 over the interval [l, u]. - - Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum - approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the - two interior equioscillation nodes q, r until convergence. Returns the - closed-form equioscillating solution when l ≈ u. - - Raises ValueError if any intermediate value (a, b, c, E, q, r) is non-finite - (NaN or inf). Raises RuntimeError if convergence is not reached within - max_iter iterations. - """ - assert 0 <= l <= u - if 1 - 5e-6 <= l / u: - return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5) - q = (3 * l + u) / 4 - r = (l + 3 * u) / 4 - E = inf - for _ in range(max_iter): - old_E = E - LHS = np.array( - [ - [l, l**3, l**5, 1], - [q, q**3, q**5, -1], - [r, r**3, r**5, 1], - [u, u**3, u**5, -1], - ] - ) - a, b, c, E = np.linalg.solve(LHS, np.ones(4)) - if not np.all(np.isfinite([a, b, c, E])): - raise ValueError( - f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}" - ) - q, r = np.sqrt( - (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / (10 * c) - ) - if not np.all(np.isfinite([q, r])): - raise ValueError(f"_optimal_quintic: non-finite node update q={q}, r={r}") - if abs(old_E - E) <= 1e-15: - break - else: - raise RuntimeError( - f"_optimal_quintic: did not converge after {max_iter} iterations" - ) - return float(a), float(b), float(c) - - -def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0): - """ - Compute the Polar Express coefficient series for `num_iters` quintic iterations. - - Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that - compose to map singular values from [l, 1] toward 1. At each step: - 1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion` - prevents near-zero singular values from stalling by raising the effective - lower bound; if it is active (cushion*u > l), the coefficients are - rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u]. - 2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the - last iteration, providing numerical headroom at the cost of a slightly slower - final convergence step. - 3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1). - - Returns a list of (a, b, c) tuples, one per iteration. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - u = 1 - assert 0 <= l <= u - safety_factor = 1 + safety_factor_eps - coefficients = [] - for iter in range(num_iters): - a, b, c = _optimal_quintic(max(l, cushion * u), u) - if cushion * u > l: - pl = a * l + b * l**3 + c * l**5 - pu = a * u + b * u**3 + c * u**5 - rescaler = 2 / (pl + pu) - a *= rescaler - b *= rescaler - c *= rescaler - if iter < num_iters - 1: - a /= safety_factor - b /= safety_factor**3 - c /= safety_factor**5 - coefficients.append((a, b, c)) - l = a * l + b * l**3 + c * l**5 - u = 2 - l - return coefficients - - -# Precomputed Polar Express coefficients (a, b, c) for 10 quintic Newton-Schulz -# iterations. Each tuple is the minimax-optimal (Remez/equioscillation) odd quintic -# approximant to x->1 over the current singular-value interval, computed once at -# import time and reused across all optimizer steps. -# -# Contrast with the former hardcoded NS coefficients (5 fixed tuples): -# - Former: empirically tuned to maximize slope at zero; did not converge -# singular values to 1, yielding US'V^T with S' ~ Uniform(0.5, 1.5) instead -# of the true polar factor UV^T. -# - Polar Express: analytically optimal per step, adapting to the shrinking -# singular-value interval [l, u] as iterations progress; converges all -# singular values to 1, producing the exact polar factor UV^T. -_coeffs_list = _optimal_composition( - l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02 -) - - -# This code is adapted from: -# KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py) -# NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress) -# matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon) -@torch.no_grad() -def _zeropower_via_newtonschulz5(G, steps): - """ - Compute the polar factor of G via the Polar Express method. - - Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c) - are the Polar Express coefficients from `_coeffs_list`. Each step is the - optimal odd quintic approximant to x -> 1 over the current singular-value - interval, minimizing the maximum approximation error (Remez / minimax criterion). - The composition maps singular values from [l, 1] to near 1, producing the - polar factor (orthogonal factor in the polar decomposition G = UP). - - `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2, - cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE - X = G # no manual typecast - - if G.size(0) > G.size(1): - X = X.T - - X = X / (X.norm() + 1e-7) - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - # Perform the NS iterations - for a, b, c in hs: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(0) > G.size(1): - X = X.T - - return X - - -@torch.no_grad() -def _zeropower_via_newtonschulz5_batched(G, steps): - """Batched polar factor computation for 3D (E, out, in) tensors. - - Same algorithm as ``_zeropower_via_newtonschulz5`` but uses - ``torch.bmm`` / ``torch.baddbmm`` instead of the 2D Triton kernel, - processing all E expert matrices in a single batched call. - """ - assert len(G.shape) == 3 - assert G.dtype == COMM_DTYPE - X = G - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - # Per-expert Frobenius norm. - X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) - - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - for a, b, c in hs: - buf1 = torch.bmm(X, X.transpose(-2, -1)) - buf2 = torch.bmm(buf1, buf1.transpose(-2, -1)) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.baddbmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - return X - - -_ns_per_shape: dict[tuple[int, ...], callable] = {} -_use_compile = True - - -def set_ns_compile(enabled: bool): - """Toggle torch.compile for Newton-Schulz iteration.""" - global _use_compile - _use_compile = enabled - - -def zeropower_via_newtonschulz5(G, steps=5): - if not _use_compile: - return _zeropower_via_newtonschulz5(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() - - -def zeropower_via_newtonschulz5_batched(G, steps=5): - """Compile-cached batched Newton-Schulz for 3D expert tensors.""" - if not _use_compile: - return _zeropower_via_newtonschulz5_batched(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile( - _zeropower_via_newtonschulz5_batched, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() diff --git a/build/torch28-cxx11-cu129-x86_64-linux/optimizer/__init__.py b/build/torch28-cxx11-cu129-x86_64-linux/optimizer/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/optimizer/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch28-cxx11-cu129-x86_64-linux/pipeline.py b/build/torch28-cxx11-cu129-x86_64-linux/pipeline.py deleted file mode 100644 index c0c2d515856182d8d15ad27dd4e4e093b29397d6..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/pipeline.py +++ /dev/null @@ -1,468 +0,0 @@ -import logging -from typing import Generator - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -from .core import _muon_state, adjust_lr_for_muon -from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5 -from .qk_clip import compute_scales - -logger = logging.getLogger(__name__) - -# ====================================================================== -# Stage helpers -# ====================================================================== - - -def _launch_gather( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Allocate gather buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_gather``). - gathered_grads: ``{id(p): empty_tensor}`` for owned params, - ``None`` for non-owned. - recv_counts: Per-source-rank element counts. - """ - # Allocate gathered-grad buffers - gathered_grads: dict[int, torch.Tensor | None] = {} - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - gathered_grads[id(p)] = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - else: - gathered_grads[id(p)] = None - - # Build send buffer – batch grad copies via torch.cat - # (1-2 fused kernels vs N individual narrow().copy_() calls). - send_counts = [0] * num_ranks - for p in params: - state = param_to_state[id(p)] - send_counts[state.worker_rank] += state.rank_numels[rank] - - total_send = sum(send_counts) - if total_send > 0: - # Group grad slices by destination rank in a single pass. - dst_to_grads = [[] for _ in range(num_ranks)] - for p in params: - state = param_to_state[id(p)] - n = state.rank_numels[rank] - if n > 0: - g = p.grad.to_local() - dst_to_grads[state.worker_rank].append(g.reshape(-1)) - - # Flatten in dst order and cat once. - all_slices = [] - for dst in range(num_ranks): - all_slices.extend(dst_to_grads[dst]) - send_buf = torch.cat(all_slices) - if send_buf.dtype != COMM_DTYPE: - send_buf = send_buf.to(COMM_DTYPE) - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += state.rank_numels[src] - recv_counts[src] = total - - recv_buf = torch.empty(sum(recv_counts), dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - logger.debug(f"send_buf size: {send_buf.numel()}, " - f"recv_buf size: {recv_buf.numel()}, " - f"recv_counts: {recv_counts}, " - f"send_counts: {send_counts}, " - f"process_group: {str(process_group)}") - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, gathered_grads, recv_counts - - -def _complete_gather( - recv_buf: torch.Tensor, - recv_counts: list[int], - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - param_to_state: dict[int, _muon_state], - rank: int, -) -> None: - """Reconstruct gathered grads from the recv buffer (in-place).""" - off = 0 - for src in range(len(recv_counts)): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - - indices = state.rank_indices[src] - - shard_view = gathered_grads[id(p)][indices] - n = shard_view.numel() - if n == 0: - continue - - sg = recv_buf.narrow(0, off + inner_off, n) - sg = sg.reshape(shard_view.shape) - gathered_grads[id(p)][indices] = sg - - inner_off += n - assert inner_off == block - off += block - - -def _compute_ns( - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - ns_steps: int, -) -> dict[int, torch.Tensor | None]: - """Run Newton-Schulz orthogonalization on owned parameters. - - Returns: - computed_us: ``{id(p): orthogonalized_update}`` for owned params. - """ - computed_us: dict[int, torch.Tensor | None] = {} - for p in owned_params: - u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps) - gathered_grads[id(p)] = None # free gathered grad - computed_us[id(p)] = u - return computed_us - - -def _launch_scatter( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, - computed_us: dict[int, torch.Tensor | None], -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor], list[int]]: - """Allocate scatter buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_scatter``). - scattered_us: Empty dict, populated by ``_complete_scatter`` with - zero-copy views into ``recv_buf``. - recv_counts: Per-source-rank element counts. - """ - # scattered_us is populated by _complete_scatter with zero-copy views - # into recv_buf, avoiding N empty_like allocations + N copy_ calls. - # Pre-seed entries for params whose local shard is empty (rank_numels == 0) - # so _update_params can iterate all params without KeyError. - scattered_us: dict[int, torch.Tensor] = {} - for p in params: - if param_to_state[id(p)].rank_numels[rank] == 0: - scattered_us[id(p)] = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - # Build send buffer – batch via torch.cat - # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls). - send_counts = [0] * num_ranks - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - for dst_rank in range(num_ranks): - send_counts[dst_rank] += state.rank_numels[dst_rank] - - total_send = sum(send_counts) - if total_send > 0: - # Cache u_full conversions to avoid redundant .to() per dst_rank. - u_fulls = {} - for p in owned_params: - u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous() - - # Collect slices in dst order (matches all-to-all send layout). - all_slices = [] - for dst_rank in range(num_ranks): - for p in owned_params: - state = param_to_state[id(p)] - su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten() - if su.numel() > 0: - all_slices.append(su) - - send_buf = torch.cat(all_slices) if all_slices else torch.empty( - 0, dtype=COMM_DTYPE, device="cuda") - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += state.rank_numels[rank] - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, scattered_us, recv_counts - - -def _complete_scatter( - recv_buf: torch.Tensor, - recv_counts: list[int], - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], -) -> None: - """Populate scattered_us with zero-copy views into recv_buf. - - Instead of pre-allocating tensors and copying, we assign views directly - from ``recv_buf``. This eliminates N ``empty_like`` + N ``copy_`` calls. - The underlying storage of ``recv_buf`` is kept alive through the views - until ``scattered_us`` is cleared after ``_update_params``. - """ - off = 0 - for src in range(len(recv_counts)): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = state.rank_numels[rank] - if n == 0: - continue - - scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - - inner_off += n - - assert inner_off == block - off += block - - -def _update_params( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], - lr: float, - weight_decay: float, -) -> None: - """Apply weight decay, Muon update, and optional QK clipping. - - Uses batched ``_foreach_mul_`` for weight decay and batched - ``_foreach_add_`` for the Muon update, grouping parameters by - adjusted_lr to minimize kernel launches while preserving float32 - precision for the alpha scaling. - """ - if not params: - return - - # Batched weight decay: p *= (1 - lr * wd) — single fused kernel. - p_locals = [p._local_tensor for p in params] - torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay) - - # Group params by adjusted_lr so _foreach_add_ can use a single - # alpha per group (preserves float32 precision for alpha scaling). - lr_groups: dict[float, tuple[list, list]] = {} - for p in params: - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - if adjusted_lr not in lr_groups: - lr_groups[adjusted_lr] = ([], []) - lr_groups[adjusted_lr][0].append(p._local_tensor) - lr_groups[adjusted_lr][1].append(scattered_us[id(p)]) - - for adjusted_lr, (p_group, u_group) in lr_groups.items(): - torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr) - - # QK clipping – applied directly on the local tensor to - # avoid DTensor sharding-propagation issues with _StridedShard. - for p in params: - state = param_to_state[id(p)] - if state.qk_clip_state is None: - continue - scales_full = compute_scales(p, state.qk_clip_state) - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = state.rank_indices[rank][0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - -# ====================================================================== -# Pre-launch helper for overlapping first chunk's gather with other work. -# ====================================================================== - - -@torch.no_grad() -def prelaunch_first_gather( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - none_grad: bool, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Launch the first chunk's A2A gather early for overlap with other compute. - - Call this *before* expensive GPU work (e.g. batched expert NS) so that - the NCCL all-to-all runs concurrently on the NCCL stream while the - default stream executes compute. - - Returns the same 4-tuple that ``_launch_gather`` produces, which should - be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - with record_function("muon::prelaunch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - return work, recv_buf, gathered_grads, recv_counts - - -# ====================================================================== -# Main generator – thin orchestrator that wires stages together. -# ====================================================================== - - -@torch.no_grad() -def muon_chunk_pipeline( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - ns_steps: int, - lr: float, - weight_decay: float, - none_grad: bool, - prelaunch_gather: tuple | None = None, -) -> Generator[None, None, None]: - """Process one chunk of parameters through the full Muon pipeline. - - Stages: gather -> compute (Newton-Schulz) -> scatter -> update. - - Each ``yield`` lets :func:`run_pipeline` interleave other chunks so - that communication and computation overlap across chunks. Async - communication is launched via ``async_op=True`` and completed after - the yield with ``work.wait()``. - - Overlap happens because :func:`run_pipeline` admits one new chunk - per iteration (staggered admission). While chunk *N* does NS - compute on the default CUDA stream, chunk *N+1*'s async all-to-all - runs concurrently on the NCCL stream — no separate ``comm_stream`` - is required. - - If ``prelaunch_gather`` is provided, the gather was already launched - by :func:`prelaunch_first_gather` and we skip launching it again. - - Yields exactly **2** times: - - 1. After launching async all-to-all gather (or immediately if pre-launched). - 2. After launching async all-to-all scatter. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - if prelaunch_gather is not None: - # Gather was pre-launched; none_grad already handled by caller. - work, recv_buf, gathered_grads, recv_counts = prelaunch_gather - else: - # Normal path: launch async gather. - with record_function("muon::launch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - yield # --- YIELD 1: other chunks can launch their gather --- - - with record_function("muon::wait_gather"): - work.wait() - _complete_gather(recv_buf, recv_counts, owned_params, gathered_grads, - param_to_state, rank) - del recv_buf - - # Stage 3: Newton-Schulz orthogonalization. - with record_function("muon::newton_schulz"): - computed_us = _compute_ns(owned_params, gathered_grads, ns_steps) - gathered_grads.clear() - - # Stages 4-5: launch async scatter. - with record_function("muon::launch_scatter"): - work, recv_buf, scattered_us, recv_counts = _launch_scatter( - params, owned_params, param_to_state, rank, num_ranks, - process_group, computed_us) - computed_us.clear() - - yield # --- YIELD 2: other chunks can launch their scatter --- - - with record_function("muon::wait_scatter"): - work.wait() - _complete_scatter(recv_buf, recv_counts, params, param_to_state, rank, - scattered_us) - del recv_buf - - # Stage 6: apply parameter updates. - with record_function("muon::update_params"): - _update_params(params, param_to_state, rank, scattered_us, lr, - weight_decay) - scattered_us.clear() diff --git a/build/torch28-cxx11-cu129-x86_64-linux/qk_clip.py b/build/torch28-cxx11-cu129-x86_64-linux/qk_clip.py deleted file mode 100644 index 2aba711b3004b7f09e7141da7ef834bd61cc2430..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/qk_clip.py +++ /dev/null @@ -1,198 +0,0 @@ -import logging -import math -from dataclasses import dataclass - -import torch -from torch.distributed.tensor import DTensor - -from .core import normalize_fqn - -logger = logging.getLogger(__name__) - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - and return (kind, layer_index). - - Supported kinds: - MHA/GQA: 'wq', 'wk', 'q_proj', 'k_proj' - MLA: 'wq_b' (Q up-proj), 'wkv_b' (KV up-proj) - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.1.attn.wq_b.weight' -> ('wq_b', 1) - 'model.0.attn.wkv_b.weight' -> ('wkv_b', 0) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = normalize_fqn(name).split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: str | None # 'wq'/'q_proj'/'wq_b' or 'wk'/'k_proj'/'wkv_b' or None - indices: list[int] # which heads to consider for clipping - head_dim: int # from config (qk_head_dim for MLA wq_b) - threshold: float # from config - logit: torch.Tensor | None - - # MLA-specific fields - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 - - -def get_qk_clip_info(clip_config, n, qk_logits): - """Extract QK clipping info for a named parameter. - - Args: - clip_config: QK clipping configuration dict (or None). - MHA/GQA keys: head_dim, threshold, q_indices, k_indices - MLA extra keys: is_mla=True, qk_nope_head_dim, qk_rope_head_dim, v_head_dim - n: Parameter name string. - qk_logits: Dict mapping layer indices to logit tensors (or None). - - Returns: - QKClipInfo instance with clipping configuration for this parameter. - """ - if clip_config is None: - return None - - head_dim = clip_config.get('head_dim') - threshold = clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - # In TP settings, qk_logits may be DTensor - # We convert it to full tensor here for simplicity - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - -def compute_scales(p, qk_clip_state): - """Compute per-head scaling factors for QK clipping. - - Returns scales tensor (√γ per head) if any head exceeds threshold, else None. - For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim. - """ - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - # Check if any head exceeds threshold before allocating. - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - - if not head_scales: - return None - - # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows - if qk_clip_state.is_mla and kind == 'wkv_b': - effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim - else: - effective_head_dim = head_dim - - H_global = p.shape[0] // effective_head_dim - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale - return scales_full - - -def qk_clip(p, scales, info): - """Apply per-head scaling to a Q/K projection weight matrix. - - Args: - p: Parameter (nn.Parameter or raw tensor). - scales: [n_heads] tensor, each element = √γ_h. - info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions. - - MLA sub-region scaling per Algorithm 1 (MuonClip): - wq_b: q_nope rows → √γ, q_pe rows → γ - wkv_b: k_nope rows → √γ, v rows → unchanged - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not info.is_mla: - # MHA/GQA: uniform √γ applied to all rows in each head - W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: vectorized sub-region scaling within each head - if info.kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_head_dim = qk_nope + info.qk_rope_head_dim - W_3d = W.view(-1, qk_head_dim, W.shape[1]) # [H, qk_head_dim, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # q_nope → √γ - W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1, - 1)) # q_pe → γ - - elif info.kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - W_3d = W.view(-1, kv_stride, W.shape[1]) # [H, kv_stride, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # k_nope → √γ - # v rows: not touched (k_R shared rotary unchanged) diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/__init__.py b/build/torch28-cxx11-rocm63-x86_64-linux/__init__.py deleted file mode 100644 index 239c7a65f8293e7d0df28f05fce645af56d628c0..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .muon import Muon - -__all__ = [ - "Muon", -] diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/_ops.py b/build/torch28-cxx11-rocm63-x86_64-linux/_ops.py deleted file mode 100644 index 034bff088659b5df6f6d401feb18c89dc5f33b29..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _optimizer_8d53b78_dirty -ops = torch.ops._optimizer_8d53b78_dirty - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_optimizer_8d53b78_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so b/build/torch28-cxx11-rocm63-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so deleted file mode 100755 index 0df93dd2ea24c80cbed1f029804c4e4e480a140e..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:bcf5b8838dfaf6e81fdbd52ff4638ca76abaa678f7c2cbd81cf03dc72f9cd5d2 -size 1865080 diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/adamw.py b/build/torch28-cxx11-rocm63-x86_64-linux/adamw.py deleted file mode 100644 index b5a95816a9f5b9e1889eaadae65373bfbced809a..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/adamw.py +++ /dev/null @@ -1,271 +0,0 @@ -import logging -from collections import defaultdict -from typing import cast - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -def fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: float | torch.Tensor, - weight_decay: float, - eps: float, - maximize: bool, -) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: dict | None = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) - - -def _to_local(t): - """Unwrap DTensor to local tensor for fused ops.""" - return t._local_tensor if isinstance(t, DTensor) else t - - -# --------------------------------------------------------------------------- -# Caches for eliminating per-step Python overhead. -# -# Placement grouping and tensor list assembly are identical every step -# (params don't change placement, moment/step tensors are the same objects -# after initialisation). We cache them keyed by id() of the param list -# stored in param_groups (stable across steps). -# -# Only gradients change each step and must be collected fresh. -# --------------------------------------------------------------------------- - -# id(group["params"]) → dict[placement_key, list[param]] -_placement_cache: dict[int, dict[tuple, list]] = {} - -# id(placement_group_list) → (params_local, moment1, moment2, state_steps) -_tensor_cache: dict[int, tuple[list, list, list, list]] = {} - - -def _step_adamw_params_slow(optimizer_state, params, group): - """Uncached fallback for the rare case where some params lack grads.""" - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - g = p.grad - if g is None: - continue - state = optimizer_state[p] - params_with_grads.append(_to_local(p)) - grads.append(_to_local(g)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - if not params_with_grads: - return - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw_params(optimizer_state, params, group): - """Run fused AdamW on a list of parameters sharing the same placement. - - After the first call, cached tensor lists (params_local, moment1, - moment2, state_steps) are reused — only gradients are collected fresh. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - params: List of parameters to update. - group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay. - """ - # Collect grads — the only thing that changes each step. - with record_function("adamw::collect_grads"): - grads = [] - for p in params: - g = p.grad - if g is None: - # Rare: fall back to slow path that filters per-param. - _step_adamw_params_slow(optimizer_state, params, group) - return - grads.append(_to_local(g)) - - tensor_key = id(params) - if tensor_key not in _tensor_cache: - with record_function("adamw::init_tensor_cache"): - params_local = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - state = optimizer_state[p] - params_local.append(_to_local(p)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(p.grad) - state["moment2"] = torch.zeros_like(p.grad) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - _tensor_cache[tensor_key] = (params_local, moment1, moment2, - state_steps) - - params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key] - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - with record_function("adamw::fused_adamw"): - fused_adamw( - params_local, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw(optimizer_state, group): - """Dispatch AdamW step, grouping parameters by type and placement. - - Placement grouping is cached after the first call since params never - change their placement between steps. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - group: Parameter group dict. - """ - params = group["params"] - placement_key = id(params) - - if placement_key not in _placement_cache: - with record_function("adamw::group_by_placement"): - placement_to_params: dict[tuple, - list[torch.Tensor]] = defaultdict(list) - for p in params: - match p: - case DTensor(): - logger.debug( - "[AdamW] DTensor param: shape=%s, placements=%s, " - "mesh=%s, grad=%s", p.shape, p.placements, - p.device_mesh.mesh_dim_names, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple( - [p.placements, p.device_mesh])].append(p) - case torch.Tensor(): - logger.debug( - "[AdamW] plain param: shape=%s, grad=%s", p.shape, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple([torch.Tensor, - None])].append(p) - - logger.debug("[AdamW] %d placement groups, %d total params", - len(placement_to_params), len(params)) - - _placement_cache[placement_key] = dict(placement_to_params) - - for group_params in _placement_cache[placement_key].values(): - step_adamw_params(optimizer_state, group_params, group) diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/async_utils.py b/build/torch28-cxx11-rocm63-x86_64-linux/async_utils.py deleted file mode 100644 index a45c530ac9cad88e3555ec1047a6aa59f225347e..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/async_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import logging -from typing import Generator - -logger = logging.getLogger(__name__) - - -class _Task: - """Internal: wraps a generator, advances one yield at a time.""" - - def __init__(self, generator: Generator[None, None, None], index: int): - self._generator = generator - self._index = index - self._steps_completed = 0 - self.step() # run to first yield - - def step(self) -> bool: - try: - next(self._generator) - self._steps_completed += 1 - logger.debug("pipeline[%d] completed stage %d", self._index, - self._steps_completed) - return True - except StopIteration: - logger.debug("pipeline[%d] finished after %d stages", self._index, - self._steps_completed) - return False - - def close(self): - self._generator.close() - - -def run_pipeline( - pipelines: Generator[Generator[None, None, None], None, None], - max_concurrent: int, -) -> None: - """Run generator-based pipelines with bounded concurrency. - - Each pipeline is a generator that yields at stage boundaries. - The runtime interleaves pipelines so communication and computation - overlap across chunks. - """ - if max_concurrent <= 0: - raise ValueError(f"max_concurrent must be > 0, got {max_concurrent}") - - have_new = True - task_index = 0 - previous_tasks: list[_Task] = [] - - try: - while have_new or previous_tasks: - running_tasks: list[_Task] = [] - - # Admit one new pipeline per iteration (staggered admission). - # Admitting one at a time ensures that while chunk N does NS - # compute on the default stream, chunk N+1's NCCL all-to-all - # runs concurrently on the NCCL stream — creating real - # communication/computation overlap on the GPU. - if have_new and len(previous_tasks) < max_concurrent: - try: - gen = next(pipelines) - task = _Task(gen, task_index) - task_index += 1 - running_tasks.append(task) - except StopIteration: - have_new = False - - # Advance every previously-yielded task by one step. - for task in previous_tasks: - if task.step(): - running_tasks.append(task) - - previous_tasks = running_tasks - except BaseException: - # Clean up all in-flight generators to release GPU resources. - for task in previous_tasks: - task.close() - raise diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/core.py b/build/torch28-cxx11-rocm63-x86_64-linux/core.py deleted file mode 100644 index c69d515afef305ad0ed66374095fa2d2468d99cc..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/core.py +++ /dev/null @@ -1,219 +0,0 @@ -import logging -import math -from dataclasses import dataclass -from typing import List - -import torch -from torch.distributed import ProcessGroup -from torch.distributed.tensor import DTensor - -# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into -# parameter FQNs. Activation checkpointing similarly inserts -# "_checkpoint_wrapped_module". Strip these so name-based matching (skip_keys, -# expert_keys, QK layer parsing) works regardless of wrapper nesting. -_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) - -logger = logging.getLogger(__name__) - - -def normalize_fqn(name: str) -> str: - """Strip torch.compile / checkpoint wrapper components from a parameter FQN.""" - return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS) - - -@dataclass -class _muon_state: - worker_rank: int - process_group: ProcessGroup - rank_indices: dict[int, tuple] # local_rank -> per-dim indices - rank_numels: dict[int, int] # local_rank -> numel - name: str - qk_clip_state: torch.Tensor | None = None - - -def _batch_momentum( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update (no nesterov).""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - - -def _batch_momentum_nesterov( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update with nesterov correction.""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - nesterov_terms = torch._foreach_mul(momentum_bufs, momentum) - torch._foreach_add_(grads, nesterov_terms) - - -_compiled_momentum: dict[bool, callable] = {} -_use_momentum_compile = True - - -def set_momentum_compile(enabled: bool): - """Toggle torch.compile for batched momentum.""" - global _use_momentum_compile - _use_momentum_compile = enabled - - -def batch_pre_ortho( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, - nesterov: bool, -) -> None: - """Batched momentum update on lists of plain tensors. - - Mirrors dion's ``muon_update_pre_orthogonalize``. - Inputs must be plain CUDA tensors (not DTensor). - Modifies ``momentum_bufs`` and (for nesterov) ``grads`` in-place. - - When compile is enabled, uses separately compiled functions for - nesterov=True/False to avoid graph breaks from the branch. - """ - fn = _batch_momentum_nesterov if nesterov else _batch_momentum - if _use_momentum_compile: - if nesterov not in _compiled_momentum: - _compiled_momentum[nesterov] = torch.compile(fn) - fn = _compiled_momentum[nesterov] - fn(grads, momentum_bufs, momentum) - - -def _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay): - """Weight-decay + update on plain tensors. - - Not compiled: per-param @torch.compile caused ~0.25ms TorchDynamo cache - lookup per call × 256+ params = massive overhead. The pipeline path uses - batched _foreach_* ops instead; this function remains for base() and - distributed_muon(). - """ - p_data.mul_(1 - lr * weight_decay) - p_data.add_(u_data, alpha=-adjusted_lr) - - -def update_p(p, u, lr, adjusted_lr, weight_decay): - """Apply weight decay and orthogonalized update to parameter. - - Args: - p: Parameter (torch.nn.Parameter or DTensor). - u: Orthogonalized update tensor. - lr: Base learning rate. - adjusted_lr: Size-adjusted learning rate. - weight_decay: Weight decay coefficient. - """ - # Unwrap Parameter -> underlying data tensor. - p_data = p.data if isinstance(p, torch.nn.Parameter) else p - # Unwrap DTensor -> local CUDA tensor for compiled kernel. - if isinstance(p_data, DTensor): - p_data = p_data._local_tensor - u_data = u._local_tensor if isinstance(u, DTensor) else u - _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay) - - -def adjust_lr_for_muon(lr, param_shape): - """Scale learning rate based on parameter matrix dimensions. - - Args: - lr: Base learning rate. - param_shape: Shape of the parameter tensor. - - Returns: - Adjusted learning rate. - """ - A, B = param_shape[:2] - # We adjust the learning rate and weight decay based on the size of the parameter matrix - # as described in the paper - adjusted_ratio = 0.2 * math.sqrt(max(A, B)) - adjusted_lr = lr * adjusted_ratio - return adjusted_lr - - -def _match_key(parts, key): - """Check if key matches as contiguous components in parts. - - Single-component keys (e.g. "experts") match any single component. - Multi-component keys (e.g. "experts.w1") match as a contiguous subsequence. - """ - key_parts = key.split(".") - key_len = len(key_parts) - if key_len == 1: - return key in parts - return any(parts[i:i + key_len] == key_parts - for i in range(len(parts) - key_len + 1)) - - -def is_expert_param(name, expert_keys): - """Check if a parameter name matches any expert key (component-level).""" - if not expert_keys: - return False - parts = normalize_fqn(name).split(".") - return any(_match_key(parts, key) for key in expert_keys) - - -def default_is_muon(name, x, expert_keys=None): - normalized = normalize_fqn(name) - parts = normalized.split(".") - skip_keys = [ - "embed_tokens", - "lm_head", - "tok_embeddings", - "output", - "mhc_attn", - "mhc_ffn", - "lambda_proj", - ] - if any(key in parts for key in skip_keys): - logger.info( - "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d", - normalized, name, x.ndim) - return False - effective_ndim = x.ndim - is_expert = is_expert_param(name, expert_keys) - if is_expert: - effective_ndim -= 1 - result = effective_ndim >= 2 - logger.info( - "[is_muon] %s (orig: %s): ndim=%d, expert=%s, effective_ndim=%d → %s", - normalized, name, x.ndim, is_expert, effective_ndim, - "Muon" if result else "AdamW") - return result - - -def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None): - if is_muon_func is None: - is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys) - - muon_params, muon_names = [], [] - non_muon_params, non_muon_names = [], [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - non_muon_names.append(n) - - logger.info("[param_groups] expert_keys=%s, Muon=%d, AdamW=%d", - expert_keys, len(muon_names), len(non_muon_names)) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/cpu_offload.py b/build/torch28-cxx11-rocm63-x86_64-linux/cpu_offload.py deleted file mode 100644 index fb5e69154a1d4a6c884491413a37a9acf0f66c80..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/cpu_offload.py +++ /dev/null @@ -1,206 +0,0 @@ -"""CPU offloading for optimizer states. - -Manages a pinned CPU memory pool and async CUDA streams to offload -optimizer state tensors (momentum buffers, Adam moments) to CPU between -optimizer steps, freeing GPU memory. - -All tracked tensors are packed into a single flat pinned CPU buffer -(per dtype). D2H and H2D copies are performed per-tensor directly -between individual GPU tensors and their slice of the CPU flat buffer -— no GPU staging buffer is allocated, so there is **no temporary GPU -memory spike** during offload or reload. - -Individual tensor storages are freed after offload via -``untyped_storage().resize_(0)``, preserving tensor identity so -downstream caches remain valid. -""" - -import logging -from collections import defaultdict - -import torch -from torch.distributed.tensor import DTensor - -logger = logging.getLogger(__name__) - - -class CPUOffloadPool: - """Pinned CPU memory pool for async optimizer state offloading. - - Tracked tensors are grouped by dtype. Each group gets a single flat - pinned CPU buffer. D2H / H2D copies are per-tensor (into slices of - the flat buffer) to avoid allocating a GPU staging buffer. - """ - - def __init__(self): - self._managed: list[torch.Tensor] = [] - self._storage_nbytes: dict[int, int] = {} # id(t) → bytes - - # Per-dtype group: populated on first offload. - # dtype → dict with keys: - # "indices" : list[int] managed-list indices - # "offsets" : list[tuple[int,int]] (start, numel) in flat buf - # "total" : int total numel - # "cpu_flat" : Tensor pinned CPU buffer - self._groups: dict[torch.dtype, dict] = {} - - self._offload_stream: torch.cuda.Stream | None = None - self._device: torch.device | None = None - self._initialized: bool = False - self._logged: bool = False - - # ------------------------------------------------------------------ - @staticmethod - def _local(t: torch.Tensor) -> torch.Tensor: - """Unwrap DTensor to its local CUDA tensor.""" - return t._local_tensor if isinstance(t, DTensor) else t - - def _ensure_stream(self): - if self._offload_stream is None: - self._offload_stream = torch.cuda.Stream(device=self._device) - - # ------------------------------------------------------------------ - def track(self, tensor: torch.Tensor): - """Register a GPU tensor for CPU offloading. Idempotent.""" - tid = id(tensor) - if tid in self._storage_nbytes: - return - local = self._local(tensor) - if self._device is None: - self._device = local.device - storage = local.untyped_storage() - # Skip tensors with empty storage (e.g. empty FSDP shards) - if storage.size() == 0: - return - self._storage_nbytes[tid] = storage.size() - self._managed.append(tensor) - - # ------------------------------------------------------------------ - def _init_buffers(self): - """Build per-dtype flat buffers on first offload.""" - # Group managed tensors by dtype. - dtype_map: dict[torch.dtype, list[tuple[int, int]]] = defaultdict(list) - for idx, t in enumerate(self._managed): - local = self._local(t) - dtype_map[local.dtype].append((idx, local.numel())) - - total_cpu_bytes = 0 - for dtype, entries in dtype_map.items(): - offsets: list[tuple[int, int]] = [] - indices: list[int] = [] - off = 0 - for idx, n in entries: - indices.append(idx) - offsets.append((off, n)) - off += n - cpu_flat = torch.empty(off, dtype=dtype, device="cpu", pin_memory=True) - self._groups[dtype] = { - "indices": indices, - "offsets": offsets, - "total": off, - "cpu_flat": cpu_flat, - } - total_cpu_bytes += off * cpu_flat.element_size() - - self._initialized = True - logger.info( - "[CPUOffload] Pool initialized: %d tensors, %d dtype group(s), " - "%.2f MB pinned CPU memory", - len(self._managed), - len(self._groups), - total_cpu_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def offload(self): - """Per-tensor async D2H into CPU flat buffer, then free GPU storage.""" - if not self._managed: - return - if not self._initialized: - self._init_buffers() - self._ensure_stream() - - # Offload stream waits for compute to finish. - compute_event = torch.cuda.current_stream(self._device).record_event() - self._offload_stream.wait_event(compute_event) - - offloaded_bytes = 0 - - # Per-tensor D2H copies directly into CPU flat buffer slices. - # No GPU staging buffer → no temporary GPU memory spike. - with torch.cuda.stream(self._offload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - cpu_flat[off : off + n].copy_(local.reshape(-1), non_blocking=True) - - offloaded_bytes += grp["total"] * cpu_flat.element_size() - - # Wait for all D2H copies to land, then free GPU storage. - self._offload_stream.synchronize() - for t in self._managed: - storage = self._local(t).untyped_storage() - if storage.size() != 0: - storage.resize_(0) - else: - raise RuntimeError( - f"Tensor storage is already freed (size=0) before offload. " - f"This indicates a double-free or external interference. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - - if not self._logged: - logger.info( - "[CPUOffload] Offloaded %.2f MB (GPU → CPU)", - offloaded_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def reload(self): - """Per-tensor H2D from CPU flat buffer on the default stream. - - Runs on the current (default) CUDA stream to avoid stream - interaction issues with the parallel Muon pipeline. Since - pinned CPU memory is the source, the copies overlap with - GPU idle time between steps. - """ - if not self._managed or not self._initialized: - return - - reloaded_bytes = 0 - - # Re-allocate all GPU storages first. - for t in self._managed: - local = self._local(t) - storage = local.untyped_storage() - if storage.size() != 0: - raise RuntimeError( - f"Storage should have been freed (size=0) before reload, " - f"but got size={storage.size()}. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - storage.resize_(self._storage_nbytes[id(t)]) - - # Per-tensor H2D copies from CPU flat buffer slices. - # non_blocking=True with pinned source allows DMA overlap. - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off : off + n], non_blocking=True) - - reloaded_bytes += grp["total"] * cpu_flat.element_size() - - if not self._logged: - logger.info( - "[CPUOffload] Reloaded %.2f MB (CPU → GPU)", reloaded_bytes / (1024**2) - ) diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/distributed/utils.py b/build/torch28-cxx11-rocm63-x86_64-linux/distributed/utils.py deleted file mode 100644 index 683096bcd50489e7a09959753fdbc29ba0f3b6a6..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/distributed/utils.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import (Placement, Shard, - _StridedShard) - - -def _is_shard(placement: Placement) -> bool: - """Check if a placement is a shard type (Shard or _StridedShard). - - In PyTorch 2.10+, _StridedShard no longer inherits from Shard, so - ``placement.is_shard()`` returns False for _StridedShard. This helper - handles both old and new hierarchies. - """ - return isinstance(placement, (Shard, _StridedShard)) - - -def get_slices_of_dtensor( - target: DTensor | torch.Tensor, - local_rank: int, - shard_mesh: DeviceMesh, - shard_placements: tuple[Placement], -) -> tuple[slice | torch.Tensor, ...]: - """ - Get per-dimension indices for a given rank's shard of the target tensor. - - Uses ``Shard.local_shard_size_and_offset`` and - ``_StridedShard.local_shard_size_and_offset`` for correct handling of - both contiguous and strided (non-contiguous) sharding. - - Args: - target (DTensor | torch.Tensor): The target tensor (for its shape). - local_rank (int): The local rank within the shard group. - shard_mesh (DeviceMesh): The shard mesh (only shard dimensions). - shard_placements (tuple[Placement]): The shard placements. - - Returns: - A tuple of indices (one per tensor dim). Each element is either: - - A ``slice`` (for contiguous or unsharded dims) - - A 1-D ``torch.LongTensor`` of indices (for strided sharding) - """ - - # find the global rank of the local rank in the shard mesh - rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank] - - rank_coords = (shard_mesh.mesh == rank).nonzero() - - assert len(rank_coords) == 1 - rank_coords = tuple(rank_coords[0].tolist()) - - assert len(rank_coords) == len(shard_placements) - - # Track per-shard-dim indices. - # None means "not yet sharded on this dim". - dim_indices: dict[int, torch.Tensor] = {} - - # Caution: Assuming replicate-to-shard of the shard mesh goes with - # left-to-right sharding. This is ensured by the sorting logic of - # construct_shard_mesh function. - for mesh_dim_idx, (rank_coord, placement) in enumerate( - zip(rank_coords, shard_placements)): - assert _is_shard(placement) - - num_chunks = shard_mesh.mesh.shape[mesh_dim_idx] - shard_dim = placement.dim - - # Current effective size on this dim (may already be sub-sharded) - if shard_dim in dim_indices: - curr_size = len(dim_indices[shard_dim]) - else: - curr_size = target.size()[shard_dim] - - # Compute indices for this level of sharding - if isinstance(placement, _StridedShard): - _shard_size, offsets = _StridedShard.local_shard_size_and_offset( - placement, - curr_size, - num_chunks, - rank_coord, - return_first_offset=False) - new_indices = torch.tensor(offsets, dtype=torch.long) - else: - shard_size, offset = Shard.local_shard_size_and_offset( - curr_size, num_chunks, rank_coord) - new_indices = torch.arange(offset, - offset + shard_size, - dtype=torch.long) - - # Compose with previous indices on this dim - if shard_dim in dim_indices: - dim_indices[shard_dim] = dim_indices[shard_dim][new_indices] - else: - dim_indices[shard_dim] = new_indices - - # Build result tuple - result: list[slice | torch.Tensor] = [] - for d in range(len(target.size())): - if d not in dim_indices: - result.append(slice(None)) - else: - indices = dim_indices[d] - # Convert contiguous indices to slice for efficiency - if len(indices) > 0: - start = indices[0].item() - expected = torch.arange(start, - start + len(indices), - dtype=torch.long) - if torch.equal(indices, expected): - result.append(slice(start, start + len(indices))) - else: - result.append(indices) - else: - result.append(slice(0, 0)) - - return tuple(result) - - -_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, - ProcessGroup]] = dict() - - -def construct_shard_mesh( - placements: tuple[Placement], - mesh: DeviceMesh, -) -> tuple[DeviceMesh, ProcessGroup, tuple[Placement, ...]]: - """Construct shard sub-mesh and ProcessGroup for all-to-all communication. - - Given a DTensor's placements and device mesh, extracts the "shard group" - — the set of ranks that together hold all shards of the same replica — - and creates a ProcessGroup for all-to-all among them. - - Steps: - 1. Sort placements: Replicate first, then Shard by (dim, granularity). - 2. Permute the mesh tensor to match the sorted order. - 3. Collapse Replicate dims → list of shard sub-meshes (one per replica). - 4. Create/retrieve a cached ProcessGroup for the current rank's sub-mesh. - - Example — 8 GPUs, mesh shape (2, 2, 2), - placements ``[Shard(0), Replicate, _StridedShard(0)]``:: - - Step 1 — Sort: [Replicate, _StridedShard(0), Shard(0)] - Permutation: [1, 2, 0] - - Step 2 — Permute mesh dims by [1, 2, 0]: - Original: Permuted: - [[[0,1],[2,3]], [[[0,2],[1,3]], - [[4,5],[6,7]]] [[4,6],[5,7]]] - - Step 3 — Unbind replicate dim (dim 0), giving 2 shard sub-meshes: - sub-mesh 0 = [[0,2],[1,3]] (replica group 0) - sub-mesh 1 = [[4,6],[5,7]] (replica group 1) - shard_placements = (_StridedShard(0), Shard(0)) - - Step 4 — Rank 0 → ProcessGroup([0,1,4,5]) - Rank 2 → ProcessGroup([2,3,6,7]) - - Returns: - ``(shard_mesh, process_group, shard_placements)`` - """ - my_rank = dist.get_rank() - assert mesh.mesh.device.type == 'cpu' - - # -- Fast path: 1D all-shard mesh → reuse existing PG. ---------------- - # Reuses the mesh's existing ProcessGroup directly, avoiding the - # overhead of dist.new_group(). The standard path below also handles - # subset calls safely via use_local_synchronization=True, but this - # fast path is still beneficial for the common 1D shard case. - if mesh.ndim == 1 and len(placements) == 1 and _is_shard(placements[0]): - key = (*mesh.mesh.shape, *mesh.mesh.flatten().tolist()) - if key not in _ranks_to_dist_cache: - _ranks_to_dist_cache[key] = (mesh, mesh.get_group()) - return (*_ranks_to_dist_cache[key], tuple(placements)) - - mesh_tensor = mesh.mesh.clone() - - # -- Step 1: Sort placements (Replicate first, then Shard by dim). ------ - # _StridedShard comes BEFORE regular Shard on the same dim so that - # get_slices_of_dtensor applies the outer sharding first, matching - # DTensor's left-to-right (outer-to-inner) composition order. - def _sort_key(item): - index, placement = item - assert not placement.is_partial(), "Partial placement not supported" - if placement.is_replicate(): - return (-1, 0, index) - assert _is_shard(placement), f"Unsupported: {type(placement)}" - split = (-1 / placement.split_factor if isinstance( - placement, _StridedShard) else 0) - return (placement.dim, split, index) - - indexed = sorted(enumerate(placements), key=_sort_key) - perm, sorted_placements = zip(*indexed) - - # -- Step 2: Permute mesh to match sorted placement order. -------------- - sorted_mesh = mesh_tensor.permute(perm) - - # -- Step 3: Collapse replicate dims → list of shard sub-meshes. -------- - # E.g. mesh (2, 3, 4, 4) with [R, R, S(0), S(1)] → 6 sub-meshes of (4, 4) - num_rep = sum(1 for p in sorted_placements if p.is_replicate()) - if num_rep > 0: - if num_rep > 1: - sorted_mesh = sorted_mesh.flatten(0, num_rep - 1) - shard_meshes = list(torch.unbind(sorted_mesh, dim=0)) - else: - shard_meshes = [sorted_mesh] - shard_placements = sorted_placements[num_rep:] - assert len(shard_placements) == len(set(shard_placements)) - - # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. -- - # Each rank only creates the group it belongs to, using - # use_local_synchronization=True so that only group members need to - # coordinate. This avoids deadlocks when different PP stages call - # construct_shard_mesh for different parameters. - def _cache_key(t: torch.Tensor) -> tuple: - return (*t.shape, *t.flatten().tolist()) - - my_key = None - for sm in shard_meshes: - if (my_rank == sm).any().item(): - key = _cache_key(sm) - assert my_key is None, "Rank appears in multiple shard groups" - my_key = key - if key not in _ranks_to_dist_cache: - pg = dist.new_group(sm.flatten().tolist(), - use_local_synchronization=True) - _ranks_to_dist_cache[key] = ( - DeviceMesh(device_type="cuda", mesh=sm), - pg, - ) - - return (*_ranks_to_dist_cache[my_key], shard_placements) diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/matmul_transpose_triton.py b/build/torch28-cxx11-rocm63-x86_64-linux/matmul_transpose_triton.py deleted file mode 100644 index 792de23d82c3fb45fe33d397ab9b76a0787259d0..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/matmul_transpose_triton.py +++ /dev/null @@ -1,122 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], - restore_value=['y'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -@torch.library.custom_op("muon::matmul_transpose_assign", - mutates_args=("d_out", )) -def matmul_transpose_assign(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """Compute d_out = d_in @ d_in.T using an optimized Triton kernel.""" - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -@matmul_transpose_assign.register_fake -def _(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """FakeTensor impl: d_out is already allocated, mutation is declared.""" - pass diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/metadata.json b/build/torch28-cxx11-rocm63-x86_64-linux/metadata.json deleted file mode 100644 index c55a35717622f1dd5c8ba376ea3a814cbcc10d78..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/metadata.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/muon.py b/build/torch28-cxx11-rocm63-x86_64-linux/muon.py deleted file mode 100644 index 14c0e22471fa6d47a51ed95e0e0c341dc18d5194..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/muon.py +++ /dev/null @@ -1,1068 +0,0 @@ -import logging -import types -from collections import defaultdict -from typing import Any - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Replicate, Shard -from torch.profiler import record_function - -from .adamw import _placement_cache, _tensor_cache, step_adamw -from .async_utils import run_pipeline -from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho, - get_default_muon_param_groups, is_expert_param, update_p) -from .cpu_offload import CPUOffloadPool -from .distributed.utils import (_is_shard, construct_shard_mesh, - get_slices_of_dtensor) -from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO, - _zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5_batched) -from .pipeline import muon_chunk_pipeline, prelaunch_first_gather -from .qk_clip import compute_scales, get_qk_clip_info, qk_clip - -logger = logging.getLogger(__name__) - - -def _expand_expert_params(names, params, expert_keys): - """Expand expert params by splitting on dim 0 (expert dimension). - - Params whose name matches any key in ``expert_keys`` are treated as - expert-parallel tensors. Their outermost dimension is the expert - dimension: an ``(E, out, in)`` tensor becomes ``E`` separate 2D - ``nn.Parameter`` views so that in-place updates propagate back to - the original storage. - - Non-expert params with ``ndim > 2`` trigger an ``AssertionError`` — - if they are expert params, their key must be added to ``expert_keys``. - - The grad must already be set on each expert param (e.g. after momentum). - - For DTensor expert params, placements that shard on dim 0 (expert dim) - are consumed by the split. Non-dim-0 shard placements (e.g. TP) are - preserved: each 2D slice is wrapped as a DTensor on the corresponding - submesh so the parallel pipeline handles the TP communication. - """ - expanded_names = [] - expanded_params = [] - - for n, p in zip(names, params): - is_expert = is_expert_param(n, expert_keys) - is_dtensor = isinstance(p.data, DTensor) - - if is_expert: - if is_dtensor: - logger.debug( - "[expand_expert] %s: expert DTensor, shape=%s, " - "placements=%s, mesh=%s, local_shape=%s", n, p.shape, - p.placements, p.device_mesh.mesh_dim_names, - p.to_local().shape) - else: - logger.debug( - "[expand_expert] %s: expert plain tensor, shape=%s", n, - p.data.shape) - - if not is_expert: - assert p.data.ndim <= 2, ( - f"Param {n} has ndim={p.data.ndim} but does not match " - f"expert_keys={expert_keys}. If this is an expert param, " - f"add its key to expert_keys.") - expanded_names.append(n) - expanded_params.append(p) - continue - - g = p.grad - assert g is not None, ( - f"Expert param {n} must have grad set before expansion") - - tp_mesh = None - tp_placements_2d = None - - if is_dtensor: - local_data = p.to_local() - local_grad = g.to_local() if isinstance(g, DTensor) else g - - # Find non-dim-0 shard placements (e.g. TP sharding). - # After splitting on dim 0, Shard(k) becomes Shard(k-1). - tp_dim_indices = [] - tp_placements_2d = [] - for i, pl in enumerate(p.placements): - if _is_shard(pl) and pl.dim != 0: - tp_dim_indices.append(i) - tp_placements_2d.append(Shard(pl.dim - 1)) - - if tp_dim_indices: - tp_dim_names = tuple(p.device_mesh.mesh_dim_names[i] - for i in tp_dim_indices) - if len(tp_dim_names) == 1: - tp_mesh = p.device_mesh[tp_dim_names[0]] - else: - tp_mesh = p.device_mesh[tp_dim_names] - else: - local_data = p.data - local_grad = g - - # Expand: split dim 0, reshape each slice to 2D. - num_local_experts = local_data.shape[0] - for i in range(num_local_experts): - slice_data = local_data[i] - slice_grad = local_grad[i] - - if tp_mesh is not None: - # Wrap as DTensor on TP submesh so the pipeline handles - # TP communication (gather/scatter across TP ranks). - dt_data = DTensor.from_local(slice_data, - device_mesh=tp_mesh, - placements=tp_placements_2d) - dt_grad = DTensor.from_local(slice_grad, - device_mesh=tp_mesh, - placements=tp_placements_2d) - expert_param = torch.nn.Parameter(dt_data, requires_grad=False) - expert_param.grad = dt_grad - else: - expert_param = torch.nn.Parameter(slice_data, - requires_grad=False) - expert_param.grad = slice_grad - - expanded_names.append(f"{n}[{i}]") - expanded_params.append(expert_param) - - p.grad = None # allow expert grad storage to be freed after pipeline - - return expanded_names, expanded_params - - -class Muon(torch.optim.Optimizer): - """ - Muon - MomentUm Orthogonalized by Newton-schulz - - Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- - processing step, in which each 2D parameter's update is replaced with the nearest orthogonal - matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has - the advantage that it can be stably run in bfloat16 on the GPU. - - Some warnings: - - We believe this optimizer is unlikely to work well for training with small batch size. - - We believe it may not work well for finetuning pretrained models, but we haven't tested this. - - Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. - lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) - momentum: The momentum used by the internal SGD. (0.95 is a good default) - nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) - ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. - Parameters that are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW instead. - adamw_lr: The learning rate for the internal AdamW. - adamw_betas: The betas for the internal AdamW. - adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - warmup_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher warmup_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. - chunk_size : Batch size of parameters to process in each - all2all gather/compute/scatter step. - Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified. - use_distributed_muon: Use distributed muon by Liu et al. (2024). - For testing purpose only. - expert_keys: List of strings to identify expert-parallel parameters. - If any key appears in a parameter's name, its outermost - dimension is treated as the expert dimension and expanded - into per-expert 2D params for Muon. For example, - ``expert_keys=["experts"]`` matches any param whose name - contains "experts". 3D+ params not matched by any key - will raise an error. - """ - - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config=None, - warmup_step=5, - chunk_size=-1, - use_distributed_muon=False, - expert_keys=None): - defaults = dict( - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - nesterov=nesterov, - ns_steps=ns_steps, - adamw_betas=adamw_betas, - adamw_eps=adamw_eps, - none_grad=none_grad, - use_muon=True, - ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) - super().__init__(params, defaults) - - self.debug = debug - self.clip_config = clip_config if clip_config is not None else { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100, - } - self.warmup_step = warmup_step - self.chunk_size = chunk_size - self.use_distributed_muon = use_distributed_muon - self.expert_keys = expert_keys - self.cpu_offload = False - self._cpu_offload_pool: CPUOffloadPool | None = None - self._offload_initialized = False - self._parallel_cache: dict[tuple[str, ...], dict] = {} - self._expert_expand_cache: dict[tuple[int, ...], dict] = {} - - def _calc_flops(self, G, steps): - assert len(G.shape) == 2 - M, N = G.shape - if M > N: - M, N = N, M - - return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) - - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - p.placements, p.device_mesh) - - return shard_mesh, shard_pg, shard_placements - - def init_state_and_assign_params(self, names, params, group, qk_logits): - param_to_state = {} - param_to_flops = {} - - total_flops = 0 - for p in params: - g = p.grad - if g is None: - continue - assert g.ndim == 2, "Muon only supports 2D parameters." - - flops = self._calc_flops(g, group["ns_steps"]) - param_to_flops[id(p)] = flops - total_flops += flops - - if self.debug: - logger.debug("Total TFLOPs for Muon: %.2f TFLOPs", - total_flops / 1e12) - - paired = list(zip(names, params)) - - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) - - round_robin = 0 - mesh = ordered_params[0].device_mesh - placements = ordered_params[0].placements - - shard_mesh, shard_pg, shard_placements = self.get_shard_mesh( - ordered_params[0]) - shard_mesh_flattened = shard_mesh.mesh.flatten() - num_ranks = dist.get_world_size(group=shard_pg) - - for n, p in zip(ordered_names, ordered_params): - if mesh != p.device_mesh: - raise ValueError("All parameters must be on the same mesh.") - if placements != p.placements: - raise ValueError("All parameters must have same placements.") - - worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks - round_robin = (round_robin + 1) % len(shard_mesh_flattened) - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - # Precompute per-rank indices and numels for all-to-all. - rank_indices: dict[int, tuple] = {} - rank_numels: dict[int, int] = {} - for r in range(num_ranks): - indices = get_slices_of_dtensor(p, r, shard_mesh, - shard_placements) - rank_indices[r] = indices - numel = 1 - for idx, dim_size in zip(indices, p.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) - rank_numels[r] = numel - - param_to_state[id(p)] = _muon_state( - worker_rank=worker_rank, - process_group=shard_pg, - rank_indices=rank_indices, - rank_numels=rank_numels, - name=n, - qk_clip_state=qk_clip_state, - ) - - return param_to_state, ordered_params - - def base(self, names, params, group, lr, weight_decay, qk_logits): - # Momentum is already applied by _step_muon before this method. - for n, p in zip(names, params): - g = p.grad - if g is None: - continue - - u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) - - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - def distributed_muon( - self, - names: list[str], - params: list[torch.nn.Parameter], - group: dict[str, Any], - lr: float, - weight_decay: float, - qk_logits: list[torch.Tensor | DTensor] | None, - ): - """Batched Distributed Muon — for testing/correctness verification only. - - Uses all-gather to reconstruct full tensors, computes Newton-Schulz on - the full grad, then slices back to local shards. This is simpler but - slower than the parallel pipeline (all2all) path, so it serves as a - reference implementation for verifying correctness. - """ - with record_function("distributed_muon"): - # Momentum is already applied by _step_muon before this method. - ns_steps = group["ns_steps"] - - # Separate plain tensors (no communication) from DTensors. - plain_names, plain_params = [], [] - dtensor_names, dtensor_params = [], [] - for n, p in zip(names, params): - if p.grad is None: - continue - if isinstance(p.data, DTensor): - dtensor_names.append(n) - dtensor_params.append(p) - else: - plain_names.append(n) - plain_params.append(p) - - # Process plain tensors per-param (no communication). - for n, p in zip(plain_names, plain_params): - u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE), - steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, - qk_logits) - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - if not dtensor_params: - return - - # Group DTensors by (placements, mesh) for batched all-gather. - placement_groups: dict[tuple, - tuple[list, - list]] = defaultdict(lambda: ([], [])) - for n, p in zip(dtensor_names, dtensor_params): - key = (p.placements, p.device_mesh) - placement_groups[key][0].append(n) - placement_groups[key][1].append(p) - - logger.info( - "distributed_muon: %d placement groups, %d total dtensors", - len(placement_groups), len(dtensor_params)) - - for (placements, mesh), (grp_names, - grp_params) in placement_groups.items(): - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - placements, mesh) - rank = dist.get_rank(shard_pg) - world_size = dist.get_world_size(shard_pg) - - logger.info(" group: %d params, placements=%s, world_size=%d", - len(grp_params), placements, world_size) - - # Separate params that can be batched (all shard dims evenly - # divisible) from those needing per-param full_tensor - # (e.g. MoE gate weights with fewer rows than shard ranks). - # all_gather_into_tensor requires equal buffer sizes across - # ranks, so uneven splits must use DTensor full_tensor(). - batch_names, batch_params = [], [] - single_names, single_params = [], [] - for n, p in zip(grp_names, grp_params): - even = all(p.shape[pl.dim] % - shard_mesh.mesh.shape[dim_idx] == 0 - for dim_idx, pl in enumerate(shard_placements)) - if even: - batch_names.append(n) - batch_params.append(p) - else: - single_names.append(n) - single_params.append(p) - - # Process uneven-split params per-param via full_tensor(). - for n, p in zip(single_names, single_params): - with record_function("distributed_muon::newton_schulz"): - g_full = p.grad.full_tensor().to(COMM_DTYPE) - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - if not batch_params: - continue - - logger.info(" batched=%d, single=%d", len(batch_params), - len(single_params)) - - # Concat all local grad shards into a single flat buffer. - with record_function("distributed_muon::gather"): - grad_locals = [ - p.grad.to_local().to(COMM_DTYPE).flatten() - for p in batch_params - ] - numels = [g.numel() for g in grad_locals] - grad_concat = torch.cat(grad_locals) - del grad_locals - - # Single all-gather (replaces N separate full_tensor). - grad_gathered = torch.empty( - grad_concat.numel() * world_size, - dtype=COMM_DTYPE, - device="cuda", - ) - dist.all_gather_into_tensor(grad_gathered, - grad_concat, - group=shard_pg) - - total_numel = grad_concat.numel() - del grad_concat - - # Precompute per-param offsets within the concat buffer. - offsets = [] - off = 0 - for ne in numels: - offsets.append(off) - off += ne - - # Per-param: reconstruct full grad → NS → local update. - for i, (n, p) in enumerate(zip(batch_names, batch_params)): - with record_function("distributed_muon::newton_schulz"): - g_full = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - for r in range(world_size): - r_start = r * total_numel + offsets[i] - shard = grad_gathered[r_start:r_start + numels[i]] - indices = get_slices_of_dtensor( - p, r, shard_mesh, shard_placements) - g_full[indices] = shard.reshape( - g_full[indices].shape) - - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - def _setup_parallel(self, names, params, group, qk_logits): - """Compute (or retrieve cached) parallel pipeline metadata. - - Returns: - (ordered_params, param_to_state, rank, chunk_size) - """ - cache_key = tuple(names) - - if cache_key not in self._parallel_cache: - # First call: compute metadata and populate cache. - param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - shard_pg = param_to_state[id(ordered_params[0])].process_group - rank = dist.get_rank(group=shard_pg) - - if self.chunk_size == -1: - shard_ranks = dist.get_world_size(shard_pg) - chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO - elif self.chunk_size > 0: - chunk_size = self.chunk_size - else: - raise ValueError( - "chunk_size must be -1 or a positive integer.") - - ordered_names = [ - param_to_state[id(p)].name for p in ordered_params - ] - name_to_state = { - param_to_state[id(p)].name: param_to_state[id(p)] - for p in ordered_params - } - self._parallel_cache[cache_key] = { - 'ordered_names': ordered_names, - 'name_to_state': name_to_state, - 'rank': rank, - 'chunk_size': chunk_size, - } - else: - # Cached path: rebuild param_to_state with current id(p) keys. - cache = self._parallel_cache[cache_key] - rank = cache['rank'] - chunk_size = cache['chunk_size'] - - name_to_param = dict(zip(names, params)) - ordered_params = [name_to_param[n] for n in cache['ordered_names']] - - param_to_state = {} - for p, n in zip(ordered_params, cache['ordered_names']): - cached_state = cache['name_to_state'][n] - param_to_state[id(p)] = _muon_state( - worker_rank=cached_state.worker_rank, - process_group=cached_state.process_group, - rank_indices=cached_state.rank_indices, - rank_numels=cached_state.rank_numels, - name=n, - qk_clip_state=get_qk_clip_info(self.clip_config, n, - qk_logits), - ) - - return ordered_params, param_to_state, rank, chunk_size - - def parallel(self, - names, - params, - group, - lr, - weight_decay, - qk_logits, - prelaunch_gather=None): - """ - Perform a parallel optimization step using Muon. - - Parameters are chunked and each chunk is processed by a - :func:`muon_chunk_pipeline` generator. :func:`run_pipeline` - interleaves multiple chunks so that communication and computation - overlap across chunks (the same overlap previously achieved by the - warmup + main-loop index scheduling). - - If ``prelaunch_gather`` is provided, it is passed to the first - chunk's generator to skip re-launching the already in-flight - A2A gather. - """ - - # Momentum is already applied by _step_muon before this method. - - ordered_params, param_to_state, rank, chunk_size = ( - self._setup_parallel(names, params, group, qk_logits)) - - def pipelines(): - first = True - for start in range(0, len(ordered_params), chunk_size): - chunk = ordered_params[start:start + chunk_size] - if chunk: - kwargs = dict( - params=chunk, - param_to_state=param_to_state, - rank=rank, - ns_steps=group["ns_steps"], - lr=lr, - weight_decay=weight_decay, - none_grad=group["none_grad"], - ) - if first and prelaunch_gather is not None: - kwargs['prelaunch_gather'] = prelaunch_gather - first = False - yield muon_chunk_pipeline(**kwargs) - - with record_function("muon::pipeline"): - run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1) - - def _step_muon(self, group, qk_logits=None): - params = group["params"] - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - # Apply momentum to all params before routing/expansion. - # Batched using _foreach_* ops (compiled, fullgraph=True). - with record_function("muon::momentum"): - active_params = [p for p in params if p.grad is not None] - if active_params: - # Ensure momentum buffers exist (avoid zeros_like when already present). - for p in active_params: - if "momentum_buffer" not in self.state[p]: - self.state[p]["momentum_buffer"] = torch.zeros_like( - p.grad) - - # Extract local tensors for compiled batch function. - local_grads = [ - p.grad._local_tensor - if isinstance(p.grad, DTensor) else p.grad - for p in active_params - ] - local_bufs = [ - self.state[p]["momentum_buffer"]._local_tensor - if isinstance(self.state[p]["momentum_buffer"], DTensor) - else self.state[p]["momentum_buffer"] - for p in active_params - ] - - # Wrap momentum as tensor for torch.compile. - batch_pre_ortho(local_grads, local_bufs, - torch.tensor(momentum), group["nesterov"]) - - # For non-nesterov, the result is the momentum buffer. - if not group["nesterov"]: - for p in active_params: - p.grad = self.state[p]["momentum_buffer"] - - # Identify batched experts for deferred NS. - # Detection is cheap (condition checks only); actual NS compute is - # deferred so it can overlap with the first chunk's A2A gather. - deferred_expert_work = [] - if self.expert_keys: - batched_expert_indices = [] - for i, (n, p) in enumerate(zip(names, params)): - if not (is_expert_param(n, self.expert_keys) - and p.grad is not None): - continue - # Eligible: plain tensor, or DTensor with no non-dim-0 shards. - if isinstance(p.data, DTensor): - has_tp = any( - _is_shard(pl) and pl.dim != 0 for pl in p.placements) - if has_tp: - continue - batched_expert_indices.append(i) - - if batched_expert_indices: - # Save refs for deferred NS; free grads from param list. - for i in batched_expert_indices: - p = params[i] - g = p.grad - local_g = (g._local_tensor - if isinstance(g, DTensor) else g) - local_data = (p.data._local_tensor if isinstance( - p.data, DTensor) else p.data) - deferred_expert_work.append((local_data, local_g)) - p.grad = None - - # Remove batched experts from lists before expansion. - keep = sorted( - set(range(len(params))) - set(batched_expert_indices)) - names = [names[i] for i in keep] - params = [params[i] for i in keep] - - def _run_deferred_expert_ns(): - """Execute deferred batched expert NS.""" - if not deferred_expert_work: - return - with record_function("muon::batched_expert_ns"): - ns_steps = group["ns_steps"] - for local_data, local_g in deferred_expert_work: - u = zeropower_via_newtonschulz5_batched( - local_g.to(COMM_DTYPE), steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:]) - local_data.mul_(1 - lr * weight_decay) - local_data.add_(u, alpha=-adjusted_lr) - - # Expand expert params by splitting on dim 0. - logger.debug("[_step_muon] before expand: %d params, expert_keys=%s", - len(params), self.expert_keys) - if self.expert_keys: - cache_key = tuple(id(p) for p in params) - cache = self._expert_expand_cache.get(cache_key) - - if cache is None: - # Cold path: full expansion + build cache metadata. - exp_names, exp_params = _expand_expert_params( - names, params, self.expert_keys) - - # Build per-expert-group info for hot-path grad updates. - grad_info = [] - exp_idx = 0 - for orig_idx, (n, p) in enumerate(zip(names, params)): - if not is_expert_param(n, self.expert_keys): - exp_idx += 1 - continue - - is_dt = isinstance(p.data, DTensor) - num_experts = (p.to_local() if is_dt else p.data).shape[0] - - # Detect TP mesh from the first expanded expert param. - tp_mesh = None - tp_pls = None - sample = exp_params[exp_idx] - if isinstance(sample.data, DTensor): - tp_mesh = sample.data.device_mesh - tp_pls = list(sample.data.placements) - - grad_info.append((orig_idx, num_experts, exp_idx, is_dt, - tp_mesh, tp_pls)) - exp_idx += num_experts - - self._expert_expand_cache[cache_key] = { - 'names': exp_names, - 'params': exp_params, - 'grad_info': grad_info, - } - names, params = exp_names, exp_params - else: - # Hot path: reuse cached params, only update expert grads. - for (orig_idx, num_experts, exp_start, is_dt, tp_mesh, - tp_pls) in cache['grad_info']: - p = params[orig_idx] - g = p.grad - local_grad = (g.to_local() - if is_dt and isinstance(g, DTensor) else g) - for i in range(num_experts): - expert_p = cache['params'][exp_start + i] - sg = local_grad[i] - if tp_mesh is not None: - expert_p.grad = DTensor.from_local( - sg, device_mesh=tp_mesh, placements=tp_pls) - else: - expert_p.grad = sg - p.grad = None - - names = cache['names'] - params = cache['params'] - else: - names, params = _expand_expert_params(names, params, - self.expert_keys) - logger.debug("[_step_muon] after expand: %d params", len(params)) - - param_dtensors = [] - name_dtensors = [] - - param_tensors = [] - name_tensors = [] - - # distributed_muon is a reference implementation for testing only. - # The parallel pipeline (all2all) path below is the production path. - if self.use_distributed_muon: - _run_deferred_expert_ns() - self.distributed_muon(names=names, - params=params, - group=group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits) - return - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - logger.debug( - "[route] %s → base (DTensor all-Replicate), " - "shape=%s, placements=%s", n, p.shape, p.placements) - param_tensors.append(p) - name_tensors.append(n) - else: - logger.debug( - "[route] %s → parallel (DTensor), shape=%s, " - "placements=%s, mesh=%s", n, p.shape, p.placements, - p.device_mesh.mesh_dim_names) - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - logger.debug("[route] %s → base (plain tensor), shape=%s", n, - p.data.shape) - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError(f"Unsupported parameter type: {type(p.data)}") - - logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, " - f"{len(param_tensors)} Tensors → base") - - def group_dtensors(dtensors, names): - # To support different placements, we group parameters by placements - # and run parallel Muon on each group. - - placement_to_params = defaultdict(lambda: ([], [])) - - assert len(dtensors) == len(names) - for p, n in zip(dtensors, names): - placement_to_params[tuple([p.placements, - p.device_mesh])][0].append(n) - placement_to_params[tuple([p.placements, - p.device_mesh])][1].append(p) - return placement_to_params - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - dtensor_group = group_dtensors(param_dtensors, name_dtensors) - - # Pre-launch the first chunk's A2A gather so that the NCCL - # communication overlaps with the (deferred) batched expert NS - # compute on the default CUDA stream. - prelaunch = None - if deferred_expert_work: - first_names, first_params = next(iter(dtensor_group.values())) - ordered, pts, rnk, csz = self._setup_parallel( - first_names, first_params, group, qk_logits) - first_chunk = ordered[:csz] - if first_chunk: - prelaunch = prelaunch_first_gather(first_chunk, pts, rnk, - group["none_grad"]) - - _run_deferred_expert_ns() - - first_group = True - for _, (names, params) in dtensor_group.items(): - pg = prelaunch if first_group else None - first_group = False - self.parallel( - names, - params, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - prelaunch_gather=pg, - ) - else: - _run_deferred_expert_ns() - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - ) - - def _register_states_for_offload(self): - """Register all optimizer state tensors with the CPU offload pool. - - Called once after the first step when states have been lazily created. - Offloads all param states (momentum buffers for Muon, moment1/moment2 - for AdamW) to free GPU memory between steps. - """ - pool = self._cpu_offload_pool - tracked = 0 - for group in self.param_groups: - for p in group["params"]: - if p not in self.state: - continue - state = self.state[p] - if group.get("use_muon", False): - if "momentum_buffer" in state: - pool.track(state["momentum_buffer"]) - tracked += 1 - else: - if "moment1" in state: - pool.track(state["moment1"]) - if "moment2" in state: - pool.track(state["moment2"]) - tracked += 1 - logger.info("[CPUOffload] Registered %d param states for offload", - tracked) - - @torch.no_grad - def step(self, closure=None, qk_logits=None): - """Perform a single optimization step. - - Args: - closure (Callable, optional): A closure that reevaluates the model - and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). - """ - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # H2D: reload optimizer states from CPU before computation. - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload() - - logger.debug("[Muon.step] expert_keys=%s, %d param groups", - self.expert_keys, len(self.param_groups)) - - for i, group in enumerate(self.param_groups): - if group["use_muon"]: - logger.debug("[Muon.step] group %d: use_muon=True, %d params", - i, len(group["params"])) - self._step_muon(group, qk_logits=qk_logits) - else: - logger.debug( - "[Muon.step] group %d: use_muon=False (AdamW), %d params", - i, len(group["params"])) - step_adamw(self.state, group) - - # D2H: offload optimizer states to CPU after computation. - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - return loss - - # ------------------------------------------------------------------ - # CPU offload public helpers - # ------------------------------------------------------------------ - - def turn_on_cpu_offload(self): - """Enable CPU offload for optimizer states.""" - if self.cpu_offload: - return - logger.info("[Muon] turn_on_cpu_offload") - self.cpu_offload = True - if not self.state: - return - self._cpu_offload_pool = CPUOffloadPool() - self._offload_initialized = False - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_off_cpu_offload(self): - """Disable CPU offload and keep optimizer states resident on GPU.""" - if not self.cpu_offload: - return - logger.info("[Muon] turn_off_cpu_offload") - if self._offload_initialized: - self._cpu_offload_pool.reload() - torch.cuda.current_stream().synchronize() - self._cpu_offload_pool = None - self._offload_initialized = False - self.cpu_offload = False - - # ------------------------------------------------------------------ - # Checkpoint support for cpu_offload - # ------------------------------------------------------------------ - - def state_dict(self) -> dict: - if self.cpu_offload: - raise RuntimeError( - "Muon.state_dict() requires turn_off_cpu_offload() before checkpoint save." - ) - return super().state_dict() - - def load_state_dict(self, state_dict: dict) -> None: - if self.cpu_offload: - raise RuntimeError( - "Muon.load_state_dict() requires turn_off_cpu_offload() before checkpoint load." - ) - super().load_state_dict(state_dict) - - # Invalidate adamw.py's module-level tensor caches so that - # the next step rebuilds them with the newly loaded state tensors. - _placement_cache.clear() - _tensor_cache.clear() diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/newton_schulz.py b/build/torch28-cxx11-rocm63-x86_64-linux/newton_schulz.py deleted file mode 100644 index d939264b69a34e7a3fa78859f34dc265a1159d59..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/newton_schulz.py +++ /dev/null @@ -1,240 +0,0 @@ -from itertools import repeat -from math import inf, sqrt - -import numpy as np -import torch - -from .matmul_transpose_triton import matmul_transpose_assign - -COMM_DTYPE = torch.bfloat16 -DEFAULT_CHUNK_SIZE_RATIO = 4 - - -def _optimal_quintic(l, u, max_iter=1000): - """ - Use the simplified Remez algorithm to find the optimal odd quintic approximant - to the constant function x -> 1 over the interval [l, u]. - - Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum - approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the - two interior equioscillation nodes q, r until convergence. Returns the - closed-form equioscillating solution when l ≈ u. - - Raises ValueError if any intermediate value (a, b, c, E, q, r) is non-finite - (NaN or inf). Raises RuntimeError if convergence is not reached within - max_iter iterations. - """ - assert 0 <= l <= u - if 1 - 5e-6 <= l / u: - return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5) - q = (3 * l + u) / 4 - r = (l + 3 * u) / 4 - E = inf - for _ in range(max_iter): - old_E = E - LHS = np.array( - [ - [l, l**3, l**5, 1], - [q, q**3, q**5, -1], - [r, r**3, r**5, 1], - [u, u**3, u**5, -1], - ] - ) - a, b, c, E = np.linalg.solve(LHS, np.ones(4)) - if not np.all(np.isfinite([a, b, c, E])): - raise ValueError( - f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}" - ) - q, r = np.sqrt( - (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / (10 * c) - ) - if not np.all(np.isfinite([q, r])): - raise ValueError(f"_optimal_quintic: non-finite node update q={q}, r={r}") - if abs(old_E - E) <= 1e-15: - break - else: - raise RuntimeError( - f"_optimal_quintic: did not converge after {max_iter} iterations" - ) - return float(a), float(b), float(c) - - -def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0): - """ - Compute the Polar Express coefficient series for `num_iters` quintic iterations. - - Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that - compose to map singular values from [l, 1] toward 1. At each step: - 1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion` - prevents near-zero singular values from stalling by raising the effective - lower bound; if it is active (cushion*u > l), the coefficients are - rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u]. - 2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the - last iteration, providing numerical headroom at the cost of a slightly slower - final convergence step. - 3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1). - - Returns a list of (a, b, c) tuples, one per iteration. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - u = 1 - assert 0 <= l <= u - safety_factor = 1 + safety_factor_eps - coefficients = [] - for iter in range(num_iters): - a, b, c = _optimal_quintic(max(l, cushion * u), u) - if cushion * u > l: - pl = a * l + b * l**3 + c * l**5 - pu = a * u + b * u**3 + c * u**5 - rescaler = 2 / (pl + pu) - a *= rescaler - b *= rescaler - c *= rescaler - if iter < num_iters - 1: - a /= safety_factor - b /= safety_factor**3 - c /= safety_factor**5 - coefficients.append((a, b, c)) - l = a * l + b * l**3 + c * l**5 - u = 2 - l - return coefficients - - -# Precomputed Polar Express coefficients (a, b, c) for 10 quintic Newton-Schulz -# iterations. Each tuple is the minimax-optimal (Remez/equioscillation) odd quintic -# approximant to x->1 over the current singular-value interval, computed once at -# import time and reused across all optimizer steps. -# -# Contrast with the former hardcoded NS coefficients (5 fixed tuples): -# - Former: empirically tuned to maximize slope at zero; did not converge -# singular values to 1, yielding US'V^T with S' ~ Uniform(0.5, 1.5) instead -# of the true polar factor UV^T. -# - Polar Express: analytically optimal per step, adapting to the shrinking -# singular-value interval [l, u] as iterations progress; converges all -# singular values to 1, producing the exact polar factor UV^T. -_coeffs_list = _optimal_composition( - l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02 -) - - -# This code is adapted from: -# KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py) -# NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress) -# matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon) -@torch.no_grad() -def _zeropower_via_newtonschulz5(G, steps): - """ - Compute the polar factor of G via the Polar Express method. - - Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c) - are the Polar Express coefficients from `_coeffs_list`. Each step is the - optimal odd quintic approximant to x -> 1 over the current singular-value - interval, minimizing the maximum approximation error (Remez / minimax criterion). - The composition maps singular values from [l, 1] to near 1, producing the - polar factor (orthogonal factor in the polar decomposition G = UP). - - `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2, - cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE - X = G # no manual typecast - - if G.size(0) > G.size(1): - X = X.T - - X = X / (X.norm() + 1e-7) - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - # Perform the NS iterations - for a, b, c in hs: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(0) > G.size(1): - X = X.T - - return X - - -@torch.no_grad() -def _zeropower_via_newtonschulz5_batched(G, steps): - """Batched polar factor computation for 3D (E, out, in) tensors. - - Same algorithm as ``_zeropower_via_newtonschulz5`` but uses - ``torch.bmm`` / ``torch.baddbmm`` instead of the 2D Triton kernel, - processing all E expert matrices in a single batched call. - """ - assert len(G.shape) == 3 - assert G.dtype == COMM_DTYPE - X = G - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - # Per-expert Frobenius norm. - X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) - - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - for a, b, c in hs: - buf1 = torch.bmm(X, X.transpose(-2, -1)) - buf2 = torch.bmm(buf1, buf1.transpose(-2, -1)) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.baddbmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - return X - - -_ns_per_shape: dict[tuple[int, ...], callable] = {} -_use_compile = True - - -def set_ns_compile(enabled: bool): - """Toggle torch.compile for Newton-Schulz iteration.""" - global _use_compile - _use_compile = enabled - - -def zeropower_via_newtonschulz5(G, steps=5): - if not _use_compile: - return _zeropower_via_newtonschulz5(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() - - -def zeropower_via_newtonschulz5_batched(G, steps=5): - """Compile-cached batched Newton-Schulz for 3D expert tensors.""" - if not _use_compile: - return _zeropower_via_newtonschulz5_batched(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile( - _zeropower_via_newtonschulz5_batched, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/optimizer/__init__.py b/build/torch28-cxx11-rocm63-x86_64-linux/optimizer/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/optimizer/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/pipeline.py b/build/torch28-cxx11-rocm63-x86_64-linux/pipeline.py deleted file mode 100644 index c0c2d515856182d8d15ad27dd4e4e093b29397d6..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/pipeline.py +++ /dev/null @@ -1,468 +0,0 @@ -import logging -from typing import Generator - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -from .core import _muon_state, adjust_lr_for_muon -from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5 -from .qk_clip import compute_scales - -logger = logging.getLogger(__name__) - -# ====================================================================== -# Stage helpers -# ====================================================================== - - -def _launch_gather( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Allocate gather buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_gather``). - gathered_grads: ``{id(p): empty_tensor}`` for owned params, - ``None`` for non-owned. - recv_counts: Per-source-rank element counts. - """ - # Allocate gathered-grad buffers - gathered_grads: dict[int, torch.Tensor | None] = {} - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - gathered_grads[id(p)] = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - else: - gathered_grads[id(p)] = None - - # Build send buffer – batch grad copies via torch.cat - # (1-2 fused kernels vs N individual narrow().copy_() calls). - send_counts = [0] * num_ranks - for p in params: - state = param_to_state[id(p)] - send_counts[state.worker_rank] += state.rank_numels[rank] - - total_send = sum(send_counts) - if total_send > 0: - # Group grad slices by destination rank in a single pass. - dst_to_grads = [[] for _ in range(num_ranks)] - for p in params: - state = param_to_state[id(p)] - n = state.rank_numels[rank] - if n > 0: - g = p.grad.to_local() - dst_to_grads[state.worker_rank].append(g.reshape(-1)) - - # Flatten in dst order and cat once. - all_slices = [] - for dst in range(num_ranks): - all_slices.extend(dst_to_grads[dst]) - send_buf = torch.cat(all_slices) - if send_buf.dtype != COMM_DTYPE: - send_buf = send_buf.to(COMM_DTYPE) - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += state.rank_numels[src] - recv_counts[src] = total - - recv_buf = torch.empty(sum(recv_counts), dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - logger.debug(f"send_buf size: {send_buf.numel()}, " - f"recv_buf size: {recv_buf.numel()}, " - f"recv_counts: {recv_counts}, " - f"send_counts: {send_counts}, " - f"process_group: {str(process_group)}") - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, gathered_grads, recv_counts - - -def _complete_gather( - recv_buf: torch.Tensor, - recv_counts: list[int], - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - param_to_state: dict[int, _muon_state], - rank: int, -) -> None: - """Reconstruct gathered grads from the recv buffer (in-place).""" - off = 0 - for src in range(len(recv_counts)): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - - indices = state.rank_indices[src] - - shard_view = gathered_grads[id(p)][indices] - n = shard_view.numel() - if n == 0: - continue - - sg = recv_buf.narrow(0, off + inner_off, n) - sg = sg.reshape(shard_view.shape) - gathered_grads[id(p)][indices] = sg - - inner_off += n - assert inner_off == block - off += block - - -def _compute_ns( - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - ns_steps: int, -) -> dict[int, torch.Tensor | None]: - """Run Newton-Schulz orthogonalization on owned parameters. - - Returns: - computed_us: ``{id(p): orthogonalized_update}`` for owned params. - """ - computed_us: dict[int, torch.Tensor | None] = {} - for p in owned_params: - u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps) - gathered_grads[id(p)] = None # free gathered grad - computed_us[id(p)] = u - return computed_us - - -def _launch_scatter( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, - computed_us: dict[int, torch.Tensor | None], -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor], list[int]]: - """Allocate scatter buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_scatter``). - scattered_us: Empty dict, populated by ``_complete_scatter`` with - zero-copy views into ``recv_buf``. - recv_counts: Per-source-rank element counts. - """ - # scattered_us is populated by _complete_scatter with zero-copy views - # into recv_buf, avoiding N empty_like allocations + N copy_ calls. - # Pre-seed entries for params whose local shard is empty (rank_numels == 0) - # so _update_params can iterate all params without KeyError. - scattered_us: dict[int, torch.Tensor] = {} - for p in params: - if param_to_state[id(p)].rank_numels[rank] == 0: - scattered_us[id(p)] = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - # Build send buffer – batch via torch.cat - # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls). - send_counts = [0] * num_ranks - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - for dst_rank in range(num_ranks): - send_counts[dst_rank] += state.rank_numels[dst_rank] - - total_send = sum(send_counts) - if total_send > 0: - # Cache u_full conversions to avoid redundant .to() per dst_rank. - u_fulls = {} - for p in owned_params: - u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous() - - # Collect slices in dst order (matches all-to-all send layout). - all_slices = [] - for dst_rank in range(num_ranks): - for p in owned_params: - state = param_to_state[id(p)] - su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten() - if su.numel() > 0: - all_slices.append(su) - - send_buf = torch.cat(all_slices) if all_slices else torch.empty( - 0, dtype=COMM_DTYPE, device="cuda") - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += state.rank_numels[rank] - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, scattered_us, recv_counts - - -def _complete_scatter( - recv_buf: torch.Tensor, - recv_counts: list[int], - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], -) -> None: - """Populate scattered_us with zero-copy views into recv_buf. - - Instead of pre-allocating tensors and copying, we assign views directly - from ``recv_buf``. This eliminates N ``empty_like`` + N ``copy_`` calls. - The underlying storage of ``recv_buf`` is kept alive through the views - until ``scattered_us`` is cleared after ``_update_params``. - """ - off = 0 - for src in range(len(recv_counts)): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = state.rank_numels[rank] - if n == 0: - continue - - scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - - inner_off += n - - assert inner_off == block - off += block - - -def _update_params( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], - lr: float, - weight_decay: float, -) -> None: - """Apply weight decay, Muon update, and optional QK clipping. - - Uses batched ``_foreach_mul_`` for weight decay and batched - ``_foreach_add_`` for the Muon update, grouping parameters by - adjusted_lr to minimize kernel launches while preserving float32 - precision for the alpha scaling. - """ - if not params: - return - - # Batched weight decay: p *= (1 - lr * wd) — single fused kernel. - p_locals = [p._local_tensor for p in params] - torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay) - - # Group params by adjusted_lr so _foreach_add_ can use a single - # alpha per group (preserves float32 precision for alpha scaling). - lr_groups: dict[float, tuple[list, list]] = {} - for p in params: - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - if adjusted_lr not in lr_groups: - lr_groups[adjusted_lr] = ([], []) - lr_groups[adjusted_lr][0].append(p._local_tensor) - lr_groups[adjusted_lr][1].append(scattered_us[id(p)]) - - for adjusted_lr, (p_group, u_group) in lr_groups.items(): - torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr) - - # QK clipping – applied directly on the local tensor to - # avoid DTensor sharding-propagation issues with _StridedShard. - for p in params: - state = param_to_state[id(p)] - if state.qk_clip_state is None: - continue - scales_full = compute_scales(p, state.qk_clip_state) - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = state.rank_indices[rank][0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - -# ====================================================================== -# Pre-launch helper for overlapping first chunk's gather with other work. -# ====================================================================== - - -@torch.no_grad() -def prelaunch_first_gather( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - none_grad: bool, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Launch the first chunk's A2A gather early for overlap with other compute. - - Call this *before* expensive GPU work (e.g. batched expert NS) so that - the NCCL all-to-all runs concurrently on the NCCL stream while the - default stream executes compute. - - Returns the same 4-tuple that ``_launch_gather`` produces, which should - be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - with record_function("muon::prelaunch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - return work, recv_buf, gathered_grads, recv_counts - - -# ====================================================================== -# Main generator – thin orchestrator that wires stages together. -# ====================================================================== - - -@torch.no_grad() -def muon_chunk_pipeline( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - ns_steps: int, - lr: float, - weight_decay: float, - none_grad: bool, - prelaunch_gather: tuple | None = None, -) -> Generator[None, None, None]: - """Process one chunk of parameters through the full Muon pipeline. - - Stages: gather -> compute (Newton-Schulz) -> scatter -> update. - - Each ``yield`` lets :func:`run_pipeline` interleave other chunks so - that communication and computation overlap across chunks. Async - communication is launched via ``async_op=True`` and completed after - the yield with ``work.wait()``. - - Overlap happens because :func:`run_pipeline` admits one new chunk - per iteration (staggered admission). While chunk *N* does NS - compute on the default CUDA stream, chunk *N+1*'s async all-to-all - runs concurrently on the NCCL stream — no separate ``comm_stream`` - is required. - - If ``prelaunch_gather`` is provided, the gather was already launched - by :func:`prelaunch_first_gather` and we skip launching it again. - - Yields exactly **2** times: - - 1. After launching async all-to-all gather (or immediately if pre-launched). - 2. After launching async all-to-all scatter. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - if prelaunch_gather is not None: - # Gather was pre-launched; none_grad already handled by caller. - work, recv_buf, gathered_grads, recv_counts = prelaunch_gather - else: - # Normal path: launch async gather. - with record_function("muon::launch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - yield # --- YIELD 1: other chunks can launch their gather --- - - with record_function("muon::wait_gather"): - work.wait() - _complete_gather(recv_buf, recv_counts, owned_params, gathered_grads, - param_to_state, rank) - del recv_buf - - # Stage 3: Newton-Schulz orthogonalization. - with record_function("muon::newton_schulz"): - computed_us = _compute_ns(owned_params, gathered_grads, ns_steps) - gathered_grads.clear() - - # Stages 4-5: launch async scatter. - with record_function("muon::launch_scatter"): - work, recv_buf, scattered_us, recv_counts = _launch_scatter( - params, owned_params, param_to_state, rank, num_ranks, - process_group, computed_us) - computed_us.clear() - - yield # --- YIELD 2: other chunks can launch their scatter --- - - with record_function("muon::wait_scatter"): - work.wait() - _complete_scatter(recv_buf, recv_counts, params, param_to_state, rank, - scattered_us) - del recv_buf - - # Stage 6: apply parameter updates. - with record_function("muon::update_params"): - _update_params(params, param_to_state, rank, scattered_us, lr, - weight_decay) - scattered_us.clear() diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/qk_clip.py b/build/torch28-cxx11-rocm63-x86_64-linux/qk_clip.py deleted file mode 100644 index 2aba711b3004b7f09e7141da7ef834bd61cc2430..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/qk_clip.py +++ /dev/null @@ -1,198 +0,0 @@ -import logging -import math -from dataclasses import dataclass - -import torch -from torch.distributed.tensor import DTensor - -from .core import normalize_fqn - -logger = logging.getLogger(__name__) - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - and return (kind, layer_index). - - Supported kinds: - MHA/GQA: 'wq', 'wk', 'q_proj', 'k_proj' - MLA: 'wq_b' (Q up-proj), 'wkv_b' (KV up-proj) - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.1.attn.wq_b.weight' -> ('wq_b', 1) - 'model.0.attn.wkv_b.weight' -> ('wkv_b', 0) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = normalize_fqn(name).split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: str | None # 'wq'/'q_proj'/'wq_b' or 'wk'/'k_proj'/'wkv_b' or None - indices: list[int] # which heads to consider for clipping - head_dim: int # from config (qk_head_dim for MLA wq_b) - threshold: float # from config - logit: torch.Tensor | None - - # MLA-specific fields - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 - - -def get_qk_clip_info(clip_config, n, qk_logits): - """Extract QK clipping info for a named parameter. - - Args: - clip_config: QK clipping configuration dict (or None). - MHA/GQA keys: head_dim, threshold, q_indices, k_indices - MLA extra keys: is_mla=True, qk_nope_head_dim, qk_rope_head_dim, v_head_dim - n: Parameter name string. - qk_logits: Dict mapping layer indices to logit tensors (or None). - - Returns: - QKClipInfo instance with clipping configuration for this parameter. - """ - if clip_config is None: - return None - - head_dim = clip_config.get('head_dim') - threshold = clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - # In TP settings, qk_logits may be DTensor - # We convert it to full tensor here for simplicity - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - -def compute_scales(p, qk_clip_state): - """Compute per-head scaling factors for QK clipping. - - Returns scales tensor (√γ per head) if any head exceeds threshold, else None. - For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim. - """ - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - # Check if any head exceeds threshold before allocating. - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - - if not head_scales: - return None - - # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows - if qk_clip_state.is_mla and kind == 'wkv_b': - effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim - else: - effective_head_dim = head_dim - - H_global = p.shape[0] // effective_head_dim - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale - return scales_full - - -def qk_clip(p, scales, info): - """Apply per-head scaling to a Q/K projection weight matrix. - - Args: - p: Parameter (nn.Parameter or raw tensor). - scales: [n_heads] tensor, each element = √γ_h. - info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions. - - MLA sub-region scaling per Algorithm 1 (MuonClip): - wq_b: q_nope rows → √γ, q_pe rows → γ - wkv_b: k_nope rows → √γ, v rows → unchanged - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not info.is_mla: - # MHA/GQA: uniform √γ applied to all rows in each head - W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: vectorized sub-region scaling within each head - if info.kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_head_dim = qk_nope + info.qk_rope_head_dim - W_3d = W.view(-1, qk_head_dim, W.shape[1]) # [H, qk_head_dim, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # q_nope → √γ - W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1, - 1)) # q_pe → γ - - elif info.kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - W_3d = W.view(-1, kv_stride, W.shape[1]) # [H, kv_stride, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # k_nope → √γ - # v rows: not touched (k_R shared rotary unchanged) diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/__init__.py b/build/torch28-cxx11-rocm64-x86_64-linux/__init__.py deleted file mode 100644 index 239c7a65f8293e7d0df28f05fce645af56d628c0..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .muon import Muon - -__all__ = [ - "Muon", -] diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/_ops.py b/build/torch28-cxx11-rocm64-x86_64-linux/_ops.py deleted file mode 100644 index 034bff088659b5df6f6d401feb18c89dc5f33b29..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _optimizer_8d53b78_dirty -ops = torch.ops._optimizer_8d53b78_dirty - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_optimizer_8d53b78_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so b/build/torch28-cxx11-rocm64-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so deleted file mode 100755 index 9d0435e26c1566f1da7309ceeb31f49e290217cb..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2f43a2025f967fcd94fcada1e8a708956f07774c522f156caab135d7162c7a91 -size 1865168 diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/adamw.py b/build/torch28-cxx11-rocm64-x86_64-linux/adamw.py deleted file mode 100644 index b5a95816a9f5b9e1889eaadae65373bfbced809a..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/adamw.py +++ /dev/null @@ -1,271 +0,0 @@ -import logging -from collections import defaultdict -from typing import cast - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -def fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: float | torch.Tensor, - weight_decay: float, - eps: float, - maximize: bool, -) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: dict | None = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) - - -def _to_local(t): - """Unwrap DTensor to local tensor for fused ops.""" - return t._local_tensor if isinstance(t, DTensor) else t - - -# --------------------------------------------------------------------------- -# Caches for eliminating per-step Python overhead. -# -# Placement grouping and tensor list assembly are identical every step -# (params don't change placement, moment/step tensors are the same objects -# after initialisation). We cache them keyed by id() of the param list -# stored in param_groups (stable across steps). -# -# Only gradients change each step and must be collected fresh. -# --------------------------------------------------------------------------- - -# id(group["params"]) → dict[placement_key, list[param]] -_placement_cache: dict[int, dict[tuple, list]] = {} - -# id(placement_group_list) → (params_local, moment1, moment2, state_steps) -_tensor_cache: dict[int, tuple[list, list, list, list]] = {} - - -def _step_adamw_params_slow(optimizer_state, params, group): - """Uncached fallback for the rare case where some params lack grads.""" - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - g = p.grad - if g is None: - continue - state = optimizer_state[p] - params_with_grads.append(_to_local(p)) - grads.append(_to_local(g)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - if not params_with_grads: - return - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw_params(optimizer_state, params, group): - """Run fused AdamW on a list of parameters sharing the same placement. - - After the first call, cached tensor lists (params_local, moment1, - moment2, state_steps) are reused — only gradients are collected fresh. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - params: List of parameters to update. - group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay. - """ - # Collect grads — the only thing that changes each step. - with record_function("adamw::collect_grads"): - grads = [] - for p in params: - g = p.grad - if g is None: - # Rare: fall back to slow path that filters per-param. - _step_adamw_params_slow(optimizer_state, params, group) - return - grads.append(_to_local(g)) - - tensor_key = id(params) - if tensor_key not in _tensor_cache: - with record_function("adamw::init_tensor_cache"): - params_local = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - state = optimizer_state[p] - params_local.append(_to_local(p)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(p.grad) - state["moment2"] = torch.zeros_like(p.grad) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - _tensor_cache[tensor_key] = (params_local, moment1, moment2, - state_steps) - - params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key] - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - with record_function("adamw::fused_adamw"): - fused_adamw( - params_local, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw(optimizer_state, group): - """Dispatch AdamW step, grouping parameters by type and placement. - - Placement grouping is cached after the first call since params never - change their placement between steps. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - group: Parameter group dict. - """ - params = group["params"] - placement_key = id(params) - - if placement_key not in _placement_cache: - with record_function("adamw::group_by_placement"): - placement_to_params: dict[tuple, - list[torch.Tensor]] = defaultdict(list) - for p in params: - match p: - case DTensor(): - logger.debug( - "[AdamW] DTensor param: shape=%s, placements=%s, " - "mesh=%s, grad=%s", p.shape, p.placements, - p.device_mesh.mesh_dim_names, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple( - [p.placements, p.device_mesh])].append(p) - case torch.Tensor(): - logger.debug( - "[AdamW] plain param: shape=%s, grad=%s", p.shape, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple([torch.Tensor, - None])].append(p) - - logger.debug("[AdamW] %d placement groups, %d total params", - len(placement_to_params), len(params)) - - _placement_cache[placement_key] = dict(placement_to_params) - - for group_params in _placement_cache[placement_key].values(): - step_adamw_params(optimizer_state, group_params, group) diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/async_utils.py b/build/torch28-cxx11-rocm64-x86_64-linux/async_utils.py deleted file mode 100644 index a45c530ac9cad88e3555ec1047a6aa59f225347e..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/async_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import logging -from typing import Generator - -logger = logging.getLogger(__name__) - - -class _Task: - """Internal: wraps a generator, advances one yield at a time.""" - - def __init__(self, generator: Generator[None, None, None], index: int): - self._generator = generator - self._index = index - self._steps_completed = 0 - self.step() # run to first yield - - def step(self) -> bool: - try: - next(self._generator) - self._steps_completed += 1 - logger.debug("pipeline[%d] completed stage %d", self._index, - self._steps_completed) - return True - except StopIteration: - logger.debug("pipeline[%d] finished after %d stages", self._index, - self._steps_completed) - return False - - def close(self): - self._generator.close() - - -def run_pipeline( - pipelines: Generator[Generator[None, None, None], None, None], - max_concurrent: int, -) -> None: - """Run generator-based pipelines with bounded concurrency. - - Each pipeline is a generator that yields at stage boundaries. - The runtime interleaves pipelines so communication and computation - overlap across chunks. - """ - if max_concurrent <= 0: - raise ValueError(f"max_concurrent must be > 0, got {max_concurrent}") - - have_new = True - task_index = 0 - previous_tasks: list[_Task] = [] - - try: - while have_new or previous_tasks: - running_tasks: list[_Task] = [] - - # Admit one new pipeline per iteration (staggered admission). - # Admitting one at a time ensures that while chunk N does NS - # compute on the default stream, chunk N+1's NCCL all-to-all - # runs concurrently on the NCCL stream — creating real - # communication/computation overlap on the GPU. - if have_new and len(previous_tasks) < max_concurrent: - try: - gen = next(pipelines) - task = _Task(gen, task_index) - task_index += 1 - running_tasks.append(task) - except StopIteration: - have_new = False - - # Advance every previously-yielded task by one step. - for task in previous_tasks: - if task.step(): - running_tasks.append(task) - - previous_tasks = running_tasks - except BaseException: - # Clean up all in-flight generators to release GPU resources. - for task in previous_tasks: - task.close() - raise diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/core.py b/build/torch28-cxx11-rocm64-x86_64-linux/core.py deleted file mode 100644 index c69d515afef305ad0ed66374095fa2d2468d99cc..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/core.py +++ /dev/null @@ -1,219 +0,0 @@ -import logging -import math -from dataclasses import dataclass -from typing import List - -import torch -from torch.distributed import ProcessGroup -from torch.distributed.tensor import DTensor - -# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into -# parameter FQNs. Activation checkpointing similarly inserts -# "_checkpoint_wrapped_module". Strip these so name-based matching (skip_keys, -# expert_keys, QK layer parsing) works regardless of wrapper nesting. -_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) - -logger = logging.getLogger(__name__) - - -def normalize_fqn(name: str) -> str: - """Strip torch.compile / checkpoint wrapper components from a parameter FQN.""" - return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS) - - -@dataclass -class _muon_state: - worker_rank: int - process_group: ProcessGroup - rank_indices: dict[int, tuple] # local_rank -> per-dim indices - rank_numels: dict[int, int] # local_rank -> numel - name: str - qk_clip_state: torch.Tensor | None = None - - -def _batch_momentum( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update (no nesterov).""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - - -def _batch_momentum_nesterov( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update with nesterov correction.""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - nesterov_terms = torch._foreach_mul(momentum_bufs, momentum) - torch._foreach_add_(grads, nesterov_terms) - - -_compiled_momentum: dict[bool, callable] = {} -_use_momentum_compile = True - - -def set_momentum_compile(enabled: bool): - """Toggle torch.compile for batched momentum.""" - global _use_momentum_compile - _use_momentum_compile = enabled - - -def batch_pre_ortho( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, - nesterov: bool, -) -> None: - """Batched momentum update on lists of plain tensors. - - Mirrors dion's ``muon_update_pre_orthogonalize``. - Inputs must be plain CUDA tensors (not DTensor). - Modifies ``momentum_bufs`` and (for nesterov) ``grads`` in-place. - - When compile is enabled, uses separately compiled functions for - nesterov=True/False to avoid graph breaks from the branch. - """ - fn = _batch_momentum_nesterov if nesterov else _batch_momentum - if _use_momentum_compile: - if nesterov not in _compiled_momentum: - _compiled_momentum[nesterov] = torch.compile(fn) - fn = _compiled_momentum[nesterov] - fn(grads, momentum_bufs, momentum) - - -def _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay): - """Weight-decay + update on plain tensors. - - Not compiled: per-param @torch.compile caused ~0.25ms TorchDynamo cache - lookup per call × 256+ params = massive overhead. The pipeline path uses - batched _foreach_* ops instead; this function remains for base() and - distributed_muon(). - """ - p_data.mul_(1 - lr * weight_decay) - p_data.add_(u_data, alpha=-adjusted_lr) - - -def update_p(p, u, lr, adjusted_lr, weight_decay): - """Apply weight decay and orthogonalized update to parameter. - - Args: - p: Parameter (torch.nn.Parameter or DTensor). - u: Orthogonalized update tensor. - lr: Base learning rate. - adjusted_lr: Size-adjusted learning rate. - weight_decay: Weight decay coefficient. - """ - # Unwrap Parameter -> underlying data tensor. - p_data = p.data if isinstance(p, torch.nn.Parameter) else p - # Unwrap DTensor -> local CUDA tensor for compiled kernel. - if isinstance(p_data, DTensor): - p_data = p_data._local_tensor - u_data = u._local_tensor if isinstance(u, DTensor) else u - _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay) - - -def adjust_lr_for_muon(lr, param_shape): - """Scale learning rate based on parameter matrix dimensions. - - Args: - lr: Base learning rate. - param_shape: Shape of the parameter tensor. - - Returns: - Adjusted learning rate. - """ - A, B = param_shape[:2] - # We adjust the learning rate and weight decay based on the size of the parameter matrix - # as described in the paper - adjusted_ratio = 0.2 * math.sqrt(max(A, B)) - adjusted_lr = lr * adjusted_ratio - return adjusted_lr - - -def _match_key(parts, key): - """Check if key matches as contiguous components in parts. - - Single-component keys (e.g. "experts") match any single component. - Multi-component keys (e.g. "experts.w1") match as a contiguous subsequence. - """ - key_parts = key.split(".") - key_len = len(key_parts) - if key_len == 1: - return key in parts - return any(parts[i:i + key_len] == key_parts - for i in range(len(parts) - key_len + 1)) - - -def is_expert_param(name, expert_keys): - """Check if a parameter name matches any expert key (component-level).""" - if not expert_keys: - return False - parts = normalize_fqn(name).split(".") - return any(_match_key(parts, key) for key in expert_keys) - - -def default_is_muon(name, x, expert_keys=None): - normalized = normalize_fqn(name) - parts = normalized.split(".") - skip_keys = [ - "embed_tokens", - "lm_head", - "tok_embeddings", - "output", - "mhc_attn", - "mhc_ffn", - "lambda_proj", - ] - if any(key in parts for key in skip_keys): - logger.info( - "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d", - normalized, name, x.ndim) - return False - effective_ndim = x.ndim - is_expert = is_expert_param(name, expert_keys) - if is_expert: - effective_ndim -= 1 - result = effective_ndim >= 2 - logger.info( - "[is_muon] %s (orig: %s): ndim=%d, expert=%s, effective_ndim=%d → %s", - normalized, name, x.ndim, is_expert, effective_ndim, - "Muon" if result else "AdamW") - return result - - -def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None): - if is_muon_func is None: - is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys) - - muon_params, muon_names = [], [] - non_muon_params, non_muon_names = [], [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - non_muon_names.append(n) - - logger.info("[param_groups] expert_keys=%s, Muon=%d, AdamW=%d", - expert_keys, len(muon_names), len(non_muon_names)) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/cpu_offload.py b/build/torch28-cxx11-rocm64-x86_64-linux/cpu_offload.py deleted file mode 100644 index fb5e69154a1d4a6c884491413a37a9acf0f66c80..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/cpu_offload.py +++ /dev/null @@ -1,206 +0,0 @@ -"""CPU offloading for optimizer states. - -Manages a pinned CPU memory pool and async CUDA streams to offload -optimizer state tensors (momentum buffers, Adam moments) to CPU between -optimizer steps, freeing GPU memory. - -All tracked tensors are packed into a single flat pinned CPU buffer -(per dtype). D2H and H2D copies are performed per-tensor directly -between individual GPU tensors and their slice of the CPU flat buffer -— no GPU staging buffer is allocated, so there is **no temporary GPU -memory spike** during offload or reload. - -Individual tensor storages are freed after offload via -``untyped_storage().resize_(0)``, preserving tensor identity so -downstream caches remain valid. -""" - -import logging -from collections import defaultdict - -import torch -from torch.distributed.tensor import DTensor - -logger = logging.getLogger(__name__) - - -class CPUOffloadPool: - """Pinned CPU memory pool for async optimizer state offloading. - - Tracked tensors are grouped by dtype. Each group gets a single flat - pinned CPU buffer. D2H / H2D copies are per-tensor (into slices of - the flat buffer) to avoid allocating a GPU staging buffer. - """ - - def __init__(self): - self._managed: list[torch.Tensor] = [] - self._storage_nbytes: dict[int, int] = {} # id(t) → bytes - - # Per-dtype group: populated on first offload. - # dtype → dict with keys: - # "indices" : list[int] managed-list indices - # "offsets" : list[tuple[int,int]] (start, numel) in flat buf - # "total" : int total numel - # "cpu_flat" : Tensor pinned CPU buffer - self._groups: dict[torch.dtype, dict] = {} - - self._offload_stream: torch.cuda.Stream | None = None - self._device: torch.device | None = None - self._initialized: bool = False - self._logged: bool = False - - # ------------------------------------------------------------------ - @staticmethod - def _local(t: torch.Tensor) -> torch.Tensor: - """Unwrap DTensor to its local CUDA tensor.""" - return t._local_tensor if isinstance(t, DTensor) else t - - def _ensure_stream(self): - if self._offload_stream is None: - self._offload_stream = torch.cuda.Stream(device=self._device) - - # ------------------------------------------------------------------ - def track(self, tensor: torch.Tensor): - """Register a GPU tensor for CPU offloading. Idempotent.""" - tid = id(tensor) - if tid in self._storage_nbytes: - return - local = self._local(tensor) - if self._device is None: - self._device = local.device - storage = local.untyped_storage() - # Skip tensors with empty storage (e.g. empty FSDP shards) - if storage.size() == 0: - return - self._storage_nbytes[tid] = storage.size() - self._managed.append(tensor) - - # ------------------------------------------------------------------ - def _init_buffers(self): - """Build per-dtype flat buffers on first offload.""" - # Group managed tensors by dtype. - dtype_map: dict[torch.dtype, list[tuple[int, int]]] = defaultdict(list) - for idx, t in enumerate(self._managed): - local = self._local(t) - dtype_map[local.dtype].append((idx, local.numel())) - - total_cpu_bytes = 0 - for dtype, entries in dtype_map.items(): - offsets: list[tuple[int, int]] = [] - indices: list[int] = [] - off = 0 - for idx, n in entries: - indices.append(idx) - offsets.append((off, n)) - off += n - cpu_flat = torch.empty(off, dtype=dtype, device="cpu", pin_memory=True) - self._groups[dtype] = { - "indices": indices, - "offsets": offsets, - "total": off, - "cpu_flat": cpu_flat, - } - total_cpu_bytes += off * cpu_flat.element_size() - - self._initialized = True - logger.info( - "[CPUOffload] Pool initialized: %d tensors, %d dtype group(s), " - "%.2f MB pinned CPU memory", - len(self._managed), - len(self._groups), - total_cpu_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def offload(self): - """Per-tensor async D2H into CPU flat buffer, then free GPU storage.""" - if not self._managed: - return - if not self._initialized: - self._init_buffers() - self._ensure_stream() - - # Offload stream waits for compute to finish. - compute_event = torch.cuda.current_stream(self._device).record_event() - self._offload_stream.wait_event(compute_event) - - offloaded_bytes = 0 - - # Per-tensor D2H copies directly into CPU flat buffer slices. - # No GPU staging buffer → no temporary GPU memory spike. - with torch.cuda.stream(self._offload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - cpu_flat[off : off + n].copy_(local.reshape(-1), non_blocking=True) - - offloaded_bytes += grp["total"] * cpu_flat.element_size() - - # Wait for all D2H copies to land, then free GPU storage. - self._offload_stream.synchronize() - for t in self._managed: - storage = self._local(t).untyped_storage() - if storage.size() != 0: - storage.resize_(0) - else: - raise RuntimeError( - f"Tensor storage is already freed (size=0) before offload. " - f"This indicates a double-free or external interference. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - - if not self._logged: - logger.info( - "[CPUOffload] Offloaded %.2f MB (GPU → CPU)", - offloaded_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def reload(self): - """Per-tensor H2D from CPU flat buffer on the default stream. - - Runs on the current (default) CUDA stream to avoid stream - interaction issues with the parallel Muon pipeline. Since - pinned CPU memory is the source, the copies overlap with - GPU idle time between steps. - """ - if not self._managed or not self._initialized: - return - - reloaded_bytes = 0 - - # Re-allocate all GPU storages first. - for t in self._managed: - local = self._local(t) - storage = local.untyped_storage() - if storage.size() != 0: - raise RuntimeError( - f"Storage should have been freed (size=0) before reload, " - f"but got size={storage.size()}. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - storage.resize_(self._storage_nbytes[id(t)]) - - # Per-tensor H2D copies from CPU flat buffer slices. - # non_blocking=True with pinned source allows DMA overlap. - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off : off + n], non_blocking=True) - - reloaded_bytes += grp["total"] * cpu_flat.element_size() - - if not self._logged: - logger.info( - "[CPUOffload] Reloaded %.2f MB (CPU → GPU)", reloaded_bytes / (1024**2) - ) diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/distributed/utils.py b/build/torch28-cxx11-rocm64-x86_64-linux/distributed/utils.py deleted file mode 100644 index 683096bcd50489e7a09959753fdbc29ba0f3b6a6..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/distributed/utils.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import (Placement, Shard, - _StridedShard) - - -def _is_shard(placement: Placement) -> bool: - """Check if a placement is a shard type (Shard or _StridedShard). - - In PyTorch 2.10+, _StridedShard no longer inherits from Shard, so - ``placement.is_shard()`` returns False for _StridedShard. This helper - handles both old and new hierarchies. - """ - return isinstance(placement, (Shard, _StridedShard)) - - -def get_slices_of_dtensor( - target: DTensor | torch.Tensor, - local_rank: int, - shard_mesh: DeviceMesh, - shard_placements: tuple[Placement], -) -> tuple[slice | torch.Tensor, ...]: - """ - Get per-dimension indices for a given rank's shard of the target tensor. - - Uses ``Shard.local_shard_size_and_offset`` and - ``_StridedShard.local_shard_size_and_offset`` for correct handling of - both contiguous and strided (non-contiguous) sharding. - - Args: - target (DTensor | torch.Tensor): The target tensor (for its shape). - local_rank (int): The local rank within the shard group. - shard_mesh (DeviceMesh): The shard mesh (only shard dimensions). - shard_placements (tuple[Placement]): The shard placements. - - Returns: - A tuple of indices (one per tensor dim). Each element is either: - - A ``slice`` (for contiguous or unsharded dims) - - A 1-D ``torch.LongTensor`` of indices (for strided sharding) - """ - - # find the global rank of the local rank in the shard mesh - rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank] - - rank_coords = (shard_mesh.mesh == rank).nonzero() - - assert len(rank_coords) == 1 - rank_coords = tuple(rank_coords[0].tolist()) - - assert len(rank_coords) == len(shard_placements) - - # Track per-shard-dim indices. - # None means "not yet sharded on this dim". - dim_indices: dict[int, torch.Tensor] = {} - - # Caution: Assuming replicate-to-shard of the shard mesh goes with - # left-to-right sharding. This is ensured by the sorting logic of - # construct_shard_mesh function. - for mesh_dim_idx, (rank_coord, placement) in enumerate( - zip(rank_coords, shard_placements)): - assert _is_shard(placement) - - num_chunks = shard_mesh.mesh.shape[mesh_dim_idx] - shard_dim = placement.dim - - # Current effective size on this dim (may already be sub-sharded) - if shard_dim in dim_indices: - curr_size = len(dim_indices[shard_dim]) - else: - curr_size = target.size()[shard_dim] - - # Compute indices for this level of sharding - if isinstance(placement, _StridedShard): - _shard_size, offsets = _StridedShard.local_shard_size_and_offset( - placement, - curr_size, - num_chunks, - rank_coord, - return_first_offset=False) - new_indices = torch.tensor(offsets, dtype=torch.long) - else: - shard_size, offset = Shard.local_shard_size_and_offset( - curr_size, num_chunks, rank_coord) - new_indices = torch.arange(offset, - offset + shard_size, - dtype=torch.long) - - # Compose with previous indices on this dim - if shard_dim in dim_indices: - dim_indices[shard_dim] = dim_indices[shard_dim][new_indices] - else: - dim_indices[shard_dim] = new_indices - - # Build result tuple - result: list[slice | torch.Tensor] = [] - for d in range(len(target.size())): - if d not in dim_indices: - result.append(slice(None)) - else: - indices = dim_indices[d] - # Convert contiguous indices to slice for efficiency - if len(indices) > 0: - start = indices[0].item() - expected = torch.arange(start, - start + len(indices), - dtype=torch.long) - if torch.equal(indices, expected): - result.append(slice(start, start + len(indices))) - else: - result.append(indices) - else: - result.append(slice(0, 0)) - - return tuple(result) - - -_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, - ProcessGroup]] = dict() - - -def construct_shard_mesh( - placements: tuple[Placement], - mesh: DeviceMesh, -) -> tuple[DeviceMesh, ProcessGroup, tuple[Placement, ...]]: - """Construct shard sub-mesh and ProcessGroup for all-to-all communication. - - Given a DTensor's placements and device mesh, extracts the "shard group" - — the set of ranks that together hold all shards of the same replica — - and creates a ProcessGroup for all-to-all among them. - - Steps: - 1. Sort placements: Replicate first, then Shard by (dim, granularity). - 2. Permute the mesh tensor to match the sorted order. - 3. Collapse Replicate dims → list of shard sub-meshes (one per replica). - 4. Create/retrieve a cached ProcessGroup for the current rank's sub-mesh. - - Example — 8 GPUs, mesh shape (2, 2, 2), - placements ``[Shard(0), Replicate, _StridedShard(0)]``:: - - Step 1 — Sort: [Replicate, _StridedShard(0), Shard(0)] - Permutation: [1, 2, 0] - - Step 2 — Permute mesh dims by [1, 2, 0]: - Original: Permuted: - [[[0,1],[2,3]], [[[0,2],[1,3]], - [[4,5],[6,7]]] [[4,6],[5,7]]] - - Step 3 — Unbind replicate dim (dim 0), giving 2 shard sub-meshes: - sub-mesh 0 = [[0,2],[1,3]] (replica group 0) - sub-mesh 1 = [[4,6],[5,7]] (replica group 1) - shard_placements = (_StridedShard(0), Shard(0)) - - Step 4 — Rank 0 → ProcessGroup([0,1,4,5]) - Rank 2 → ProcessGroup([2,3,6,7]) - - Returns: - ``(shard_mesh, process_group, shard_placements)`` - """ - my_rank = dist.get_rank() - assert mesh.mesh.device.type == 'cpu' - - # -- Fast path: 1D all-shard mesh → reuse existing PG. ---------------- - # Reuses the mesh's existing ProcessGroup directly, avoiding the - # overhead of dist.new_group(). The standard path below also handles - # subset calls safely via use_local_synchronization=True, but this - # fast path is still beneficial for the common 1D shard case. - if mesh.ndim == 1 and len(placements) == 1 and _is_shard(placements[0]): - key = (*mesh.mesh.shape, *mesh.mesh.flatten().tolist()) - if key not in _ranks_to_dist_cache: - _ranks_to_dist_cache[key] = (mesh, mesh.get_group()) - return (*_ranks_to_dist_cache[key], tuple(placements)) - - mesh_tensor = mesh.mesh.clone() - - # -- Step 1: Sort placements (Replicate first, then Shard by dim). ------ - # _StridedShard comes BEFORE regular Shard on the same dim so that - # get_slices_of_dtensor applies the outer sharding first, matching - # DTensor's left-to-right (outer-to-inner) composition order. - def _sort_key(item): - index, placement = item - assert not placement.is_partial(), "Partial placement not supported" - if placement.is_replicate(): - return (-1, 0, index) - assert _is_shard(placement), f"Unsupported: {type(placement)}" - split = (-1 / placement.split_factor if isinstance( - placement, _StridedShard) else 0) - return (placement.dim, split, index) - - indexed = sorted(enumerate(placements), key=_sort_key) - perm, sorted_placements = zip(*indexed) - - # -- Step 2: Permute mesh to match sorted placement order. -------------- - sorted_mesh = mesh_tensor.permute(perm) - - # -- Step 3: Collapse replicate dims → list of shard sub-meshes. -------- - # E.g. mesh (2, 3, 4, 4) with [R, R, S(0), S(1)] → 6 sub-meshes of (4, 4) - num_rep = sum(1 for p in sorted_placements if p.is_replicate()) - if num_rep > 0: - if num_rep > 1: - sorted_mesh = sorted_mesh.flatten(0, num_rep - 1) - shard_meshes = list(torch.unbind(sorted_mesh, dim=0)) - else: - shard_meshes = [sorted_mesh] - shard_placements = sorted_placements[num_rep:] - assert len(shard_placements) == len(set(shard_placements)) - - # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. -- - # Each rank only creates the group it belongs to, using - # use_local_synchronization=True so that only group members need to - # coordinate. This avoids deadlocks when different PP stages call - # construct_shard_mesh for different parameters. - def _cache_key(t: torch.Tensor) -> tuple: - return (*t.shape, *t.flatten().tolist()) - - my_key = None - for sm in shard_meshes: - if (my_rank == sm).any().item(): - key = _cache_key(sm) - assert my_key is None, "Rank appears in multiple shard groups" - my_key = key - if key not in _ranks_to_dist_cache: - pg = dist.new_group(sm.flatten().tolist(), - use_local_synchronization=True) - _ranks_to_dist_cache[key] = ( - DeviceMesh(device_type="cuda", mesh=sm), - pg, - ) - - return (*_ranks_to_dist_cache[my_key], shard_placements) diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/matmul_transpose_triton.py b/build/torch28-cxx11-rocm64-x86_64-linux/matmul_transpose_triton.py deleted file mode 100644 index 792de23d82c3fb45fe33d397ab9b76a0787259d0..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/matmul_transpose_triton.py +++ /dev/null @@ -1,122 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], - restore_value=['y'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -@torch.library.custom_op("muon::matmul_transpose_assign", - mutates_args=("d_out", )) -def matmul_transpose_assign(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """Compute d_out = d_in @ d_in.T using an optimized Triton kernel.""" - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -@matmul_transpose_assign.register_fake -def _(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """FakeTensor impl: d_out is already allocated, mutation is declared.""" - pass diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/metadata.json b/build/torch28-cxx11-rocm64-x86_64-linux/metadata.json deleted file mode 100644 index c55a35717622f1dd5c8ba376ea3a814cbcc10d78..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/metadata.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/muon.py b/build/torch28-cxx11-rocm64-x86_64-linux/muon.py deleted file mode 100644 index 14c0e22471fa6d47a51ed95e0e0c341dc18d5194..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/muon.py +++ /dev/null @@ -1,1068 +0,0 @@ -import logging -import types -from collections import defaultdict -from typing import Any - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Replicate, Shard -from torch.profiler import record_function - -from .adamw import _placement_cache, _tensor_cache, step_adamw -from .async_utils import run_pipeline -from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho, - get_default_muon_param_groups, is_expert_param, update_p) -from .cpu_offload import CPUOffloadPool -from .distributed.utils import (_is_shard, construct_shard_mesh, - get_slices_of_dtensor) -from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO, - _zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5_batched) -from .pipeline import muon_chunk_pipeline, prelaunch_first_gather -from .qk_clip import compute_scales, get_qk_clip_info, qk_clip - -logger = logging.getLogger(__name__) - - -def _expand_expert_params(names, params, expert_keys): - """Expand expert params by splitting on dim 0 (expert dimension). - - Params whose name matches any key in ``expert_keys`` are treated as - expert-parallel tensors. Their outermost dimension is the expert - dimension: an ``(E, out, in)`` tensor becomes ``E`` separate 2D - ``nn.Parameter`` views so that in-place updates propagate back to - the original storage. - - Non-expert params with ``ndim > 2`` trigger an ``AssertionError`` — - if they are expert params, their key must be added to ``expert_keys``. - - The grad must already be set on each expert param (e.g. after momentum). - - For DTensor expert params, placements that shard on dim 0 (expert dim) - are consumed by the split. Non-dim-0 shard placements (e.g. TP) are - preserved: each 2D slice is wrapped as a DTensor on the corresponding - submesh so the parallel pipeline handles the TP communication. - """ - expanded_names = [] - expanded_params = [] - - for n, p in zip(names, params): - is_expert = is_expert_param(n, expert_keys) - is_dtensor = isinstance(p.data, DTensor) - - if is_expert: - if is_dtensor: - logger.debug( - "[expand_expert] %s: expert DTensor, shape=%s, " - "placements=%s, mesh=%s, local_shape=%s", n, p.shape, - p.placements, p.device_mesh.mesh_dim_names, - p.to_local().shape) - else: - logger.debug( - "[expand_expert] %s: expert plain tensor, shape=%s", n, - p.data.shape) - - if not is_expert: - assert p.data.ndim <= 2, ( - f"Param {n} has ndim={p.data.ndim} but does not match " - f"expert_keys={expert_keys}. If this is an expert param, " - f"add its key to expert_keys.") - expanded_names.append(n) - expanded_params.append(p) - continue - - g = p.grad - assert g is not None, ( - f"Expert param {n} must have grad set before expansion") - - tp_mesh = None - tp_placements_2d = None - - if is_dtensor: - local_data = p.to_local() - local_grad = g.to_local() if isinstance(g, DTensor) else g - - # Find non-dim-0 shard placements (e.g. TP sharding). - # After splitting on dim 0, Shard(k) becomes Shard(k-1). - tp_dim_indices = [] - tp_placements_2d = [] - for i, pl in enumerate(p.placements): - if _is_shard(pl) and pl.dim != 0: - tp_dim_indices.append(i) - tp_placements_2d.append(Shard(pl.dim - 1)) - - if tp_dim_indices: - tp_dim_names = tuple(p.device_mesh.mesh_dim_names[i] - for i in tp_dim_indices) - if len(tp_dim_names) == 1: - tp_mesh = p.device_mesh[tp_dim_names[0]] - else: - tp_mesh = p.device_mesh[tp_dim_names] - else: - local_data = p.data - local_grad = g - - # Expand: split dim 0, reshape each slice to 2D. - num_local_experts = local_data.shape[0] - for i in range(num_local_experts): - slice_data = local_data[i] - slice_grad = local_grad[i] - - if tp_mesh is not None: - # Wrap as DTensor on TP submesh so the pipeline handles - # TP communication (gather/scatter across TP ranks). - dt_data = DTensor.from_local(slice_data, - device_mesh=tp_mesh, - placements=tp_placements_2d) - dt_grad = DTensor.from_local(slice_grad, - device_mesh=tp_mesh, - placements=tp_placements_2d) - expert_param = torch.nn.Parameter(dt_data, requires_grad=False) - expert_param.grad = dt_grad - else: - expert_param = torch.nn.Parameter(slice_data, - requires_grad=False) - expert_param.grad = slice_grad - - expanded_names.append(f"{n}[{i}]") - expanded_params.append(expert_param) - - p.grad = None # allow expert grad storage to be freed after pipeline - - return expanded_names, expanded_params - - -class Muon(torch.optim.Optimizer): - """ - Muon - MomentUm Orthogonalized by Newton-schulz - - Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- - processing step, in which each 2D parameter's update is replaced with the nearest orthogonal - matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has - the advantage that it can be stably run in bfloat16 on the GPU. - - Some warnings: - - We believe this optimizer is unlikely to work well for training with small batch size. - - We believe it may not work well for finetuning pretrained models, but we haven't tested this. - - Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. - lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) - momentum: The momentum used by the internal SGD. (0.95 is a good default) - nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) - ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. - Parameters that are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW instead. - adamw_lr: The learning rate for the internal AdamW. - adamw_betas: The betas for the internal AdamW. - adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - warmup_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher warmup_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. - chunk_size : Batch size of parameters to process in each - all2all gather/compute/scatter step. - Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified. - use_distributed_muon: Use distributed muon by Liu et al. (2024). - For testing purpose only. - expert_keys: List of strings to identify expert-parallel parameters. - If any key appears in a parameter's name, its outermost - dimension is treated as the expert dimension and expanded - into per-expert 2D params for Muon. For example, - ``expert_keys=["experts"]`` matches any param whose name - contains "experts". 3D+ params not matched by any key - will raise an error. - """ - - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config=None, - warmup_step=5, - chunk_size=-1, - use_distributed_muon=False, - expert_keys=None): - defaults = dict( - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - nesterov=nesterov, - ns_steps=ns_steps, - adamw_betas=adamw_betas, - adamw_eps=adamw_eps, - none_grad=none_grad, - use_muon=True, - ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) - super().__init__(params, defaults) - - self.debug = debug - self.clip_config = clip_config if clip_config is not None else { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100, - } - self.warmup_step = warmup_step - self.chunk_size = chunk_size - self.use_distributed_muon = use_distributed_muon - self.expert_keys = expert_keys - self.cpu_offload = False - self._cpu_offload_pool: CPUOffloadPool | None = None - self._offload_initialized = False - self._parallel_cache: dict[tuple[str, ...], dict] = {} - self._expert_expand_cache: dict[tuple[int, ...], dict] = {} - - def _calc_flops(self, G, steps): - assert len(G.shape) == 2 - M, N = G.shape - if M > N: - M, N = N, M - - return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) - - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - p.placements, p.device_mesh) - - return shard_mesh, shard_pg, shard_placements - - def init_state_and_assign_params(self, names, params, group, qk_logits): - param_to_state = {} - param_to_flops = {} - - total_flops = 0 - for p in params: - g = p.grad - if g is None: - continue - assert g.ndim == 2, "Muon only supports 2D parameters." - - flops = self._calc_flops(g, group["ns_steps"]) - param_to_flops[id(p)] = flops - total_flops += flops - - if self.debug: - logger.debug("Total TFLOPs for Muon: %.2f TFLOPs", - total_flops / 1e12) - - paired = list(zip(names, params)) - - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) - - round_robin = 0 - mesh = ordered_params[0].device_mesh - placements = ordered_params[0].placements - - shard_mesh, shard_pg, shard_placements = self.get_shard_mesh( - ordered_params[0]) - shard_mesh_flattened = shard_mesh.mesh.flatten() - num_ranks = dist.get_world_size(group=shard_pg) - - for n, p in zip(ordered_names, ordered_params): - if mesh != p.device_mesh: - raise ValueError("All parameters must be on the same mesh.") - if placements != p.placements: - raise ValueError("All parameters must have same placements.") - - worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks - round_robin = (round_robin + 1) % len(shard_mesh_flattened) - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - # Precompute per-rank indices and numels for all-to-all. - rank_indices: dict[int, tuple] = {} - rank_numels: dict[int, int] = {} - for r in range(num_ranks): - indices = get_slices_of_dtensor(p, r, shard_mesh, - shard_placements) - rank_indices[r] = indices - numel = 1 - for idx, dim_size in zip(indices, p.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) - rank_numels[r] = numel - - param_to_state[id(p)] = _muon_state( - worker_rank=worker_rank, - process_group=shard_pg, - rank_indices=rank_indices, - rank_numels=rank_numels, - name=n, - qk_clip_state=qk_clip_state, - ) - - return param_to_state, ordered_params - - def base(self, names, params, group, lr, weight_decay, qk_logits): - # Momentum is already applied by _step_muon before this method. - for n, p in zip(names, params): - g = p.grad - if g is None: - continue - - u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) - - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - def distributed_muon( - self, - names: list[str], - params: list[torch.nn.Parameter], - group: dict[str, Any], - lr: float, - weight_decay: float, - qk_logits: list[torch.Tensor | DTensor] | None, - ): - """Batched Distributed Muon — for testing/correctness verification only. - - Uses all-gather to reconstruct full tensors, computes Newton-Schulz on - the full grad, then slices back to local shards. This is simpler but - slower than the parallel pipeline (all2all) path, so it serves as a - reference implementation for verifying correctness. - """ - with record_function("distributed_muon"): - # Momentum is already applied by _step_muon before this method. - ns_steps = group["ns_steps"] - - # Separate plain tensors (no communication) from DTensors. - plain_names, plain_params = [], [] - dtensor_names, dtensor_params = [], [] - for n, p in zip(names, params): - if p.grad is None: - continue - if isinstance(p.data, DTensor): - dtensor_names.append(n) - dtensor_params.append(p) - else: - plain_names.append(n) - plain_params.append(p) - - # Process plain tensors per-param (no communication). - for n, p in zip(plain_names, plain_params): - u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE), - steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, - qk_logits) - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - if not dtensor_params: - return - - # Group DTensors by (placements, mesh) for batched all-gather. - placement_groups: dict[tuple, - tuple[list, - list]] = defaultdict(lambda: ([], [])) - for n, p in zip(dtensor_names, dtensor_params): - key = (p.placements, p.device_mesh) - placement_groups[key][0].append(n) - placement_groups[key][1].append(p) - - logger.info( - "distributed_muon: %d placement groups, %d total dtensors", - len(placement_groups), len(dtensor_params)) - - for (placements, mesh), (grp_names, - grp_params) in placement_groups.items(): - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - placements, mesh) - rank = dist.get_rank(shard_pg) - world_size = dist.get_world_size(shard_pg) - - logger.info(" group: %d params, placements=%s, world_size=%d", - len(grp_params), placements, world_size) - - # Separate params that can be batched (all shard dims evenly - # divisible) from those needing per-param full_tensor - # (e.g. MoE gate weights with fewer rows than shard ranks). - # all_gather_into_tensor requires equal buffer sizes across - # ranks, so uneven splits must use DTensor full_tensor(). - batch_names, batch_params = [], [] - single_names, single_params = [], [] - for n, p in zip(grp_names, grp_params): - even = all(p.shape[pl.dim] % - shard_mesh.mesh.shape[dim_idx] == 0 - for dim_idx, pl in enumerate(shard_placements)) - if even: - batch_names.append(n) - batch_params.append(p) - else: - single_names.append(n) - single_params.append(p) - - # Process uneven-split params per-param via full_tensor(). - for n, p in zip(single_names, single_params): - with record_function("distributed_muon::newton_schulz"): - g_full = p.grad.full_tensor().to(COMM_DTYPE) - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - if not batch_params: - continue - - logger.info(" batched=%d, single=%d", len(batch_params), - len(single_params)) - - # Concat all local grad shards into a single flat buffer. - with record_function("distributed_muon::gather"): - grad_locals = [ - p.grad.to_local().to(COMM_DTYPE).flatten() - for p in batch_params - ] - numels = [g.numel() for g in grad_locals] - grad_concat = torch.cat(grad_locals) - del grad_locals - - # Single all-gather (replaces N separate full_tensor). - grad_gathered = torch.empty( - grad_concat.numel() * world_size, - dtype=COMM_DTYPE, - device="cuda", - ) - dist.all_gather_into_tensor(grad_gathered, - grad_concat, - group=shard_pg) - - total_numel = grad_concat.numel() - del grad_concat - - # Precompute per-param offsets within the concat buffer. - offsets = [] - off = 0 - for ne in numels: - offsets.append(off) - off += ne - - # Per-param: reconstruct full grad → NS → local update. - for i, (n, p) in enumerate(zip(batch_names, batch_params)): - with record_function("distributed_muon::newton_schulz"): - g_full = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - for r in range(world_size): - r_start = r * total_numel + offsets[i] - shard = grad_gathered[r_start:r_start + numels[i]] - indices = get_slices_of_dtensor( - p, r, shard_mesh, shard_placements) - g_full[indices] = shard.reshape( - g_full[indices].shape) - - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - def _setup_parallel(self, names, params, group, qk_logits): - """Compute (or retrieve cached) parallel pipeline metadata. - - Returns: - (ordered_params, param_to_state, rank, chunk_size) - """ - cache_key = tuple(names) - - if cache_key not in self._parallel_cache: - # First call: compute metadata and populate cache. - param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - shard_pg = param_to_state[id(ordered_params[0])].process_group - rank = dist.get_rank(group=shard_pg) - - if self.chunk_size == -1: - shard_ranks = dist.get_world_size(shard_pg) - chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO - elif self.chunk_size > 0: - chunk_size = self.chunk_size - else: - raise ValueError( - "chunk_size must be -1 or a positive integer.") - - ordered_names = [ - param_to_state[id(p)].name for p in ordered_params - ] - name_to_state = { - param_to_state[id(p)].name: param_to_state[id(p)] - for p in ordered_params - } - self._parallel_cache[cache_key] = { - 'ordered_names': ordered_names, - 'name_to_state': name_to_state, - 'rank': rank, - 'chunk_size': chunk_size, - } - else: - # Cached path: rebuild param_to_state with current id(p) keys. - cache = self._parallel_cache[cache_key] - rank = cache['rank'] - chunk_size = cache['chunk_size'] - - name_to_param = dict(zip(names, params)) - ordered_params = [name_to_param[n] for n in cache['ordered_names']] - - param_to_state = {} - for p, n in zip(ordered_params, cache['ordered_names']): - cached_state = cache['name_to_state'][n] - param_to_state[id(p)] = _muon_state( - worker_rank=cached_state.worker_rank, - process_group=cached_state.process_group, - rank_indices=cached_state.rank_indices, - rank_numels=cached_state.rank_numels, - name=n, - qk_clip_state=get_qk_clip_info(self.clip_config, n, - qk_logits), - ) - - return ordered_params, param_to_state, rank, chunk_size - - def parallel(self, - names, - params, - group, - lr, - weight_decay, - qk_logits, - prelaunch_gather=None): - """ - Perform a parallel optimization step using Muon. - - Parameters are chunked and each chunk is processed by a - :func:`muon_chunk_pipeline` generator. :func:`run_pipeline` - interleaves multiple chunks so that communication and computation - overlap across chunks (the same overlap previously achieved by the - warmup + main-loop index scheduling). - - If ``prelaunch_gather`` is provided, it is passed to the first - chunk's generator to skip re-launching the already in-flight - A2A gather. - """ - - # Momentum is already applied by _step_muon before this method. - - ordered_params, param_to_state, rank, chunk_size = ( - self._setup_parallel(names, params, group, qk_logits)) - - def pipelines(): - first = True - for start in range(0, len(ordered_params), chunk_size): - chunk = ordered_params[start:start + chunk_size] - if chunk: - kwargs = dict( - params=chunk, - param_to_state=param_to_state, - rank=rank, - ns_steps=group["ns_steps"], - lr=lr, - weight_decay=weight_decay, - none_grad=group["none_grad"], - ) - if first and prelaunch_gather is not None: - kwargs['prelaunch_gather'] = prelaunch_gather - first = False - yield muon_chunk_pipeline(**kwargs) - - with record_function("muon::pipeline"): - run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1) - - def _step_muon(self, group, qk_logits=None): - params = group["params"] - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - # Apply momentum to all params before routing/expansion. - # Batched using _foreach_* ops (compiled, fullgraph=True). - with record_function("muon::momentum"): - active_params = [p for p in params if p.grad is not None] - if active_params: - # Ensure momentum buffers exist (avoid zeros_like when already present). - for p in active_params: - if "momentum_buffer" not in self.state[p]: - self.state[p]["momentum_buffer"] = torch.zeros_like( - p.grad) - - # Extract local tensors for compiled batch function. - local_grads = [ - p.grad._local_tensor - if isinstance(p.grad, DTensor) else p.grad - for p in active_params - ] - local_bufs = [ - self.state[p]["momentum_buffer"]._local_tensor - if isinstance(self.state[p]["momentum_buffer"], DTensor) - else self.state[p]["momentum_buffer"] - for p in active_params - ] - - # Wrap momentum as tensor for torch.compile. - batch_pre_ortho(local_grads, local_bufs, - torch.tensor(momentum), group["nesterov"]) - - # For non-nesterov, the result is the momentum buffer. - if not group["nesterov"]: - for p in active_params: - p.grad = self.state[p]["momentum_buffer"] - - # Identify batched experts for deferred NS. - # Detection is cheap (condition checks only); actual NS compute is - # deferred so it can overlap with the first chunk's A2A gather. - deferred_expert_work = [] - if self.expert_keys: - batched_expert_indices = [] - for i, (n, p) in enumerate(zip(names, params)): - if not (is_expert_param(n, self.expert_keys) - and p.grad is not None): - continue - # Eligible: plain tensor, or DTensor with no non-dim-0 shards. - if isinstance(p.data, DTensor): - has_tp = any( - _is_shard(pl) and pl.dim != 0 for pl in p.placements) - if has_tp: - continue - batched_expert_indices.append(i) - - if batched_expert_indices: - # Save refs for deferred NS; free grads from param list. - for i in batched_expert_indices: - p = params[i] - g = p.grad - local_g = (g._local_tensor - if isinstance(g, DTensor) else g) - local_data = (p.data._local_tensor if isinstance( - p.data, DTensor) else p.data) - deferred_expert_work.append((local_data, local_g)) - p.grad = None - - # Remove batched experts from lists before expansion. - keep = sorted( - set(range(len(params))) - set(batched_expert_indices)) - names = [names[i] for i in keep] - params = [params[i] for i in keep] - - def _run_deferred_expert_ns(): - """Execute deferred batched expert NS.""" - if not deferred_expert_work: - return - with record_function("muon::batched_expert_ns"): - ns_steps = group["ns_steps"] - for local_data, local_g in deferred_expert_work: - u = zeropower_via_newtonschulz5_batched( - local_g.to(COMM_DTYPE), steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:]) - local_data.mul_(1 - lr * weight_decay) - local_data.add_(u, alpha=-adjusted_lr) - - # Expand expert params by splitting on dim 0. - logger.debug("[_step_muon] before expand: %d params, expert_keys=%s", - len(params), self.expert_keys) - if self.expert_keys: - cache_key = tuple(id(p) for p in params) - cache = self._expert_expand_cache.get(cache_key) - - if cache is None: - # Cold path: full expansion + build cache metadata. - exp_names, exp_params = _expand_expert_params( - names, params, self.expert_keys) - - # Build per-expert-group info for hot-path grad updates. - grad_info = [] - exp_idx = 0 - for orig_idx, (n, p) in enumerate(zip(names, params)): - if not is_expert_param(n, self.expert_keys): - exp_idx += 1 - continue - - is_dt = isinstance(p.data, DTensor) - num_experts = (p.to_local() if is_dt else p.data).shape[0] - - # Detect TP mesh from the first expanded expert param. - tp_mesh = None - tp_pls = None - sample = exp_params[exp_idx] - if isinstance(sample.data, DTensor): - tp_mesh = sample.data.device_mesh - tp_pls = list(sample.data.placements) - - grad_info.append((orig_idx, num_experts, exp_idx, is_dt, - tp_mesh, tp_pls)) - exp_idx += num_experts - - self._expert_expand_cache[cache_key] = { - 'names': exp_names, - 'params': exp_params, - 'grad_info': grad_info, - } - names, params = exp_names, exp_params - else: - # Hot path: reuse cached params, only update expert grads. - for (orig_idx, num_experts, exp_start, is_dt, tp_mesh, - tp_pls) in cache['grad_info']: - p = params[orig_idx] - g = p.grad - local_grad = (g.to_local() - if is_dt and isinstance(g, DTensor) else g) - for i in range(num_experts): - expert_p = cache['params'][exp_start + i] - sg = local_grad[i] - if tp_mesh is not None: - expert_p.grad = DTensor.from_local( - sg, device_mesh=tp_mesh, placements=tp_pls) - else: - expert_p.grad = sg - p.grad = None - - names = cache['names'] - params = cache['params'] - else: - names, params = _expand_expert_params(names, params, - self.expert_keys) - logger.debug("[_step_muon] after expand: %d params", len(params)) - - param_dtensors = [] - name_dtensors = [] - - param_tensors = [] - name_tensors = [] - - # distributed_muon is a reference implementation for testing only. - # The parallel pipeline (all2all) path below is the production path. - if self.use_distributed_muon: - _run_deferred_expert_ns() - self.distributed_muon(names=names, - params=params, - group=group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits) - return - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - logger.debug( - "[route] %s → base (DTensor all-Replicate), " - "shape=%s, placements=%s", n, p.shape, p.placements) - param_tensors.append(p) - name_tensors.append(n) - else: - logger.debug( - "[route] %s → parallel (DTensor), shape=%s, " - "placements=%s, mesh=%s", n, p.shape, p.placements, - p.device_mesh.mesh_dim_names) - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - logger.debug("[route] %s → base (plain tensor), shape=%s", n, - p.data.shape) - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError(f"Unsupported parameter type: {type(p.data)}") - - logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, " - f"{len(param_tensors)} Tensors → base") - - def group_dtensors(dtensors, names): - # To support different placements, we group parameters by placements - # and run parallel Muon on each group. - - placement_to_params = defaultdict(lambda: ([], [])) - - assert len(dtensors) == len(names) - for p, n in zip(dtensors, names): - placement_to_params[tuple([p.placements, - p.device_mesh])][0].append(n) - placement_to_params[tuple([p.placements, - p.device_mesh])][1].append(p) - return placement_to_params - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - dtensor_group = group_dtensors(param_dtensors, name_dtensors) - - # Pre-launch the first chunk's A2A gather so that the NCCL - # communication overlaps with the (deferred) batched expert NS - # compute on the default CUDA stream. - prelaunch = None - if deferred_expert_work: - first_names, first_params = next(iter(dtensor_group.values())) - ordered, pts, rnk, csz = self._setup_parallel( - first_names, first_params, group, qk_logits) - first_chunk = ordered[:csz] - if first_chunk: - prelaunch = prelaunch_first_gather(first_chunk, pts, rnk, - group["none_grad"]) - - _run_deferred_expert_ns() - - first_group = True - for _, (names, params) in dtensor_group.items(): - pg = prelaunch if first_group else None - first_group = False - self.parallel( - names, - params, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - prelaunch_gather=pg, - ) - else: - _run_deferred_expert_ns() - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - ) - - def _register_states_for_offload(self): - """Register all optimizer state tensors with the CPU offload pool. - - Called once after the first step when states have been lazily created. - Offloads all param states (momentum buffers for Muon, moment1/moment2 - for AdamW) to free GPU memory between steps. - """ - pool = self._cpu_offload_pool - tracked = 0 - for group in self.param_groups: - for p in group["params"]: - if p not in self.state: - continue - state = self.state[p] - if group.get("use_muon", False): - if "momentum_buffer" in state: - pool.track(state["momentum_buffer"]) - tracked += 1 - else: - if "moment1" in state: - pool.track(state["moment1"]) - if "moment2" in state: - pool.track(state["moment2"]) - tracked += 1 - logger.info("[CPUOffload] Registered %d param states for offload", - tracked) - - @torch.no_grad - def step(self, closure=None, qk_logits=None): - """Perform a single optimization step. - - Args: - closure (Callable, optional): A closure that reevaluates the model - and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). - """ - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # H2D: reload optimizer states from CPU before computation. - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload() - - logger.debug("[Muon.step] expert_keys=%s, %d param groups", - self.expert_keys, len(self.param_groups)) - - for i, group in enumerate(self.param_groups): - if group["use_muon"]: - logger.debug("[Muon.step] group %d: use_muon=True, %d params", - i, len(group["params"])) - self._step_muon(group, qk_logits=qk_logits) - else: - logger.debug( - "[Muon.step] group %d: use_muon=False (AdamW), %d params", - i, len(group["params"])) - step_adamw(self.state, group) - - # D2H: offload optimizer states to CPU after computation. - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - return loss - - # ------------------------------------------------------------------ - # CPU offload public helpers - # ------------------------------------------------------------------ - - def turn_on_cpu_offload(self): - """Enable CPU offload for optimizer states.""" - if self.cpu_offload: - return - logger.info("[Muon] turn_on_cpu_offload") - self.cpu_offload = True - if not self.state: - return - self._cpu_offload_pool = CPUOffloadPool() - self._offload_initialized = False - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_off_cpu_offload(self): - """Disable CPU offload and keep optimizer states resident on GPU.""" - if not self.cpu_offload: - return - logger.info("[Muon] turn_off_cpu_offload") - if self._offload_initialized: - self._cpu_offload_pool.reload() - torch.cuda.current_stream().synchronize() - self._cpu_offload_pool = None - self._offload_initialized = False - self.cpu_offload = False - - # ------------------------------------------------------------------ - # Checkpoint support for cpu_offload - # ------------------------------------------------------------------ - - def state_dict(self) -> dict: - if self.cpu_offload: - raise RuntimeError( - "Muon.state_dict() requires turn_off_cpu_offload() before checkpoint save." - ) - return super().state_dict() - - def load_state_dict(self, state_dict: dict) -> None: - if self.cpu_offload: - raise RuntimeError( - "Muon.load_state_dict() requires turn_off_cpu_offload() before checkpoint load." - ) - super().load_state_dict(state_dict) - - # Invalidate adamw.py's module-level tensor caches so that - # the next step rebuilds them with the newly loaded state tensors. - _placement_cache.clear() - _tensor_cache.clear() diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/newton_schulz.py b/build/torch28-cxx11-rocm64-x86_64-linux/newton_schulz.py deleted file mode 100644 index d939264b69a34e7a3fa78859f34dc265a1159d59..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/newton_schulz.py +++ /dev/null @@ -1,240 +0,0 @@ -from itertools import repeat -from math import inf, sqrt - -import numpy as np -import torch - -from .matmul_transpose_triton import matmul_transpose_assign - -COMM_DTYPE = torch.bfloat16 -DEFAULT_CHUNK_SIZE_RATIO = 4 - - -def _optimal_quintic(l, u, max_iter=1000): - """ - Use the simplified Remez algorithm to find the optimal odd quintic approximant - to the constant function x -> 1 over the interval [l, u]. - - Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum - approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the - two interior equioscillation nodes q, r until convergence. Returns the - closed-form equioscillating solution when l ≈ u. - - Raises ValueError if any intermediate value (a, b, c, E, q, r) is non-finite - (NaN or inf). Raises RuntimeError if convergence is not reached within - max_iter iterations. - """ - assert 0 <= l <= u - if 1 - 5e-6 <= l / u: - return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5) - q = (3 * l + u) / 4 - r = (l + 3 * u) / 4 - E = inf - for _ in range(max_iter): - old_E = E - LHS = np.array( - [ - [l, l**3, l**5, 1], - [q, q**3, q**5, -1], - [r, r**3, r**5, 1], - [u, u**3, u**5, -1], - ] - ) - a, b, c, E = np.linalg.solve(LHS, np.ones(4)) - if not np.all(np.isfinite([a, b, c, E])): - raise ValueError( - f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}" - ) - q, r = np.sqrt( - (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / (10 * c) - ) - if not np.all(np.isfinite([q, r])): - raise ValueError(f"_optimal_quintic: non-finite node update q={q}, r={r}") - if abs(old_E - E) <= 1e-15: - break - else: - raise RuntimeError( - f"_optimal_quintic: did not converge after {max_iter} iterations" - ) - return float(a), float(b), float(c) - - -def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0): - """ - Compute the Polar Express coefficient series for `num_iters` quintic iterations. - - Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that - compose to map singular values from [l, 1] toward 1. At each step: - 1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion` - prevents near-zero singular values from stalling by raising the effective - lower bound; if it is active (cushion*u > l), the coefficients are - rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u]. - 2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the - last iteration, providing numerical headroom at the cost of a slightly slower - final convergence step. - 3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1). - - Returns a list of (a, b, c) tuples, one per iteration. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - u = 1 - assert 0 <= l <= u - safety_factor = 1 + safety_factor_eps - coefficients = [] - for iter in range(num_iters): - a, b, c = _optimal_quintic(max(l, cushion * u), u) - if cushion * u > l: - pl = a * l + b * l**3 + c * l**5 - pu = a * u + b * u**3 + c * u**5 - rescaler = 2 / (pl + pu) - a *= rescaler - b *= rescaler - c *= rescaler - if iter < num_iters - 1: - a /= safety_factor - b /= safety_factor**3 - c /= safety_factor**5 - coefficients.append((a, b, c)) - l = a * l + b * l**3 + c * l**5 - u = 2 - l - return coefficients - - -# Precomputed Polar Express coefficients (a, b, c) for 10 quintic Newton-Schulz -# iterations. Each tuple is the minimax-optimal (Remez/equioscillation) odd quintic -# approximant to x->1 over the current singular-value interval, computed once at -# import time and reused across all optimizer steps. -# -# Contrast with the former hardcoded NS coefficients (5 fixed tuples): -# - Former: empirically tuned to maximize slope at zero; did not converge -# singular values to 1, yielding US'V^T with S' ~ Uniform(0.5, 1.5) instead -# of the true polar factor UV^T. -# - Polar Express: analytically optimal per step, adapting to the shrinking -# singular-value interval [l, u] as iterations progress; converges all -# singular values to 1, producing the exact polar factor UV^T. -_coeffs_list = _optimal_composition( - l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02 -) - - -# This code is adapted from: -# KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py) -# NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress) -# matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon) -@torch.no_grad() -def _zeropower_via_newtonschulz5(G, steps): - """ - Compute the polar factor of G via the Polar Express method. - - Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c) - are the Polar Express coefficients from `_coeffs_list`. Each step is the - optimal odd quintic approximant to x -> 1 over the current singular-value - interval, minimizing the maximum approximation error (Remez / minimax criterion). - The composition maps singular values from [l, 1] to near 1, producing the - polar factor (orthogonal factor in the polar decomposition G = UP). - - `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2, - cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE - X = G # no manual typecast - - if G.size(0) > G.size(1): - X = X.T - - X = X / (X.norm() + 1e-7) - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - # Perform the NS iterations - for a, b, c in hs: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(0) > G.size(1): - X = X.T - - return X - - -@torch.no_grad() -def _zeropower_via_newtonschulz5_batched(G, steps): - """Batched polar factor computation for 3D (E, out, in) tensors. - - Same algorithm as ``_zeropower_via_newtonschulz5`` but uses - ``torch.bmm`` / ``torch.baddbmm`` instead of the 2D Triton kernel, - processing all E expert matrices in a single batched call. - """ - assert len(G.shape) == 3 - assert G.dtype == COMM_DTYPE - X = G - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - # Per-expert Frobenius norm. - X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) - - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - for a, b, c in hs: - buf1 = torch.bmm(X, X.transpose(-2, -1)) - buf2 = torch.bmm(buf1, buf1.transpose(-2, -1)) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.baddbmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - return X - - -_ns_per_shape: dict[tuple[int, ...], callable] = {} -_use_compile = True - - -def set_ns_compile(enabled: bool): - """Toggle torch.compile for Newton-Schulz iteration.""" - global _use_compile - _use_compile = enabled - - -def zeropower_via_newtonschulz5(G, steps=5): - if not _use_compile: - return _zeropower_via_newtonschulz5(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() - - -def zeropower_via_newtonschulz5_batched(G, steps=5): - """Compile-cached batched Newton-Schulz for 3D expert tensors.""" - if not _use_compile: - return _zeropower_via_newtonschulz5_batched(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile( - _zeropower_via_newtonschulz5_batched, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/optimizer/__init__.py b/build/torch28-cxx11-rocm64-x86_64-linux/optimizer/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/optimizer/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/pipeline.py b/build/torch28-cxx11-rocm64-x86_64-linux/pipeline.py deleted file mode 100644 index c0c2d515856182d8d15ad27dd4e4e093b29397d6..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/pipeline.py +++ /dev/null @@ -1,468 +0,0 @@ -import logging -from typing import Generator - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -from .core import _muon_state, adjust_lr_for_muon -from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5 -from .qk_clip import compute_scales - -logger = logging.getLogger(__name__) - -# ====================================================================== -# Stage helpers -# ====================================================================== - - -def _launch_gather( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Allocate gather buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_gather``). - gathered_grads: ``{id(p): empty_tensor}`` for owned params, - ``None`` for non-owned. - recv_counts: Per-source-rank element counts. - """ - # Allocate gathered-grad buffers - gathered_grads: dict[int, torch.Tensor | None] = {} - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - gathered_grads[id(p)] = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - else: - gathered_grads[id(p)] = None - - # Build send buffer – batch grad copies via torch.cat - # (1-2 fused kernels vs N individual narrow().copy_() calls). - send_counts = [0] * num_ranks - for p in params: - state = param_to_state[id(p)] - send_counts[state.worker_rank] += state.rank_numels[rank] - - total_send = sum(send_counts) - if total_send > 0: - # Group grad slices by destination rank in a single pass. - dst_to_grads = [[] for _ in range(num_ranks)] - for p in params: - state = param_to_state[id(p)] - n = state.rank_numels[rank] - if n > 0: - g = p.grad.to_local() - dst_to_grads[state.worker_rank].append(g.reshape(-1)) - - # Flatten in dst order and cat once. - all_slices = [] - for dst in range(num_ranks): - all_slices.extend(dst_to_grads[dst]) - send_buf = torch.cat(all_slices) - if send_buf.dtype != COMM_DTYPE: - send_buf = send_buf.to(COMM_DTYPE) - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += state.rank_numels[src] - recv_counts[src] = total - - recv_buf = torch.empty(sum(recv_counts), dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - logger.debug(f"send_buf size: {send_buf.numel()}, " - f"recv_buf size: {recv_buf.numel()}, " - f"recv_counts: {recv_counts}, " - f"send_counts: {send_counts}, " - f"process_group: {str(process_group)}") - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, gathered_grads, recv_counts - - -def _complete_gather( - recv_buf: torch.Tensor, - recv_counts: list[int], - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - param_to_state: dict[int, _muon_state], - rank: int, -) -> None: - """Reconstruct gathered grads from the recv buffer (in-place).""" - off = 0 - for src in range(len(recv_counts)): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - - indices = state.rank_indices[src] - - shard_view = gathered_grads[id(p)][indices] - n = shard_view.numel() - if n == 0: - continue - - sg = recv_buf.narrow(0, off + inner_off, n) - sg = sg.reshape(shard_view.shape) - gathered_grads[id(p)][indices] = sg - - inner_off += n - assert inner_off == block - off += block - - -def _compute_ns( - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - ns_steps: int, -) -> dict[int, torch.Tensor | None]: - """Run Newton-Schulz orthogonalization on owned parameters. - - Returns: - computed_us: ``{id(p): orthogonalized_update}`` for owned params. - """ - computed_us: dict[int, torch.Tensor | None] = {} - for p in owned_params: - u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps) - gathered_grads[id(p)] = None # free gathered grad - computed_us[id(p)] = u - return computed_us - - -def _launch_scatter( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, - computed_us: dict[int, torch.Tensor | None], -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor], list[int]]: - """Allocate scatter buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_scatter``). - scattered_us: Empty dict, populated by ``_complete_scatter`` with - zero-copy views into ``recv_buf``. - recv_counts: Per-source-rank element counts. - """ - # scattered_us is populated by _complete_scatter with zero-copy views - # into recv_buf, avoiding N empty_like allocations + N copy_ calls. - # Pre-seed entries for params whose local shard is empty (rank_numels == 0) - # so _update_params can iterate all params without KeyError. - scattered_us: dict[int, torch.Tensor] = {} - for p in params: - if param_to_state[id(p)].rank_numels[rank] == 0: - scattered_us[id(p)] = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - # Build send buffer – batch via torch.cat - # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls). - send_counts = [0] * num_ranks - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - for dst_rank in range(num_ranks): - send_counts[dst_rank] += state.rank_numels[dst_rank] - - total_send = sum(send_counts) - if total_send > 0: - # Cache u_full conversions to avoid redundant .to() per dst_rank. - u_fulls = {} - for p in owned_params: - u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous() - - # Collect slices in dst order (matches all-to-all send layout). - all_slices = [] - for dst_rank in range(num_ranks): - for p in owned_params: - state = param_to_state[id(p)] - su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten() - if su.numel() > 0: - all_slices.append(su) - - send_buf = torch.cat(all_slices) if all_slices else torch.empty( - 0, dtype=COMM_DTYPE, device="cuda") - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += state.rank_numels[rank] - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, scattered_us, recv_counts - - -def _complete_scatter( - recv_buf: torch.Tensor, - recv_counts: list[int], - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], -) -> None: - """Populate scattered_us with zero-copy views into recv_buf. - - Instead of pre-allocating tensors and copying, we assign views directly - from ``recv_buf``. This eliminates N ``empty_like`` + N ``copy_`` calls. - The underlying storage of ``recv_buf`` is kept alive through the views - until ``scattered_us`` is cleared after ``_update_params``. - """ - off = 0 - for src in range(len(recv_counts)): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = state.rank_numels[rank] - if n == 0: - continue - - scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - - inner_off += n - - assert inner_off == block - off += block - - -def _update_params( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], - lr: float, - weight_decay: float, -) -> None: - """Apply weight decay, Muon update, and optional QK clipping. - - Uses batched ``_foreach_mul_`` for weight decay and batched - ``_foreach_add_`` for the Muon update, grouping parameters by - adjusted_lr to minimize kernel launches while preserving float32 - precision for the alpha scaling. - """ - if not params: - return - - # Batched weight decay: p *= (1 - lr * wd) — single fused kernel. - p_locals = [p._local_tensor for p in params] - torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay) - - # Group params by adjusted_lr so _foreach_add_ can use a single - # alpha per group (preserves float32 precision for alpha scaling). - lr_groups: dict[float, tuple[list, list]] = {} - for p in params: - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - if adjusted_lr not in lr_groups: - lr_groups[adjusted_lr] = ([], []) - lr_groups[adjusted_lr][0].append(p._local_tensor) - lr_groups[adjusted_lr][1].append(scattered_us[id(p)]) - - for adjusted_lr, (p_group, u_group) in lr_groups.items(): - torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr) - - # QK clipping – applied directly on the local tensor to - # avoid DTensor sharding-propagation issues with _StridedShard. - for p in params: - state = param_to_state[id(p)] - if state.qk_clip_state is None: - continue - scales_full = compute_scales(p, state.qk_clip_state) - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = state.rank_indices[rank][0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - -# ====================================================================== -# Pre-launch helper for overlapping first chunk's gather with other work. -# ====================================================================== - - -@torch.no_grad() -def prelaunch_first_gather( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - none_grad: bool, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Launch the first chunk's A2A gather early for overlap with other compute. - - Call this *before* expensive GPU work (e.g. batched expert NS) so that - the NCCL all-to-all runs concurrently on the NCCL stream while the - default stream executes compute. - - Returns the same 4-tuple that ``_launch_gather`` produces, which should - be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - with record_function("muon::prelaunch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - return work, recv_buf, gathered_grads, recv_counts - - -# ====================================================================== -# Main generator – thin orchestrator that wires stages together. -# ====================================================================== - - -@torch.no_grad() -def muon_chunk_pipeline( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - ns_steps: int, - lr: float, - weight_decay: float, - none_grad: bool, - prelaunch_gather: tuple | None = None, -) -> Generator[None, None, None]: - """Process one chunk of parameters through the full Muon pipeline. - - Stages: gather -> compute (Newton-Schulz) -> scatter -> update. - - Each ``yield`` lets :func:`run_pipeline` interleave other chunks so - that communication and computation overlap across chunks. Async - communication is launched via ``async_op=True`` and completed after - the yield with ``work.wait()``. - - Overlap happens because :func:`run_pipeline` admits one new chunk - per iteration (staggered admission). While chunk *N* does NS - compute on the default CUDA stream, chunk *N+1*'s async all-to-all - runs concurrently on the NCCL stream — no separate ``comm_stream`` - is required. - - If ``prelaunch_gather`` is provided, the gather was already launched - by :func:`prelaunch_first_gather` and we skip launching it again. - - Yields exactly **2** times: - - 1. After launching async all-to-all gather (or immediately if pre-launched). - 2. After launching async all-to-all scatter. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - if prelaunch_gather is not None: - # Gather was pre-launched; none_grad already handled by caller. - work, recv_buf, gathered_grads, recv_counts = prelaunch_gather - else: - # Normal path: launch async gather. - with record_function("muon::launch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - yield # --- YIELD 1: other chunks can launch their gather --- - - with record_function("muon::wait_gather"): - work.wait() - _complete_gather(recv_buf, recv_counts, owned_params, gathered_grads, - param_to_state, rank) - del recv_buf - - # Stage 3: Newton-Schulz orthogonalization. - with record_function("muon::newton_schulz"): - computed_us = _compute_ns(owned_params, gathered_grads, ns_steps) - gathered_grads.clear() - - # Stages 4-5: launch async scatter. - with record_function("muon::launch_scatter"): - work, recv_buf, scattered_us, recv_counts = _launch_scatter( - params, owned_params, param_to_state, rank, num_ranks, - process_group, computed_us) - computed_us.clear() - - yield # --- YIELD 2: other chunks can launch their scatter --- - - with record_function("muon::wait_scatter"): - work.wait() - _complete_scatter(recv_buf, recv_counts, params, param_to_state, rank, - scattered_us) - del recv_buf - - # Stage 6: apply parameter updates. - with record_function("muon::update_params"): - _update_params(params, param_to_state, rank, scattered_us, lr, - weight_decay) - scattered_us.clear() diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/qk_clip.py b/build/torch28-cxx11-rocm64-x86_64-linux/qk_clip.py deleted file mode 100644 index 2aba711b3004b7f09e7141da7ef834bd61cc2430..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/qk_clip.py +++ /dev/null @@ -1,198 +0,0 @@ -import logging -import math -from dataclasses import dataclass - -import torch -from torch.distributed.tensor import DTensor - -from .core import normalize_fqn - -logger = logging.getLogger(__name__) - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - and return (kind, layer_index). - - Supported kinds: - MHA/GQA: 'wq', 'wk', 'q_proj', 'k_proj' - MLA: 'wq_b' (Q up-proj), 'wkv_b' (KV up-proj) - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.1.attn.wq_b.weight' -> ('wq_b', 1) - 'model.0.attn.wkv_b.weight' -> ('wkv_b', 0) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = normalize_fqn(name).split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: str | None # 'wq'/'q_proj'/'wq_b' or 'wk'/'k_proj'/'wkv_b' or None - indices: list[int] # which heads to consider for clipping - head_dim: int # from config (qk_head_dim for MLA wq_b) - threshold: float # from config - logit: torch.Tensor | None - - # MLA-specific fields - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 - - -def get_qk_clip_info(clip_config, n, qk_logits): - """Extract QK clipping info for a named parameter. - - Args: - clip_config: QK clipping configuration dict (or None). - MHA/GQA keys: head_dim, threshold, q_indices, k_indices - MLA extra keys: is_mla=True, qk_nope_head_dim, qk_rope_head_dim, v_head_dim - n: Parameter name string. - qk_logits: Dict mapping layer indices to logit tensors (or None). - - Returns: - QKClipInfo instance with clipping configuration for this parameter. - """ - if clip_config is None: - return None - - head_dim = clip_config.get('head_dim') - threshold = clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - # In TP settings, qk_logits may be DTensor - # We convert it to full tensor here for simplicity - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - -def compute_scales(p, qk_clip_state): - """Compute per-head scaling factors for QK clipping. - - Returns scales tensor (√γ per head) if any head exceeds threshold, else None. - For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim. - """ - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - # Check if any head exceeds threshold before allocating. - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - - if not head_scales: - return None - - # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows - if qk_clip_state.is_mla and kind == 'wkv_b': - effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim - else: - effective_head_dim = head_dim - - H_global = p.shape[0] // effective_head_dim - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale - return scales_full - - -def qk_clip(p, scales, info): - """Apply per-head scaling to a Q/K projection weight matrix. - - Args: - p: Parameter (nn.Parameter or raw tensor). - scales: [n_heads] tensor, each element = √γ_h. - info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions. - - MLA sub-region scaling per Algorithm 1 (MuonClip): - wq_b: q_nope rows → √γ, q_pe rows → γ - wkv_b: k_nope rows → √γ, v rows → unchanged - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not info.is_mla: - # MHA/GQA: uniform √γ applied to all rows in each head - W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: vectorized sub-region scaling within each head - if info.kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_head_dim = qk_nope + info.qk_rope_head_dim - W_3d = W.view(-1, qk_head_dim, W.shape[1]) # [H, qk_head_dim, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # q_nope → √γ - W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1, - 1)) # q_pe → γ - - elif info.kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - W_3d = W.view(-1, kv_stride, W.shape[1]) # [H, kv_stride, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # k_nope → √γ - # v rows: not touched (k_R shared rotary unchanged) diff --git a/build/torch29-cxx11-cu126-x86_64-linux/__init__.py b/build/torch29-cxx11-cu126-x86_64-linux/__init__.py deleted file mode 100644 index 239c7a65f8293e7d0df28f05fce645af56d628c0..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .muon import Muon - -__all__ = [ - "Muon", -] diff --git a/build/torch29-cxx11-cu126-x86_64-linux/_ops.py b/build/torch29-cxx11-cu126-x86_64-linux/_ops.py deleted file mode 100644 index 034bff088659b5df6f6d401feb18c89dc5f33b29..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _optimizer_8d53b78_dirty -ops = torch.ops._optimizer_8d53b78_dirty - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_optimizer_8d53b78_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch29-cxx11-cu126-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so b/build/torch29-cxx11-cu126-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so deleted file mode 100755 index 8b49da51c7e1e432368980c185a0456d90b49f70..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7ba2d5675067b94d6327adc7d124fa9ac534af8236b6783981e13b47a5ee603b -size 1936664 diff --git a/build/torch29-cxx11-cu126-x86_64-linux/adamw.py b/build/torch29-cxx11-cu126-x86_64-linux/adamw.py deleted file mode 100644 index b5a95816a9f5b9e1889eaadae65373bfbced809a..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/adamw.py +++ /dev/null @@ -1,271 +0,0 @@ -import logging -from collections import defaultdict -from typing import cast - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -def fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: float | torch.Tensor, - weight_decay: float, - eps: float, - maximize: bool, -) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: dict | None = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) - - -def _to_local(t): - """Unwrap DTensor to local tensor for fused ops.""" - return t._local_tensor if isinstance(t, DTensor) else t - - -# --------------------------------------------------------------------------- -# Caches for eliminating per-step Python overhead. -# -# Placement grouping and tensor list assembly are identical every step -# (params don't change placement, moment/step tensors are the same objects -# after initialisation). We cache them keyed by id() of the param list -# stored in param_groups (stable across steps). -# -# Only gradients change each step and must be collected fresh. -# --------------------------------------------------------------------------- - -# id(group["params"]) → dict[placement_key, list[param]] -_placement_cache: dict[int, dict[tuple, list]] = {} - -# id(placement_group_list) → (params_local, moment1, moment2, state_steps) -_tensor_cache: dict[int, tuple[list, list, list, list]] = {} - - -def _step_adamw_params_slow(optimizer_state, params, group): - """Uncached fallback for the rare case where some params lack grads.""" - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - g = p.grad - if g is None: - continue - state = optimizer_state[p] - params_with_grads.append(_to_local(p)) - grads.append(_to_local(g)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - if not params_with_grads: - return - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw_params(optimizer_state, params, group): - """Run fused AdamW on a list of parameters sharing the same placement. - - After the first call, cached tensor lists (params_local, moment1, - moment2, state_steps) are reused — only gradients are collected fresh. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - params: List of parameters to update. - group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay. - """ - # Collect grads — the only thing that changes each step. - with record_function("adamw::collect_grads"): - grads = [] - for p in params: - g = p.grad - if g is None: - # Rare: fall back to slow path that filters per-param. - _step_adamw_params_slow(optimizer_state, params, group) - return - grads.append(_to_local(g)) - - tensor_key = id(params) - if tensor_key not in _tensor_cache: - with record_function("adamw::init_tensor_cache"): - params_local = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - state = optimizer_state[p] - params_local.append(_to_local(p)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(p.grad) - state["moment2"] = torch.zeros_like(p.grad) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - _tensor_cache[tensor_key] = (params_local, moment1, moment2, - state_steps) - - params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key] - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - with record_function("adamw::fused_adamw"): - fused_adamw( - params_local, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw(optimizer_state, group): - """Dispatch AdamW step, grouping parameters by type and placement. - - Placement grouping is cached after the first call since params never - change their placement between steps. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - group: Parameter group dict. - """ - params = group["params"] - placement_key = id(params) - - if placement_key not in _placement_cache: - with record_function("adamw::group_by_placement"): - placement_to_params: dict[tuple, - list[torch.Tensor]] = defaultdict(list) - for p in params: - match p: - case DTensor(): - logger.debug( - "[AdamW] DTensor param: shape=%s, placements=%s, " - "mesh=%s, grad=%s", p.shape, p.placements, - p.device_mesh.mesh_dim_names, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple( - [p.placements, p.device_mesh])].append(p) - case torch.Tensor(): - logger.debug( - "[AdamW] plain param: shape=%s, grad=%s", p.shape, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple([torch.Tensor, - None])].append(p) - - logger.debug("[AdamW] %d placement groups, %d total params", - len(placement_to_params), len(params)) - - _placement_cache[placement_key] = dict(placement_to_params) - - for group_params in _placement_cache[placement_key].values(): - step_adamw_params(optimizer_state, group_params, group) diff --git a/build/torch29-cxx11-cu126-x86_64-linux/async_utils.py b/build/torch29-cxx11-cu126-x86_64-linux/async_utils.py deleted file mode 100644 index a45c530ac9cad88e3555ec1047a6aa59f225347e..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/async_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import logging -from typing import Generator - -logger = logging.getLogger(__name__) - - -class _Task: - """Internal: wraps a generator, advances one yield at a time.""" - - def __init__(self, generator: Generator[None, None, None], index: int): - self._generator = generator - self._index = index - self._steps_completed = 0 - self.step() # run to first yield - - def step(self) -> bool: - try: - next(self._generator) - self._steps_completed += 1 - logger.debug("pipeline[%d] completed stage %d", self._index, - self._steps_completed) - return True - except StopIteration: - logger.debug("pipeline[%d] finished after %d stages", self._index, - self._steps_completed) - return False - - def close(self): - self._generator.close() - - -def run_pipeline( - pipelines: Generator[Generator[None, None, None], None, None], - max_concurrent: int, -) -> None: - """Run generator-based pipelines with bounded concurrency. - - Each pipeline is a generator that yields at stage boundaries. - The runtime interleaves pipelines so communication and computation - overlap across chunks. - """ - if max_concurrent <= 0: - raise ValueError(f"max_concurrent must be > 0, got {max_concurrent}") - - have_new = True - task_index = 0 - previous_tasks: list[_Task] = [] - - try: - while have_new or previous_tasks: - running_tasks: list[_Task] = [] - - # Admit one new pipeline per iteration (staggered admission). - # Admitting one at a time ensures that while chunk N does NS - # compute on the default stream, chunk N+1's NCCL all-to-all - # runs concurrently on the NCCL stream — creating real - # communication/computation overlap on the GPU. - if have_new and len(previous_tasks) < max_concurrent: - try: - gen = next(pipelines) - task = _Task(gen, task_index) - task_index += 1 - running_tasks.append(task) - except StopIteration: - have_new = False - - # Advance every previously-yielded task by one step. - for task in previous_tasks: - if task.step(): - running_tasks.append(task) - - previous_tasks = running_tasks - except BaseException: - # Clean up all in-flight generators to release GPU resources. - for task in previous_tasks: - task.close() - raise diff --git a/build/torch29-cxx11-cu126-x86_64-linux/core.py b/build/torch29-cxx11-cu126-x86_64-linux/core.py deleted file mode 100644 index c69d515afef305ad0ed66374095fa2d2468d99cc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/core.py +++ /dev/null @@ -1,219 +0,0 @@ -import logging -import math -from dataclasses import dataclass -from typing import List - -import torch -from torch.distributed import ProcessGroup -from torch.distributed.tensor import DTensor - -# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into -# parameter FQNs. Activation checkpointing similarly inserts -# "_checkpoint_wrapped_module". Strip these so name-based matching (skip_keys, -# expert_keys, QK layer parsing) works regardless of wrapper nesting. -_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) - -logger = logging.getLogger(__name__) - - -def normalize_fqn(name: str) -> str: - """Strip torch.compile / checkpoint wrapper components from a parameter FQN.""" - return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS) - - -@dataclass -class _muon_state: - worker_rank: int - process_group: ProcessGroup - rank_indices: dict[int, tuple] # local_rank -> per-dim indices - rank_numels: dict[int, int] # local_rank -> numel - name: str - qk_clip_state: torch.Tensor | None = None - - -def _batch_momentum( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update (no nesterov).""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - - -def _batch_momentum_nesterov( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update with nesterov correction.""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - nesterov_terms = torch._foreach_mul(momentum_bufs, momentum) - torch._foreach_add_(grads, nesterov_terms) - - -_compiled_momentum: dict[bool, callable] = {} -_use_momentum_compile = True - - -def set_momentum_compile(enabled: bool): - """Toggle torch.compile for batched momentum.""" - global _use_momentum_compile - _use_momentum_compile = enabled - - -def batch_pre_ortho( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, - nesterov: bool, -) -> None: - """Batched momentum update on lists of plain tensors. - - Mirrors dion's ``muon_update_pre_orthogonalize``. - Inputs must be plain CUDA tensors (not DTensor). - Modifies ``momentum_bufs`` and (for nesterov) ``grads`` in-place. - - When compile is enabled, uses separately compiled functions for - nesterov=True/False to avoid graph breaks from the branch. - """ - fn = _batch_momentum_nesterov if nesterov else _batch_momentum - if _use_momentum_compile: - if nesterov not in _compiled_momentum: - _compiled_momentum[nesterov] = torch.compile(fn) - fn = _compiled_momentum[nesterov] - fn(grads, momentum_bufs, momentum) - - -def _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay): - """Weight-decay + update on plain tensors. - - Not compiled: per-param @torch.compile caused ~0.25ms TorchDynamo cache - lookup per call × 256+ params = massive overhead. The pipeline path uses - batched _foreach_* ops instead; this function remains for base() and - distributed_muon(). - """ - p_data.mul_(1 - lr * weight_decay) - p_data.add_(u_data, alpha=-adjusted_lr) - - -def update_p(p, u, lr, adjusted_lr, weight_decay): - """Apply weight decay and orthogonalized update to parameter. - - Args: - p: Parameter (torch.nn.Parameter or DTensor). - u: Orthogonalized update tensor. - lr: Base learning rate. - adjusted_lr: Size-adjusted learning rate. - weight_decay: Weight decay coefficient. - """ - # Unwrap Parameter -> underlying data tensor. - p_data = p.data if isinstance(p, torch.nn.Parameter) else p - # Unwrap DTensor -> local CUDA tensor for compiled kernel. - if isinstance(p_data, DTensor): - p_data = p_data._local_tensor - u_data = u._local_tensor if isinstance(u, DTensor) else u - _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay) - - -def adjust_lr_for_muon(lr, param_shape): - """Scale learning rate based on parameter matrix dimensions. - - Args: - lr: Base learning rate. - param_shape: Shape of the parameter tensor. - - Returns: - Adjusted learning rate. - """ - A, B = param_shape[:2] - # We adjust the learning rate and weight decay based on the size of the parameter matrix - # as described in the paper - adjusted_ratio = 0.2 * math.sqrt(max(A, B)) - adjusted_lr = lr * adjusted_ratio - return adjusted_lr - - -def _match_key(parts, key): - """Check if key matches as contiguous components in parts. - - Single-component keys (e.g. "experts") match any single component. - Multi-component keys (e.g. "experts.w1") match as a contiguous subsequence. - """ - key_parts = key.split(".") - key_len = len(key_parts) - if key_len == 1: - return key in parts - return any(parts[i:i + key_len] == key_parts - for i in range(len(parts) - key_len + 1)) - - -def is_expert_param(name, expert_keys): - """Check if a parameter name matches any expert key (component-level).""" - if not expert_keys: - return False - parts = normalize_fqn(name).split(".") - return any(_match_key(parts, key) for key in expert_keys) - - -def default_is_muon(name, x, expert_keys=None): - normalized = normalize_fqn(name) - parts = normalized.split(".") - skip_keys = [ - "embed_tokens", - "lm_head", - "tok_embeddings", - "output", - "mhc_attn", - "mhc_ffn", - "lambda_proj", - ] - if any(key in parts for key in skip_keys): - logger.info( - "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d", - normalized, name, x.ndim) - return False - effective_ndim = x.ndim - is_expert = is_expert_param(name, expert_keys) - if is_expert: - effective_ndim -= 1 - result = effective_ndim >= 2 - logger.info( - "[is_muon] %s (orig: %s): ndim=%d, expert=%s, effective_ndim=%d → %s", - normalized, name, x.ndim, is_expert, effective_ndim, - "Muon" if result else "AdamW") - return result - - -def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None): - if is_muon_func is None: - is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys) - - muon_params, muon_names = [], [] - non_muon_params, non_muon_names = [], [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - non_muon_names.append(n) - - logger.info("[param_groups] expert_keys=%s, Muon=%d, AdamW=%d", - expert_keys, len(muon_names), len(non_muon_names)) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] diff --git a/build/torch29-cxx11-cu126-x86_64-linux/cpu_offload.py b/build/torch29-cxx11-cu126-x86_64-linux/cpu_offload.py deleted file mode 100644 index fb5e69154a1d4a6c884491413a37a9acf0f66c80..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/cpu_offload.py +++ /dev/null @@ -1,206 +0,0 @@ -"""CPU offloading for optimizer states. - -Manages a pinned CPU memory pool and async CUDA streams to offload -optimizer state tensors (momentum buffers, Adam moments) to CPU between -optimizer steps, freeing GPU memory. - -All tracked tensors are packed into a single flat pinned CPU buffer -(per dtype). D2H and H2D copies are performed per-tensor directly -between individual GPU tensors and their slice of the CPU flat buffer -— no GPU staging buffer is allocated, so there is **no temporary GPU -memory spike** during offload or reload. - -Individual tensor storages are freed after offload via -``untyped_storage().resize_(0)``, preserving tensor identity so -downstream caches remain valid. -""" - -import logging -from collections import defaultdict - -import torch -from torch.distributed.tensor import DTensor - -logger = logging.getLogger(__name__) - - -class CPUOffloadPool: - """Pinned CPU memory pool for async optimizer state offloading. - - Tracked tensors are grouped by dtype. Each group gets a single flat - pinned CPU buffer. D2H / H2D copies are per-tensor (into slices of - the flat buffer) to avoid allocating a GPU staging buffer. - """ - - def __init__(self): - self._managed: list[torch.Tensor] = [] - self._storage_nbytes: dict[int, int] = {} # id(t) → bytes - - # Per-dtype group: populated on first offload. - # dtype → dict with keys: - # "indices" : list[int] managed-list indices - # "offsets" : list[tuple[int,int]] (start, numel) in flat buf - # "total" : int total numel - # "cpu_flat" : Tensor pinned CPU buffer - self._groups: dict[torch.dtype, dict] = {} - - self._offload_stream: torch.cuda.Stream | None = None - self._device: torch.device | None = None - self._initialized: bool = False - self._logged: bool = False - - # ------------------------------------------------------------------ - @staticmethod - def _local(t: torch.Tensor) -> torch.Tensor: - """Unwrap DTensor to its local CUDA tensor.""" - return t._local_tensor if isinstance(t, DTensor) else t - - def _ensure_stream(self): - if self._offload_stream is None: - self._offload_stream = torch.cuda.Stream(device=self._device) - - # ------------------------------------------------------------------ - def track(self, tensor: torch.Tensor): - """Register a GPU tensor for CPU offloading. Idempotent.""" - tid = id(tensor) - if tid in self._storage_nbytes: - return - local = self._local(tensor) - if self._device is None: - self._device = local.device - storage = local.untyped_storage() - # Skip tensors with empty storage (e.g. empty FSDP shards) - if storage.size() == 0: - return - self._storage_nbytes[tid] = storage.size() - self._managed.append(tensor) - - # ------------------------------------------------------------------ - def _init_buffers(self): - """Build per-dtype flat buffers on first offload.""" - # Group managed tensors by dtype. - dtype_map: dict[torch.dtype, list[tuple[int, int]]] = defaultdict(list) - for idx, t in enumerate(self._managed): - local = self._local(t) - dtype_map[local.dtype].append((idx, local.numel())) - - total_cpu_bytes = 0 - for dtype, entries in dtype_map.items(): - offsets: list[tuple[int, int]] = [] - indices: list[int] = [] - off = 0 - for idx, n in entries: - indices.append(idx) - offsets.append((off, n)) - off += n - cpu_flat = torch.empty(off, dtype=dtype, device="cpu", pin_memory=True) - self._groups[dtype] = { - "indices": indices, - "offsets": offsets, - "total": off, - "cpu_flat": cpu_flat, - } - total_cpu_bytes += off * cpu_flat.element_size() - - self._initialized = True - logger.info( - "[CPUOffload] Pool initialized: %d tensors, %d dtype group(s), " - "%.2f MB pinned CPU memory", - len(self._managed), - len(self._groups), - total_cpu_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def offload(self): - """Per-tensor async D2H into CPU flat buffer, then free GPU storage.""" - if not self._managed: - return - if not self._initialized: - self._init_buffers() - self._ensure_stream() - - # Offload stream waits for compute to finish. - compute_event = torch.cuda.current_stream(self._device).record_event() - self._offload_stream.wait_event(compute_event) - - offloaded_bytes = 0 - - # Per-tensor D2H copies directly into CPU flat buffer slices. - # No GPU staging buffer → no temporary GPU memory spike. - with torch.cuda.stream(self._offload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - cpu_flat[off : off + n].copy_(local.reshape(-1), non_blocking=True) - - offloaded_bytes += grp["total"] * cpu_flat.element_size() - - # Wait for all D2H copies to land, then free GPU storage. - self._offload_stream.synchronize() - for t in self._managed: - storage = self._local(t).untyped_storage() - if storage.size() != 0: - storage.resize_(0) - else: - raise RuntimeError( - f"Tensor storage is already freed (size=0) before offload. " - f"This indicates a double-free or external interference. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - - if not self._logged: - logger.info( - "[CPUOffload] Offloaded %.2f MB (GPU → CPU)", - offloaded_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def reload(self): - """Per-tensor H2D from CPU flat buffer on the default stream. - - Runs on the current (default) CUDA stream to avoid stream - interaction issues with the parallel Muon pipeline. Since - pinned CPU memory is the source, the copies overlap with - GPU idle time between steps. - """ - if not self._managed or not self._initialized: - return - - reloaded_bytes = 0 - - # Re-allocate all GPU storages first. - for t in self._managed: - local = self._local(t) - storage = local.untyped_storage() - if storage.size() != 0: - raise RuntimeError( - f"Storage should have been freed (size=0) before reload, " - f"but got size={storage.size()}. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - storage.resize_(self._storage_nbytes[id(t)]) - - # Per-tensor H2D copies from CPU flat buffer slices. - # non_blocking=True with pinned source allows DMA overlap. - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off : off + n], non_blocking=True) - - reloaded_bytes += grp["total"] * cpu_flat.element_size() - - if not self._logged: - logger.info( - "[CPUOffload] Reloaded %.2f MB (CPU → GPU)", reloaded_bytes / (1024**2) - ) diff --git a/build/torch29-cxx11-cu126-x86_64-linux/distributed/utils.py b/build/torch29-cxx11-cu126-x86_64-linux/distributed/utils.py deleted file mode 100644 index 683096bcd50489e7a09959753fdbc29ba0f3b6a6..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/distributed/utils.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import (Placement, Shard, - _StridedShard) - - -def _is_shard(placement: Placement) -> bool: - """Check if a placement is a shard type (Shard or _StridedShard). - - In PyTorch 2.10+, _StridedShard no longer inherits from Shard, so - ``placement.is_shard()`` returns False for _StridedShard. This helper - handles both old and new hierarchies. - """ - return isinstance(placement, (Shard, _StridedShard)) - - -def get_slices_of_dtensor( - target: DTensor | torch.Tensor, - local_rank: int, - shard_mesh: DeviceMesh, - shard_placements: tuple[Placement], -) -> tuple[slice | torch.Tensor, ...]: - """ - Get per-dimension indices for a given rank's shard of the target tensor. - - Uses ``Shard.local_shard_size_and_offset`` and - ``_StridedShard.local_shard_size_and_offset`` for correct handling of - both contiguous and strided (non-contiguous) sharding. - - Args: - target (DTensor | torch.Tensor): The target tensor (for its shape). - local_rank (int): The local rank within the shard group. - shard_mesh (DeviceMesh): The shard mesh (only shard dimensions). - shard_placements (tuple[Placement]): The shard placements. - - Returns: - A tuple of indices (one per tensor dim). Each element is either: - - A ``slice`` (for contiguous or unsharded dims) - - A 1-D ``torch.LongTensor`` of indices (for strided sharding) - """ - - # find the global rank of the local rank in the shard mesh - rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank] - - rank_coords = (shard_mesh.mesh == rank).nonzero() - - assert len(rank_coords) == 1 - rank_coords = tuple(rank_coords[0].tolist()) - - assert len(rank_coords) == len(shard_placements) - - # Track per-shard-dim indices. - # None means "not yet sharded on this dim". - dim_indices: dict[int, torch.Tensor] = {} - - # Caution: Assuming replicate-to-shard of the shard mesh goes with - # left-to-right sharding. This is ensured by the sorting logic of - # construct_shard_mesh function. - for mesh_dim_idx, (rank_coord, placement) in enumerate( - zip(rank_coords, shard_placements)): - assert _is_shard(placement) - - num_chunks = shard_mesh.mesh.shape[mesh_dim_idx] - shard_dim = placement.dim - - # Current effective size on this dim (may already be sub-sharded) - if shard_dim in dim_indices: - curr_size = len(dim_indices[shard_dim]) - else: - curr_size = target.size()[shard_dim] - - # Compute indices for this level of sharding - if isinstance(placement, _StridedShard): - _shard_size, offsets = _StridedShard.local_shard_size_and_offset( - placement, - curr_size, - num_chunks, - rank_coord, - return_first_offset=False) - new_indices = torch.tensor(offsets, dtype=torch.long) - else: - shard_size, offset = Shard.local_shard_size_and_offset( - curr_size, num_chunks, rank_coord) - new_indices = torch.arange(offset, - offset + shard_size, - dtype=torch.long) - - # Compose with previous indices on this dim - if shard_dim in dim_indices: - dim_indices[shard_dim] = dim_indices[shard_dim][new_indices] - else: - dim_indices[shard_dim] = new_indices - - # Build result tuple - result: list[slice | torch.Tensor] = [] - for d in range(len(target.size())): - if d not in dim_indices: - result.append(slice(None)) - else: - indices = dim_indices[d] - # Convert contiguous indices to slice for efficiency - if len(indices) > 0: - start = indices[0].item() - expected = torch.arange(start, - start + len(indices), - dtype=torch.long) - if torch.equal(indices, expected): - result.append(slice(start, start + len(indices))) - else: - result.append(indices) - else: - result.append(slice(0, 0)) - - return tuple(result) - - -_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, - ProcessGroup]] = dict() - - -def construct_shard_mesh( - placements: tuple[Placement], - mesh: DeviceMesh, -) -> tuple[DeviceMesh, ProcessGroup, tuple[Placement, ...]]: - """Construct shard sub-mesh and ProcessGroup for all-to-all communication. - - Given a DTensor's placements and device mesh, extracts the "shard group" - — the set of ranks that together hold all shards of the same replica — - and creates a ProcessGroup for all-to-all among them. - - Steps: - 1. Sort placements: Replicate first, then Shard by (dim, granularity). - 2. Permute the mesh tensor to match the sorted order. - 3. Collapse Replicate dims → list of shard sub-meshes (one per replica). - 4. Create/retrieve a cached ProcessGroup for the current rank's sub-mesh. - - Example — 8 GPUs, mesh shape (2, 2, 2), - placements ``[Shard(0), Replicate, _StridedShard(0)]``:: - - Step 1 — Sort: [Replicate, _StridedShard(0), Shard(0)] - Permutation: [1, 2, 0] - - Step 2 — Permute mesh dims by [1, 2, 0]: - Original: Permuted: - [[[0,1],[2,3]], [[[0,2],[1,3]], - [[4,5],[6,7]]] [[4,6],[5,7]]] - - Step 3 — Unbind replicate dim (dim 0), giving 2 shard sub-meshes: - sub-mesh 0 = [[0,2],[1,3]] (replica group 0) - sub-mesh 1 = [[4,6],[5,7]] (replica group 1) - shard_placements = (_StridedShard(0), Shard(0)) - - Step 4 — Rank 0 → ProcessGroup([0,1,4,5]) - Rank 2 → ProcessGroup([2,3,6,7]) - - Returns: - ``(shard_mesh, process_group, shard_placements)`` - """ - my_rank = dist.get_rank() - assert mesh.mesh.device.type == 'cpu' - - # -- Fast path: 1D all-shard mesh → reuse existing PG. ---------------- - # Reuses the mesh's existing ProcessGroup directly, avoiding the - # overhead of dist.new_group(). The standard path below also handles - # subset calls safely via use_local_synchronization=True, but this - # fast path is still beneficial for the common 1D shard case. - if mesh.ndim == 1 and len(placements) == 1 and _is_shard(placements[0]): - key = (*mesh.mesh.shape, *mesh.mesh.flatten().tolist()) - if key not in _ranks_to_dist_cache: - _ranks_to_dist_cache[key] = (mesh, mesh.get_group()) - return (*_ranks_to_dist_cache[key], tuple(placements)) - - mesh_tensor = mesh.mesh.clone() - - # -- Step 1: Sort placements (Replicate first, then Shard by dim). ------ - # _StridedShard comes BEFORE regular Shard on the same dim so that - # get_slices_of_dtensor applies the outer sharding first, matching - # DTensor's left-to-right (outer-to-inner) composition order. - def _sort_key(item): - index, placement = item - assert not placement.is_partial(), "Partial placement not supported" - if placement.is_replicate(): - return (-1, 0, index) - assert _is_shard(placement), f"Unsupported: {type(placement)}" - split = (-1 / placement.split_factor if isinstance( - placement, _StridedShard) else 0) - return (placement.dim, split, index) - - indexed = sorted(enumerate(placements), key=_sort_key) - perm, sorted_placements = zip(*indexed) - - # -- Step 2: Permute mesh to match sorted placement order. -------------- - sorted_mesh = mesh_tensor.permute(perm) - - # -- Step 3: Collapse replicate dims → list of shard sub-meshes. -------- - # E.g. mesh (2, 3, 4, 4) with [R, R, S(0), S(1)] → 6 sub-meshes of (4, 4) - num_rep = sum(1 for p in sorted_placements if p.is_replicate()) - if num_rep > 0: - if num_rep > 1: - sorted_mesh = sorted_mesh.flatten(0, num_rep - 1) - shard_meshes = list(torch.unbind(sorted_mesh, dim=0)) - else: - shard_meshes = [sorted_mesh] - shard_placements = sorted_placements[num_rep:] - assert len(shard_placements) == len(set(shard_placements)) - - # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. -- - # Each rank only creates the group it belongs to, using - # use_local_synchronization=True so that only group members need to - # coordinate. This avoids deadlocks when different PP stages call - # construct_shard_mesh for different parameters. - def _cache_key(t: torch.Tensor) -> tuple: - return (*t.shape, *t.flatten().tolist()) - - my_key = None - for sm in shard_meshes: - if (my_rank == sm).any().item(): - key = _cache_key(sm) - assert my_key is None, "Rank appears in multiple shard groups" - my_key = key - if key not in _ranks_to_dist_cache: - pg = dist.new_group(sm.flatten().tolist(), - use_local_synchronization=True) - _ranks_to_dist_cache[key] = ( - DeviceMesh(device_type="cuda", mesh=sm), - pg, - ) - - return (*_ranks_to_dist_cache[my_key], shard_placements) diff --git a/build/torch29-cxx11-cu126-x86_64-linux/matmul_transpose_triton.py b/build/torch29-cxx11-cu126-x86_64-linux/matmul_transpose_triton.py deleted file mode 100644 index 792de23d82c3fb45fe33d397ab9b76a0787259d0..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/matmul_transpose_triton.py +++ /dev/null @@ -1,122 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], - restore_value=['y'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -@torch.library.custom_op("muon::matmul_transpose_assign", - mutates_args=("d_out", )) -def matmul_transpose_assign(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """Compute d_out = d_in @ d_in.T using an optimized Triton kernel.""" - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -@matmul_transpose_assign.register_fake -def _(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """FakeTensor impl: d_out is already allocated, mutation is declared.""" - pass diff --git a/build/torch29-cxx11-cu126-x86_64-linux/metadata.json b/build/torch29-cxx11-cu126-x86_64-linux/metadata.json deleted file mode 100644 index c55a35717622f1dd5c8ba376ea3a814cbcc10d78..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/metadata.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch29-cxx11-cu126-x86_64-linux/muon.py b/build/torch29-cxx11-cu126-x86_64-linux/muon.py deleted file mode 100644 index 14c0e22471fa6d47a51ed95e0e0c341dc18d5194..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/muon.py +++ /dev/null @@ -1,1068 +0,0 @@ -import logging -import types -from collections import defaultdict -from typing import Any - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Replicate, Shard -from torch.profiler import record_function - -from .adamw import _placement_cache, _tensor_cache, step_adamw -from .async_utils import run_pipeline -from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho, - get_default_muon_param_groups, is_expert_param, update_p) -from .cpu_offload import CPUOffloadPool -from .distributed.utils import (_is_shard, construct_shard_mesh, - get_slices_of_dtensor) -from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO, - _zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5_batched) -from .pipeline import muon_chunk_pipeline, prelaunch_first_gather -from .qk_clip import compute_scales, get_qk_clip_info, qk_clip - -logger = logging.getLogger(__name__) - - -def _expand_expert_params(names, params, expert_keys): - """Expand expert params by splitting on dim 0 (expert dimension). - - Params whose name matches any key in ``expert_keys`` are treated as - expert-parallel tensors. Their outermost dimension is the expert - dimension: an ``(E, out, in)`` tensor becomes ``E`` separate 2D - ``nn.Parameter`` views so that in-place updates propagate back to - the original storage. - - Non-expert params with ``ndim > 2`` trigger an ``AssertionError`` — - if they are expert params, their key must be added to ``expert_keys``. - - The grad must already be set on each expert param (e.g. after momentum). - - For DTensor expert params, placements that shard on dim 0 (expert dim) - are consumed by the split. Non-dim-0 shard placements (e.g. TP) are - preserved: each 2D slice is wrapped as a DTensor on the corresponding - submesh so the parallel pipeline handles the TP communication. - """ - expanded_names = [] - expanded_params = [] - - for n, p in zip(names, params): - is_expert = is_expert_param(n, expert_keys) - is_dtensor = isinstance(p.data, DTensor) - - if is_expert: - if is_dtensor: - logger.debug( - "[expand_expert] %s: expert DTensor, shape=%s, " - "placements=%s, mesh=%s, local_shape=%s", n, p.shape, - p.placements, p.device_mesh.mesh_dim_names, - p.to_local().shape) - else: - logger.debug( - "[expand_expert] %s: expert plain tensor, shape=%s", n, - p.data.shape) - - if not is_expert: - assert p.data.ndim <= 2, ( - f"Param {n} has ndim={p.data.ndim} but does not match " - f"expert_keys={expert_keys}. If this is an expert param, " - f"add its key to expert_keys.") - expanded_names.append(n) - expanded_params.append(p) - continue - - g = p.grad - assert g is not None, ( - f"Expert param {n} must have grad set before expansion") - - tp_mesh = None - tp_placements_2d = None - - if is_dtensor: - local_data = p.to_local() - local_grad = g.to_local() if isinstance(g, DTensor) else g - - # Find non-dim-0 shard placements (e.g. TP sharding). - # After splitting on dim 0, Shard(k) becomes Shard(k-1). - tp_dim_indices = [] - tp_placements_2d = [] - for i, pl in enumerate(p.placements): - if _is_shard(pl) and pl.dim != 0: - tp_dim_indices.append(i) - tp_placements_2d.append(Shard(pl.dim - 1)) - - if tp_dim_indices: - tp_dim_names = tuple(p.device_mesh.mesh_dim_names[i] - for i in tp_dim_indices) - if len(tp_dim_names) == 1: - tp_mesh = p.device_mesh[tp_dim_names[0]] - else: - tp_mesh = p.device_mesh[tp_dim_names] - else: - local_data = p.data - local_grad = g - - # Expand: split dim 0, reshape each slice to 2D. - num_local_experts = local_data.shape[0] - for i in range(num_local_experts): - slice_data = local_data[i] - slice_grad = local_grad[i] - - if tp_mesh is not None: - # Wrap as DTensor on TP submesh so the pipeline handles - # TP communication (gather/scatter across TP ranks). - dt_data = DTensor.from_local(slice_data, - device_mesh=tp_mesh, - placements=tp_placements_2d) - dt_grad = DTensor.from_local(slice_grad, - device_mesh=tp_mesh, - placements=tp_placements_2d) - expert_param = torch.nn.Parameter(dt_data, requires_grad=False) - expert_param.grad = dt_grad - else: - expert_param = torch.nn.Parameter(slice_data, - requires_grad=False) - expert_param.grad = slice_grad - - expanded_names.append(f"{n}[{i}]") - expanded_params.append(expert_param) - - p.grad = None # allow expert grad storage to be freed after pipeline - - return expanded_names, expanded_params - - -class Muon(torch.optim.Optimizer): - """ - Muon - MomentUm Orthogonalized by Newton-schulz - - Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- - processing step, in which each 2D parameter's update is replaced with the nearest orthogonal - matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has - the advantage that it can be stably run in bfloat16 on the GPU. - - Some warnings: - - We believe this optimizer is unlikely to work well for training with small batch size. - - We believe it may not work well for finetuning pretrained models, but we haven't tested this. - - Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. - lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) - momentum: The momentum used by the internal SGD. (0.95 is a good default) - nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) - ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. - Parameters that are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW instead. - adamw_lr: The learning rate for the internal AdamW. - adamw_betas: The betas for the internal AdamW. - adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - warmup_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher warmup_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. - chunk_size : Batch size of parameters to process in each - all2all gather/compute/scatter step. - Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified. - use_distributed_muon: Use distributed muon by Liu et al. (2024). - For testing purpose only. - expert_keys: List of strings to identify expert-parallel parameters. - If any key appears in a parameter's name, its outermost - dimension is treated as the expert dimension and expanded - into per-expert 2D params for Muon. For example, - ``expert_keys=["experts"]`` matches any param whose name - contains "experts". 3D+ params not matched by any key - will raise an error. - """ - - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config=None, - warmup_step=5, - chunk_size=-1, - use_distributed_muon=False, - expert_keys=None): - defaults = dict( - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - nesterov=nesterov, - ns_steps=ns_steps, - adamw_betas=adamw_betas, - adamw_eps=adamw_eps, - none_grad=none_grad, - use_muon=True, - ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) - super().__init__(params, defaults) - - self.debug = debug - self.clip_config = clip_config if clip_config is not None else { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100, - } - self.warmup_step = warmup_step - self.chunk_size = chunk_size - self.use_distributed_muon = use_distributed_muon - self.expert_keys = expert_keys - self.cpu_offload = False - self._cpu_offload_pool: CPUOffloadPool | None = None - self._offload_initialized = False - self._parallel_cache: dict[tuple[str, ...], dict] = {} - self._expert_expand_cache: dict[tuple[int, ...], dict] = {} - - def _calc_flops(self, G, steps): - assert len(G.shape) == 2 - M, N = G.shape - if M > N: - M, N = N, M - - return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) - - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - p.placements, p.device_mesh) - - return shard_mesh, shard_pg, shard_placements - - def init_state_and_assign_params(self, names, params, group, qk_logits): - param_to_state = {} - param_to_flops = {} - - total_flops = 0 - for p in params: - g = p.grad - if g is None: - continue - assert g.ndim == 2, "Muon only supports 2D parameters." - - flops = self._calc_flops(g, group["ns_steps"]) - param_to_flops[id(p)] = flops - total_flops += flops - - if self.debug: - logger.debug("Total TFLOPs for Muon: %.2f TFLOPs", - total_flops / 1e12) - - paired = list(zip(names, params)) - - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) - - round_robin = 0 - mesh = ordered_params[0].device_mesh - placements = ordered_params[0].placements - - shard_mesh, shard_pg, shard_placements = self.get_shard_mesh( - ordered_params[0]) - shard_mesh_flattened = shard_mesh.mesh.flatten() - num_ranks = dist.get_world_size(group=shard_pg) - - for n, p in zip(ordered_names, ordered_params): - if mesh != p.device_mesh: - raise ValueError("All parameters must be on the same mesh.") - if placements != p.placements: - raise ValueError("All parameters must have same placements.") - - worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks - round_robin = (round_robin + 1) % len(shard_mesh_flattened) - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - # Precompute per-rank indices and numels for all-to-all. - rank_indices: dict[int, tuple] = {} - rank_numels: dict[int, int] = {} - for r in range(num_ranks): - indices = get_slices_of_dtensor(p, r, shard_mesh, - shard_placements) - rank_indices[r] = indices - numel = 1 - for idx, dim_size in zip(indices, p.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) - rank_numels[r] = numel - - param_to_state[id(p)] = _muon_state( - worker_rank=worker_rank, - process_group=shard_pg, - rank_indices=rank_indices, - rank_numels=rank_numels, - name=n, - qk_clip_state=qk_clip_state, - ) - - return param_to_state, ordered_params - - def base(self, names, params, group, lr, weight_decay, qk_logits): - # Momentum is already applied by _step_muon before this method. - for n, p in zip(names, params): - g = p.grad - if g is None: - continue - - u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) - - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - def distributed_muon( - self, - names: list[str], - params: list[torch.nn.Parameter], - group: dict[str, Any], - lr: float, - weight_decay: float, - qk_logits: list[torch.Tensor | DTensor] | None, - ): - """Batched Distributed Muon — for testing/correctness verification only. - - Uses all-gather to reconstruct full tensors, computes Newton-Schulz on - the full grad, then slices back to local shards. This is simpler but - slower than the parallel pipeline (all2all) path, so it serves as a - reference implementation for verifying correctness. - """ - with record_function("distributed_muon"): - # Momentum is already applied by _step_muon before this method. - ns_steps = group["ns_steps"] - - # Separate plain tensors (no communication) from DTensors. - plain_names, plain_params = [], [] - dtensor_names, dtensor_params = [], [] - for n, p in zip(names, params): - if p.grad is None: - continue - if isinstance(p.data, DTensor): - dtensor_names.append(n) - dtensor_params.append(p) - else: - plain_names.append(n) - plain_params.append(p) - - # Process plain tensors per-param (no communication). - for n, p in zip(plain_names, plain_params): - u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE), - steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, - qk_logits) - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - if not dtensor_params: - return - - # Group DTensors by (placements, mesh) for batched all-gather. - placement_groups: dict[tuple, - tuple[list, - list]] = defaultdict(lambda: ([], [])) - for n, p in zip(dtensor_names, dtensor_params): - key = (p.placements, p.device_mesh) - placement_groups[key][0].append(n) - placement_groups[key][1].append(p) - - logger.info( - "distributed_muon: %d placement groups, %d total dtensors", - len(placement_groups), len(dtensor_params)) - - for (placements, mesh), (grp_names, - grp_params) in placement_groups.items(): - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - placements, mesh) - rank = dist.get_rank(shard_pg) - world_size = dist.get_world_size(shard_pg) - - logger.info(" group: %d params, placements=%s, world_size=%d", - len(grp_params), placements, world_size) - - # Separate params that can be batched (all shard dims evenly - # divisible) from those needing per-param full_tensor - # (e.g. MoE gate weights with fewer rows than shard ranks). - # all_gather_into_tensor requires equal buffer sizes across - # ranks, so uneven splits must use DTensor full_tensor(). - batch_names, batch_params = [], [] - single_names, single_params = [], [] - for n, p in zip(grp_names, grp_params): - even = all(p.shape[pl.dim] % - shard_mesh.mesh.shape[dim_idx] == 0 - for dim_idx, pl in enumerate(shard_placements)) - if even: - batch_names.append(n) - batch_params.append(p) - else: - single_names.append(n) - single_params.append(p) - - # Process uneven-split params per-param via full_tensor(). - for n, p in zip(single_names, single_params): - with record_function("distributed_muon::newton_schulz"): - g_full = p.grad.full_tensor().to(COMM_DTYPE) - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - if not batch_params: - continue - - logger.info(" batched=%d, single=%d", len(batch_params), - len(single_params)) - - # Concat all local grad shards into a single flat buffer. - with record_function("distributed_muon::gather"): - grad_locals = [ - p.grad.to_local().to(COMM_DTYPE).flatten() - for p in batch_params - ] - numels = [g.numel() for g in grad_locals] - grad_concat = torch.cat(grad_locals) - del grad_locals - - # Single all-gather (replaces N separate full_tensor). - grad_gathered = torch.empty( - grad_concat.numel() * world_size, - dtype=COMM_DTYPE, - device="cuda", - ) - dist.all_gather_into_tensor(grad_gathered, - grad_concat, - group=shard_pg) - - total_numel = grad_concat.numel() - del grad_concat - - # Precompute per-param offsets within the concat buffer. - offsets = [] - off = 0 - for ne in numels: - offsets.append(off) - off += ne - - # Per-param: reconstruct full grad → NS → local update. - for i, (n, p) in enumerate(zip(batch_names, batch_params)): - with record_function("distributed_muon::newton_schulz"): - g_full = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - for r in range(world_size): - r_start = r * total_numel + offsets[i] - shard = grad_gathered[r_start:r_start + numels[i]] - indices = get_slices_of_dtensor( - p, r, shard_mesh, shard_placements) - g_full[indices] = shard.reshape( - g_full[indices].shape) - - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - def _setup_parallel(self, names, params, group, qk_logits): - """Compute (or retrieve cached) parallel pipeline metadata. - - Returns: - (ordered_params, param_to_state, rank, chunk_size) - """ - cache_key = tuple(names) - - if cache_key not in self._parallel_cache: - # First call: compute metadata and populate cache. - param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - shard_pg = param_to_state[id(ordered_params[0])].process_group - rank = dist.get_rank(group=shard_pg) - - if self.chunk_size == -1: - shard_ranks = dist.get_world_size(shard_pg) - chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO - elif self.chunk_size > 0: - chunk_size = self.chunk_size - else: - raise ValueError( - "chunk_size must be -1 or a positive integer.") - - ordered_names = [ - param_to_state[id(p)].name for p in ordered_params - ] - name_to_state = { - param_to_state[id(p)].name: param_to_state[id(p)] - for p in ordered_params - } - self._parallel_cache[cache_key] = { - 'ordered_names': ordered_names, - 'name_to_state': name_to_state, - 'rank': rank, - 'chunk_size': chunk_size, - } - else: - # Cached path: rebuild param_to_state with current id(p) keys. - cache = self._parallel_cache[cache_key] - rank = cache['rank'] - chunk_size = cache['chunk_size'] - - name_to_param = dict(zip(names, params)) - ordered_params = [name_to_param[n] for n in cache['ordered_names']] - - param_to_state = {} - for p, n in zip(ordered_params, cache['ordered_names']): - cached_state = cache['name_to_state'][n] - param_to_state[id(p)] = _muon_state( - worker_rank=cached_state.worker_rank, - process_group=cached_state.process_group, - rank_indices=cached_state.rank_indices, - rank_numels=cached_state.rank_numels, - name=n, - qk_clip_state=get_qk_clip_info(self.clip_config, n, - qk_logits), - ) - - return ordered_params, param_to_state, rank, chunk_size - - def parallel(self, - names, - params, - group, - lr, - weight_decay, - qk_logits, - prelaunch_gather=None): - """ - Perform a parallel optimization step using Muon. - - Parameters are chunked and each chunk is processed by a - :func:`muon_chunk_pipeline` generator. :func:`run_pipeline` - interleaves multiple chunks so that communication and computation - overlap across chunks (the same overlap previously achieved by the - warmup + main-loop index scheduling). - - If ``prelaunch_gather`` is provided, it is passed to the first - chunk's generator to skip re-launching the already in-flight - A2A gather. - """ - - # Momentum is already applied by _step_muon before this method. - - ordered_params, param_to_state, rank, chunk_size = ( - self._setup_parallel(names, params, group, qk_logits)) - - def pipelines(): - first = True - for start in range(0, len(ordered_params), chunk_size): - chunk = ordered_params[start:start + chunk_size] - if chunk: - kwargs = dict( - params=chunk, - param_to_state=param_to_state, - rank=rank, - ns_steps=group["ns_steps"], - lr=lr, - weight_decay=weight_decay, - none_grad=group["none_grad"], - ) - if first and prelaunch_gather is not None: - kwargs['prelaunch_gather'] = prelaunch_gather - first = False - yield muon_chunk_pipeline(**kwargs) - - with record_function("muon::pipeline"): - run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1) - - def _step_muon(self, group, qk_logits=None): - params = group["params"] - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - # Apply momentum to all params before routing/expansion. - # Batched using _foreach_* ops (compiled, fullgraph=True). - with record_function("muon::momentum"): - active_params = [p for p in params if p.grad is not None] - if active_params: - # Ensure momentum buffers exist (avoid zeros_like when already present). - for p in active_params: - if "momentum_buffer" not in self.state[p]: - self.state[p]["momentum_buffer"] = torch.zeros_like( - p.grad) - - # Extract local tensors for compiled batch function. - local_grads = [ - p.grad._local_tensor - if isinstance(p.grad, DTensor) else p.grad - for p in active_params - ] - local_bufs = [ - self.state[p]["momentum_buffer"]._local_tensor - if isinstance(self.state[p]["momentum_buffer"], DTensor) - else self.state[p]["momentum_buffer"] - for p in active_params - ] - - # Wrap momentum as tensor for torch.compile. - batch_pre_ortho(local_grads, local_bufs, - torch.tensor(momentum), group["nesterov"]) - - # For non-nesterov, the result is the momentum buffer. - if not group["nesterov"]: - for p in active_params: - p.grad = self.state[p]["momentum_buffer"] - - # Identify batched experts for deferred NS. - # Detection is cheap (condition checks only); actual NS compute is - # deferred so it can overlap with the first chunk's A2A gather. - deferred_expert_work = [] - if self.expert_keys: - batched_expert_indices = [] - for i, (n, p) in enumerate(zip(names, params)): - if not (is_expert_param(n, self.expert_keys) - and p.grad is not None): - continue - # Eligible: plain tensor, or DTensor with no non-dim-0 shards. - if isinstance(p.data, DTensor): - has_tp = any( - _is_shard(pl) and pl.dim != 0 for pl in p.placements) - if has_tp: - continue - batched_expert_indices.append(i) - - if batched_expert_indices: - # Save refs for deferred NS; free grads from param list. - for i in batched_expert_indices: - p = params[i] - g = p.grad - local_g = (g._local_tensor - if isinstance(g, DTensor) else g) - local_data = (p.data._local_tensor if isinstance( - p.data, DTensor) else p.data) - deferred_expert_work.append((local_data, local_g)) - p.grad = None - - # Remove batched experts from lists before expansion. - keep = sorted( - set(range(len(params))) - set(batched_expert_indices)) - names = [names[i] for i in keep] - params = [params[i] for i in keep] - - def _run_deferred_expert_ns(): - """Execute deferred batched expert NS.""" - if not deferred_expert_work: - return - with record_function("muon::batched_expert_ns"): - ns_steps = group["ns_steps"] - for local_data, local_g in deferred_expert_work: - u = zeropower_via_newtonschulz5_batched( - local_g.to(COMM_DTYPE), steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:]) - local_data.mul_(1 - lr * weight_decay) - local_data.add_(u, alpha=-adjusted_lr) - - # Expand expert params by splitting on dim 0. - logger.debug("[_step_muon] before expand: %d params, expert_keys=%s", - len(params), self.expert_keys) - if self.expert_keys: - cache_key = tuple(id(p) for p in params) - cache = self._expert_expand_cache.get(cache_key) - - if cache is None: - # Cold path: full expansion + build cache metadata. - exp_names, exp_params = _expand_expert_params( - names, params, self.expert_keys) - - # Build per-expert-group info for hot-path grad updates. - grad_info = [] - exp_idx = 0 - for orig_idx, (n, p) in enumerate(zip(names, params)): - if not is_expert_param(n, self.expert_keys): - exp_idx += 1 - continue - - is_dt = isinstance(p.data, DTensor) - num_experts = (p.to_local() if is_dt else p.data).shape[0] - - # Detect TP mesh from the first expanded expert param. - tp_mesh = None - tp_pls = None - sample = exp_params[exp_idx] - if isinstance(sample.data, DTensor): - tp_mesh = sample.data.device_mesh - tp_pls = list(sample.data.placements) - - grad_info.append((orig_idx, num_experts, exp_idx, is_dt, - tp_mesh, tp_pls)) - exp_idx += num_experts - - self._expert_expand_cache[cache_key] = { - 'names': exp_names, - 'params': exp_params, - 'grad_info': grad_info, - } - names, params = exp_names, exp_params - else: - # Hot path: reuse cached params, only update expert grads. - for (orig_idx, num_experts, exp_start, is_dt, tp_mesh, - tp_pls) in cache['grad_info']: - p = params[orig_idx] - g = p.grad - local_grad = (g.to_local() - if is_dt and isinstance(g, DTensor) else g) - for i in range(num_experts): - expert_p = cache['params'][exp_start + i] - sg = local_grad[i] - if tp_mesh is not None: - expert_p.grad = DTensor.from_local( - sg, device_mesh=tp_mesh, placements=tp_pls) - else: - expert_p.grad = sg - p.grad = None - - names = cache['names'] - params = cache['params'] - else: - names, params = _expand_expert_params(names, params, - self.expert_keys) - logger.debug("[_step_muon] after expand: %d params", len(params)) - - param_dtensors = [] - name_dtensors = [] - - param_tensors = [] - name_tensors = [] - - # distributed_muon is a reference implementation for testing only. - # The parallel pipeline (all2all) path below is the production path. - if self.use_distributed_muon: - _run_deferred_expert_ns() - self.distributed_muon(names=names, - params=params, - group=group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits) - return - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - logger.debug( - "[route] %s → base (DTensor all-Replicate), " - "shape=%s, placements=%s", n, p.shape, p.placements) - param_tensors.append(p) - name_tensors.append(n) - else: - logger.debug( - "[route] %s → parallel (DTensor), shape=%s, " - "placements=%s, mesh=%s", n, p.shape, p.placements, - p.device_mesh.mesh_dim_names) - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - logger.debug("[route] %s → base (plain tensor), shape=%s", n, - p.data.shape) - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError(f"Unsupported parameter type: {type(p.data)}") - - logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, " - f"{len(param_tensors)} Tensors → base") - - def group_dtensors(dtensors, names): - # To support different placements, we group parameters by placements - # and run parallel Muon on each group. - - placement_to_params = defaultdict(lambda: ([], [])) - - assert len(dtensors) == len(names) - for p, n in zip(dtensors, names): - placement_to_params[tuple([p.placements, - p.device_mesh])][0].append(n) - placement_to_params[tuple([p.placements, - p.device_mesh])][1].append(p) - return placement_to_params - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - dtensor_group = group_dtensors(param_dtensors, name_dtensors) - - # Pre-launch the first chunk's A2A gather so that the NCCL - # communication overlaps with the (deferred) batched expert NS - # compute on the default CUDA stream. - prelaunch = None - if deferred_expert_work: - first_names, first_params = next(iter(dtensor_group.values())) - ordered, pts, rnk, csz = self._setup_parallel( - first_names, first_params, group, qk_logits) - first_chunk = ordered[:csz] - if first_chunk: - prelaunch = prelaunch_first_gather(first_chunk, pts, rnk, - group["none_grad"]) - - _run_deferred_expert_ns() - - first_group = True - for _, (names, params) in dtensor_group.items(): - pg = prelaunch if first_group else None - first_group = False - self.parallel( - names, - params, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - prelaunch_gather=pg, - ) - else: - _run_deferred_expert_ns() - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - ) - - def _register_states_for_offload(self): - """Register all optimizer state tensors with the CPU offload pool. - - Called once after the first step when states have been lazily created. - Offloads all param states (momentum buffers for Muon, moment1/moment2 - for AdamW) to free GPU memory between steps. - """ - pool = self._cpu_offload_pool - tracked = 0 - for group in self.param_groups: - for p in group["params"]: - if p not in self.state: - continue - state = self.state[p] - if group.get("use_muon", False): - if "momentum_buffer" in state: - pool.track(state["momentum_buffer"]) - tracked += 1 - else: - if "moment1" in state: - pool.track(state["moment1"]) - if "moment2" in state: - pool.track(state["moment2"]) - tracked += 1 - logger.info("[CPUOffload] Registered %d param states for offload", - tracked) - - @torch.no_grad - def step(self, closure=None, qk_logits=None): - """Perform a single optimization step. - - Args: - closure (Callable, optional): A closure that reevaluates the model - and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). - """ - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # H2D: reload optimizer states from CPU before computation. - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload() - - logger.debug("[Muon.step] expert_keys=%s, %d param groups", - self.expert_keys, len(self.param_groups)) - - for i, group in enumerate(self.param_groups): - if group["use_muon"]: - logger.debug("[Muon.step] group %d: use_muon=True, %d params", - i, len(group["params"])) - self._step_muon(group, qk_logits=qk_logits) - else: - logger.debug( - "[Muon.step] group %d: use_muon=False (AdamW), %d params", - i, len(group["params"])) - step_adamw(self.state, group) - - # D2H: offload optimizer states to CPU after computation. - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - return loss - - # ------------------------------------------------------------------ - # CPU offload public helpers - # ------------------------------------------------------------------ - - def turn_on_cpu_offload(self): - """Enable CPU offload for optimizer states.""" - if self.cpu_offload: - return - logger.info("[Muon] turn_on_cpu_offload") - self.cpu_offload = True - if not self.state: - return - self._cpu_offload_pool = CPUOffloadPool() - self._offload_initialized = False - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_off_cpu_offload(self): - """Disable CPU offload and keep optimizer states resident on GPU.""" - if not self.cpu_offload: - return - logger.info("[Muon] turn_off_cpu_offload") - if self._offload_initialized: - self._cpu_offload_pool.reload() - torch.cuda.current_stream().synchronize() - self._cpu_offload_pool = None - self._offload_initialized = False - self.cpu_offload = False - - # ------------------------------------------------------------------ - # Checkpoint support for cpu_offload - # ------------------------------------------------------------------ - - def state_dict(self) -> dict: - if self.cpu_offload: - raise RuntimeError( - "Muon.state_dict() requires turn_off_cpu_offload() before checkpoint save." - ) - return super().state_dict() - - def load_state_dict(self, state_dict: dict) -> None: - if self.cpu_offload: - raise RuntimeError( - "Muon.load_state_dict() requires turn_off_cpu_offload() before checkpoint load." - ) - super().load_state_dict(state_dict) - - # Invalidate adamw.py's module-level tensor caches so that - # the next step rebuilds them with the newly loaded state tensors. - _placement_cache.clear() - _tensor_cache.clear() diff --git a/build/torch29-cxx11-cu126-x86_64-linux/newton_schulz.py b/build/torch29-cxx11-cu126-x86_64-linux/newton_schulz.py deleted file mode 100644 index d939264b69a34e7a3fa78859f34dc265a1159d59..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/newton_schulz.py +++ /dev/null @@ -1,240 +0,0 @@ -from itertools import repeat -from math import inf, sqrt - -import numpy as np -import torch - -from .matmul_transpose_triton import matmul_transpose_assign - -COMM_DTYPE = torch.bfloat16 -DEFAULT_CHUNK_SIZE_RATIO = 4 - - -def _optimal_quintic(l, u, max_iter=1000): - """ - Use the simplified Remez algorithm to find the optimal odd quintic approximant - to the constant function x -> 1 over the interval [l, u]. - - Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum - approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the - two interior equioscillation nodes q, r until convergence. Returns the - closed-form equioscillating solution when l ≈ u. - - Raises ValueError if any intermediate value (a, b, c, E, q, r) is non-finite - (NaN or inf). Raises RuntimeError if convergence is not reached within - max_iter iterations. - """ - assert 0 <= l <= u - if 1 - 5e-6 <= l / u: - return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5) - q = (3 * l + u) / 4 - r = (l + 3 * u) / 4 - E = inf - for _ in range(max_iter): - old_E = E - LHS = np.array( - [ - [l, l**3, l**5, 1], - [q, q**3, q**5, -1], - [r, r**3, r**5, 1], - [u, u**3, u**5, -1], - ] - ) - a, b, c, E = np.linalg.solve(LHS, np.ones(4)) - if not np.all(np.isfinite([a, b, c, E])): - raise ValueError( - f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}" - ) - q, r = np.sqrt( - (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / (10 * c) - ) - if not np.all(np.isfinite([q, r])): - raise ValueError(f"_optimal_quintic: non-finite node update q={q}, r={r}") - if abs(old_E - E) <= 1e-15: - break - else: - raise RuntimeError( - f"_optimal_quintic: did not converge after {max_iter} iterations" - ) - return float(a), float(b), float(c) - - -def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0): - """ - Compute the Polar Express coefficient series for `num_iters` quintic iterations. - - Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that - compose to map singular values from [l, 1] toward 1. At each step: - 1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion` - prevents near-zero singular values from stalling by raising the effective - lower bound; if it is active (cushion*u > l), the coefficients are - rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u]. - 2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the - last iteration, providing numerical headroom at the cost of a slightly slower - final convergence step. - 3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1). - - Returns a list of (a, b, c) tuples, one per iteration. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - u = 1 - assert 0 <= l <= u - safety_factor = 1 + safety_factor_eps - coefficients = [] - for iter in range(num_iters): - a, b, c = _optimal_quintic(max(l, cushion * u), u) - if cushion * u > l: - pl = a * l + b * l**3 + c * l**5 - pu = a * u + b * u**3 + c * u**5 - rescaler = 2 / (pl + pu) - a *= rescaler - b *= rescaler - c *= rescaler - if iter < num_iters - 1: - a /= safety_factor - b /= safety_factor**3 - c /= safety_factor**5 - coefficients.append((a, b, c)) - l = a * l + b * l**3 + c * l**5 - u = 2 - l - return coefficients - - -# Precomputed Polar Express coefficients (a, b, c) for 10 quintic Newton-Schulz -# iterations. Each tuple is the minimax-optimal (Remez/equioscillation) odd quintic -# approximant to x->1 over the current singular-value interval, computed once at -# import time and reused across all optimizer steps. -# -# Contrast with the former hardcoded NS coefficients (5 fixed tuples): -# - Former: empirically tuned to maximize slope at zero; did not converge -# singular values to 1, yielding US'V^T with S' ~ Uniform(0.5, 1.5) instead -# of the true polar factor UV^T. -# - Polar Express: analytically optimal per step, adapting to the shrinking -# singular-value interval [l, u] as iterations progress; converges all -# singular values to 1, producing the exact polar factor UV^T. -_coeffs_list = _optimal_composition( - l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02 -) - - -# This code is adapted from: -# KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py) -# NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress) -# matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon) -@torch.no_grad() -def _zeropower_via_newtonschulz5(G, steps): - """ - Compute the polar factor of G via the Polar Express method. - - Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c) - are the Polar Express coefficients from `_coeffs_list`. Each step is the - optimal odd quintic approximant to x -> 1 over the current singular-value - interval, minimizing the maximum approximation error (Remez / minimax criterion). - The composition maps singular values from [l, 1] to near 1, producing the - polar factor (orthogonal factor in the polar decomposition G = UP). - - `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2, - cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE - X = G # no manual typecast - - if G.size(0) > G.size(1): - X = X.T - - X = X / (X.norm() + 1e-7) - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - # Perform the NS iterations - for a, b, c in hs: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(0) > G.size(1): - X = X.T - - return X - - -@torch.no_grad() -def _zeropower_via_newtonschulz5_batched(G, steps): - """Batched polar factor computation for 3D (E, out, in) tensors. - - Same algorithm as ``_zeropower_via_newtonschulz5`` but uses - ``torch.bmm`` / ``torch.baddbmm`` instead of the 2D Triton kernel, - processing all E expert matrices in a single batched call. - """ - assert len(G.shape) == 3 - assert G.dtype == COMM_DTYPE - X = G - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - # Per-expert Frobenius norm. - X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) - - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - for a, b, c in hs: - buf1 = torch.bmm(X, X.transpose(-2, -1)) - buf2 = torch.bmm(buf1, buf1.transpose(-2, -1)) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.baddbmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - return X - - -_ns_per_shape: dict[tuple[int, ...], callable] = {} -_use_compile = True - - -def set_ns_compile(enabled: bool): - """Toggle torch.compile for Newton-Schulz iteration.""" - global _use_compile - _use_compile = enabled - - -def zeropower_via_newtonschulz5(G, steps=5): - if not _use_compile: - return _zeropower_via_newtonschulz5(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() - - -def zeropower_via_newtonschulz5_batched(G, steps=5): - """Compile-cached batched Newton-Schulz for 3D expert tensors.""" - if not _use_compile: - return _zeropower_via_newtonschulz5_batched(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile( - _zeropower_via_newtonschulz5_batched, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() diff --git a/build/torch29-cxx11-cu126-x86_64-linux/optimizer/__init__.py b/build/torch29-cxx11-cu126-x86_64-linux/optimizer/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/optimizer/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu126-x86_64-linux/pipeline.py b/build/torch29-cxx11-cu126-x86_64-linux/pipeline.py deleted file mode 100644 index c0c2d515856182d8d15ad27dd4e4e093b29397d6..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/pipeline.py +++ /dev/null @@ -1,468 +0,0 @@ -import logging -from typing import Generator - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -from .core import _muon_state, adjust_lr_for_muon -from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5 -from .qk_clip import compute_scales - -logger = logging.getLogger(__name__) - -# ====================================================================== -# Stage helpers -# ====================================================================== - - -def _launch_gather( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Allocate gather buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_gather``). - gathered_grads: ``{id(p): empty_tensor}`` for owned params, - ``None`` for non-owned. - recv_counts: Per-source-rank element counts. - """ - # Allocate gathered-grad buffers - gathered_grads: dict[int, torch.Tensor | None] = {} - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - gathered_grads[id(p)] = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - else: - gathered_grads[id(p)] = None - - # Build send buffer – batch grad copies via torch.cat - # (1-2 fused kernels vs N individual narrow().copy_() calls). - send_counts = [0] * num_ranks - for p in params: - state = param_to_state[id(p)] - send_counts[state.worker_rank] += state.rank_numels[rank] - - total_send = sum(send_counts) - if total_send > 0: - # Group grad slices by destination rank in a single pass. - dst_to_grads = [[] for _ in range(num_ranks)] - for p in params: - state = param_to_state[id(p)] - n = state.rank_numels[rank] - if n > 0: - g = p.grad.to_local() - dst_to_grads[state.worker_rank].append(g.reshape(-1)) - - # Flatten in dst order and cat once. - all_slices = [] - for dst in range(num_ranks): - all_slices.extend(dst_to_grads[dst]) - send_buf = torch.cat(all_slices) - if send_buf.dtype != COMM_DTYPE: - send_buf = send_buf.to(COMM_DTYPE) - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += state.rank_numels[src] - recv_counts[src] = total - - recv_buf = torch.empty(sum(recv_counts), dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - logger.debug(f"send_buf size: {send_buf.numel()}, " - f"recv_buf size: {recv_buf.numel()}, " - f"recv_counts: {recv_counts}, " - f"send_counts: {send_counts}, " - f"process_group: {str(process_group)}") - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, gathered_grads, recv_counts - - -def _complete_gather( - recv_buf: torch.Tensor, - recv_counts: list[int], - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - param_to_state: dict[int, _muon_state], - rank: int, -) -> None: - """Reconstruct gathered grads from the recv buffer (in-place).""" - off = 0 - for src in range(len(recv_counts)): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - - indices = state.rank_indices[src] - - shard_view = gathered_grads[id(p)][indices] - n = shard_view.numel() - if n == 0: - continue - - sg = recv_buf.narrow(0, off + inner_off, n) - sg = sg.reshape(shard_view.shape) - gathered_grads[id(p)][indices] = sg - - inner_off += n - assert inner_off == block - off += block - - -def _compute_ns( - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - ns_steps: int, -) -> dict[int, torch.Tensor | None]: - """Run Newton-Schulz orthogonalization on owned parameters. - - Returns: - computed_us: ``{id(p): orthogonalized_update}`` for owned params. - """ - computed_us: dict[int, torch.Tensor | None] = {} - for p in owned_params: - u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps) - gathered_grads[id(p)] = None # free gathered grad - computed_us[id(p)] = u - return computed_us - - -def _launch_scatter( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, - computed_us: dict[int, torch.Tensor | None], -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor], list[int]]: - """Allocate scatter buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_scatter``). - scattered_us: Empty dict, populated by ``_complete_scatter`` with - zero-copy views into ``recv_buf``. - recv_counts: Per-source-rank element counts. - """ - # scattered_us is populated by _complete_scatter with zero-copy views - # into recv_buf, avoiding N empty_like allocations + N copy_ calls. - # Pre-seed entries for params whose local shard is empty (rank_numels == 0) - # so _update_params can iterate all params without KeyError. - scattered_us: dict[int, torch.Tensor] = {} - for p in params: - if param_to_state[id(p)].rank_numels[rank] == 0: - scattered_us[id(p)] = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - # Build send buffer – batch via torch.cat - # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls). - send_counts = [0] * num_ranks - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - for dst_rank in range(num_ranks): - send_counts[dst_rank] += state.rank_numels[dst_rank] - - total_send = sum(send_counts) - if total_send > 0: - # Cache u_full conversions to avoid redundant .to() per dst_rank. - u_fulls = {} - for p in owned_params: - u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous() - - # Collect slices in dst order (matches all-to-all send layout). - all_slices = [] - for dst_rank in range(num_ranks): - for p in owned_params: - state = param_to_state[id(p)] - su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten() - if su.numel() > 0: - all_slices.append(su) - - send_buf = torch.cat(all_slices) if all_slices else torch.empty( - 0, dtype=COMM_DTYPE, device="cuda") - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += state.rank_numels[rank] - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, scattered_us, recv_counts - - -def _complete_scatter( - recv_buf: torch.Tensor, - recv_counts: list[int], - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], -) -> None: - """Populate scattered_us with zero-copy views into recv_buf. - - Instead of pre-allocating tensors and copying, we assign views directly - from ``recv_buf``. This eliminates N ``empty_like`` + N ``copy_`` calls. - The underlying storage of ``recv_buf`` is kept alive through the views - until ``scattered_us`` is cleared after ``_update_params``. - """ - off = 0 - for src in range(len(recv_counts)): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = state.rank_numels[rank] - if n == 0: - continue - - scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - - inner_off += n - - assert inner_off == block - off += block - - -def _update_params( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], - lr: float, - weight_decay: float, -) -> None: - """Apply weight decay, Muon update, and optional QK clipping. - - Uses batched ``_foreach_mul_`` for weight decay and batched - ``_foreach_add_`` for the Muon update, grouping parameters by - adjusted_lr to minimize kernel launches while preserving float32 - precision for the alpha scaling. - """ - if not params: - return - - # Batched weight decay: p *= (1 - lr * wd) — single fused kernel. - p_locals = [p._local_tensor for p in params] - torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay) - - # Group params by adjusted_lr so _foreach_add_ can use a single - # alpha per group (preserves float32 precision for alpha scaling). - lr_groups: dict[float, tuple[list, list]] = {} - for p in params: - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - if adjusted_lr not in lr_groups: - lr_groups[adjusted_lr] = ([], []) - lr_groups[adjusted_lr][0].append(p._local_tensor) - lr_groups[adjusted_lr][1].append(scattered_us[id(p)]) - - for adjusted_lr, (p_group, u_group) in lr_groups.items(): - torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr) - - # QK clipping – applied directly on the local tensor to - # avoid DTensor sharding-propagation issues with _StridedShard. - for p in params: - state = param_to_state[id(p)] - if state.qk_clip_state is None: - continue - scales_full = compute_scales(p, state.qk_clip_state) - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = state.rank_indices[rank][0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - -# ====================================================================== -# Pre-launch helper for overlapping first chunk's gather with other work. -# ====================================================================== - - -@torch.no_grad() -def prelaunch_first_gather( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - none_grad: bool, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Launch the first chunk's A2A gather early for overlap with other compute. - - Call this *before* expensive GPU work (e.g. batched expert NS) so that - the NCCL all-to-all runs concurrently on the NCCL stream while the - default stream executes compute. - - Returns the same 4-tuple that ``_launch_gather`` produces, which should - be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - with record_function("muon::prelaunch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - return work, recv_buf, gathered_grads, recv_counts - - -# ====================================================================== -# Main generator – thin orchestrator that wires stages together. -# ====================================================================== - - -@torch.no_grad() -def muon_chunk_pipeline( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - ns_steps: int, - lr: float, - weight_decay: float, - none_grad: bool, - prelaunch_gather: tuple | None = None, -) -> Generator[None, None, None]: - """Process one chunk of parameters through the full Muon pipeline. - - Stages: gather -> compute (Newton-Schulz) -> scatter -> update. - - Each ``yield`` lets :func:`run_pipeline` interleave other chunks so - that communication and computation overlap across chunks. Async - communication is launched via ``async_op=True`` and completed after - the yield with ``work.wait()``. - - Overlap happens because :func:`run_pipeline` admits one new chunk - per iteration (staggered admission). While chunk *N* does NS - compute on the default CUDA stream, chunk *N+1*'s async all-to-all - runs concurrently on the NCCL stream — no separate ``comm_stream`` - is required. - - If ``prelaunch_gather`` is provided, the gather was already launched - by :func:`prelaunch_first_gather` and we skip launching it again. - - Yields exactly **2** times: - - 1. After launching async all-to-all gather (or immediately if pre-launched). - 2. After launching async all-to-all scatter. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - if prelaunch_gather is not None: - # Gather was pre-launched; none_grad already handled by caller. - work, recv_buf, gathered_grads, recv_counts = prelaunch_gather - else: - # Normal path: launch async gather. - with record_function("muon::launch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - yield # --- YIELD 1: other chunks can launch their gather --- - - with record_function("muon::wait_gather"): - work.wait() - _complete_gather(recv_buf, recv_counts, owned_params, gathered_grads, - param_to_state, rank) - del recv_buf - - # Stage 3: Newton-Schulz orthogonalization. - with record_function("muon::newton_schulz"): - computed_us = _compute_ns(owned_params, gathered_grads, ns_steps) - gathered_grads.clear() - - # Stages 4-5: launch async scatter. - with record_function("muon::launch_scatter"): - work, recv_buf, scattered_us, recv_counts = _launch_scatter( - params, owned_params, param_to_state, rank, num_ranks, - process_group, computed_us) - computed_us.clear() - - yield # --- YIELD 2: other chunks can launch their scatter --- - - with record_function("muon::wait_scatter"): - work.wait() - _complete_scatter(recv_buf, recv_counts, params, param_to_state, rank, - scattered_us) - del recv_buf - - # Stage 6: apply parameter updates. - with record_function("muon::update_params"): - _update_params(params, param_to_state, rank, scattered_us, lr, - weight_decay) - scattered_us.clear() diff --git a/build/torch29-cxx11-cu126-x86_64-linux/qk_clip.py b/build/torch29-cxx11-cu126-x86_64-linux/qk_clip.py deleted file mode 100644 index 2aba711b3004b7f09e7141da7ef834bd61cc2430..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/qk_clip.py +++ /dev/null @@ -1,198 +0,0 @@ -import logging -import math -from dataclasses import dataclass - -import torch -from torch.distributed.tensor import DTensor - -from .core import normalize_fqn - -logger = logging.getLogger(__name__) - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - and return (kind, layer_index). - - Supported kinds: - MHA/GQA: 'wq', 'wk', 'q_proj', 'k_proj' - MLA: 'wq_b' (Q up-proj), 'wkv_b' (KV up-proj) - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.1.attn.wq_b.weight' -> ('wq_b', 1) - 'model.0.attn.wkv_b.weight' -> ('wkv_b', 0) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = normalize_fqn(name).split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: str | None # 'wq'/'q_proj'/'wq_b' or 'wk'/'k_proj'/'wkv_b' or None - indices: list[int] # which heads to consider for clipping - head_dim: int # from config (qk_head_dim for MLA wq_b) - threshold: float # from config - logit: torch.Tensor | None - - # MLA-specific fields - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 - - -def get_qk_clip_info(clip_config, n, qk_logits): - """Extract QK clipping info for a named parameter. - - Args: - clip_config: QK clipping configuration dict (or None). - MHA/GQA keys: head_dim, threshold, q_indices, k_indices - MLA extra keys: is_mla=True, qk_nope_head_dim, qk_rope_head_dim, v_head_dim - n: Parameter name string. - qk_logits: Dict mapping layer indices to logit tensors (or None). - - Returns: - QKClipInfo instance with clipping configuration for this parameter. - """ - if clip_config is None: - return None - - head_dim = clip_config.get('head_dim') - threshold = clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - # In TP settings, qk_logits may be DTensor - # We convert it to full tensor here for simplicity - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - -def compute_scales(p, qk_clip_state): - """Compute per-head scaling factors for QK clipping. - - Returns scales tensor (√γ per head) if any head exceeds threshold, else None. - For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim. - """ - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - # Check if any head exceeds threshold before allocating. - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - - if not head_scales: - return None - - # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows - if qk_clip_state.is_mla and kind == 'wkv_b': - effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim - else: - effective_head_dim = head_dim - - H_global = p.shape[0] // effective_head_dim - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale - return scales_full - - -def qk_clip(p, scales, info): - """Apply per-head scaling to a Q/K projection weight matrix. - - Args: - p: Parameter (nn.Parameter or raw tensor). - scales: [n_heads] tensor, each element = √γ_h. - info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions. - - MLA sub-region scaling per Algorithm 1 (MuonClip): - wq_b: q_nope rows → √γ, q_pe rows → γ - wkv_b: k_nope rows → √γ, v rows → unchanged - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not info.is_mla: - # MHA/GQA: uniform √γ applied to all rows in each head - W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: vectorized sub-region scaling within each head - if info.kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_head_dim = qk_nope + info.qk_rope_head_dim - W_3d = W.view(-1, qk_head_dim, W.shape[1]) # [H, qk_head_dim, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # q_nope → √γ - W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1, - 1)) # q_pe → γ - - elif info.kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - W_3d = W.view(-1, kv_stride, W.shape[1]) # [H, kv_stride, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # k_nope → √γ - # v rows: not touched (k_R shared rotary unchanged) diff --git a/build/torch29-cxx11-cu128-x86_64-linux/__init__.py b/build/torch29-cxx11-cu128-x86_64-linux/__init__.py deleted file mode 100644 index 239c7a65f8293e7d0df28f05fce645af56d628c0..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .muon import Muon - -__all__ = [ - "Muon", -] diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_ops.py b/build/torch29-cxx11-cu128-x86_64-linux/_ops.py deleted file mode 100644 index 034bff088659b5df6f6d401feb18c89dc5f33b29..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _optimizer_8d53b78_dirty -ops = torch.ops._optimizer_8d53b78_dirty - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_optimizer_8d53b78_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so b/build/torch29-cxx11-cu128-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so deleted file mode 100755 index 0ecf56af29245daebb0f2048f89e30d786c9a8a7..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5e9768c250f77d75c777d94e1f7817cb24372b8d644bc28d3edfa1a829317272 -size 1999872 diff --git a/build/torch29-cxx11-cu128-x86_64-linux/adamw.py b/build/torch29-cxx11-cu128-x86_64-linux/adamw.py deleted file mode 100644 index b5a95816a9f5b9e1889eaadae65373bfbced809a..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/adamw.py +++ /dev/null @@ -1,271 +0,0 @@ -import logging -from collections import defaultdict -from typing import cast - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -def fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: float | torch.Tensor, - weight_decay: float, - eps: float, - maximize: bool, -) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: dict | None = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) - - -def _to_local(t): - """Unwrap DTensor to local tensor for fused ops.""" - return t._local_tensor if isinstance(t, DTensor) else t - - -# --------------------------------------------------------------------------- -# Caches for eliminating per-step Python overhead. -# -# Placement grouping and tensor list assembly are identical every step -# (params don't change placement, moment/step tensors are the same objects -# after initialisation). We cache them keyed by id() of the param list -# stored in param_groups (stable across steps). -# -# Only gradients change each step and must be collected fresh. -# --------------------------------------------------------------------------- - -# id(group["params"]) → dict[placement_key, list[param]] -_placement_cache: dict[int, dict[tuple, list]] = {} - -# id(placement_group_list) → (params_local, moment1, moment2, state_steps) -_tensor_cache: dict[int, tuple[list, list, list, list]] = {} - - -def _step_adamw_params_slow(optimizer_state, params, group): - """Uncached fallback for the rare case where some params lack grads.""" - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - g = p.grad - if g is None: - continue - state = optimizer_state[p] - params_with_grads.append(_to_local(p)) - grads.append(_to_local(g)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - if not params_with_grads: - return - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw_params(optimizer_state, params, group): - """Run fused AdamW on a list of parameters sharing the same placement. - - After the first call, cached tensor lists (params_local, moment1, - moment2, state_steps) are reused — only gradients are collected fresh. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - params: List of parameters to update. - group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay. - """ - # Collect grads — the only thing that changes each step. - with record_function("adamw::collect_grads"): - grads = [] - for p in params: - g = p.grad - if g is None: - # Rare: fall back to slow path that filters per-param. - _step_adamw_params_slow(optimizer_state, params, group) - return - grads.append(_to_local(g)) - - tensor_key = id(params) - if tensor_key not in _tensor_cache: - with record_function("adamw::init_tensor_cache"): - params_local = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - state = optimizer_state[p] - params_local.append(_to_local(p)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(p.grad) - state["moment2"] = torch.zeros_like(p.grad) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - _tensor_cache[tensor_key] = (params_local, moment1, moment2, - state_steps) - - params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key] - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - with record_function("adamw::fused_adamw"): - fused_adamw( - params_local, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw(optimizer_state, group): - """Dispatch AdamW step, grouping parameters by type and placement. - - Placement grouping is cached after the first call since params never - change their placement between steps. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - group: Parameter group dict. - """ - params = group["params"] - placement_key = id(params) - - if placement_key not in _placement_cache: - with record_function("adamw::group_by_placement"): - placement_to_params: dict[tuple, - list[torch.Tensor]] = defaultdict(list) - for p in params: - match p: - case DTensor(): - logger.debug( - "[AdamW] DTensor param: shape=%s, placements=%s, " - "mesh=%s, grad=%s", p.shape, p.placements, - p.device_mesh.mesh_dim_names, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple( - [p.placements, p.device_mesh])].append(p) - case torch.Tensor(): - logger.debug( - "[AdamW] plain param: shape=%s, grad=%s", p.shape, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple([torch.Tensor, - None])].append(p) - - logger.debug("[AdamW] %d placement groups, %d total params", - len(placement_to_params), len(params)) - - _placement_cache[placement_key] = dict(placement_to_params) - - for group_params in _placement_cache[placement_key].values(): - step_adamw_params(optimizer_state, group_params, group) diff --git a/build/torch29-cxx11-cu128-x86_64-linux/async_utils.py b/build/torch29-cxx11-cu128-x86_64-linux/async_utils.py deleted file mode 100644 index a45c530ac9cad88e3555ec1047a6aa59f225347e..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/async_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import logging -from typing import Generator - -logger = logging.getLogger(__name__) - - -class _Task: - """Internal: wraps a generator, advances one yield at a time.""" - - def __init__(self, generator: Generator[None, None, None], index: int): - self._generator = generator - self._index = index - self._steps_completed = 0 - self.step() # run to first yield - - def step(self) -> bool: - try: - next(self._generator) - self._steps_completed += 1 - logger.debug("pipeline[%d] completed stage %d", self._index, - self._steps_completed) - return True - except StopIteration: - logger.debug("pipeline[%d] finished after %d stages", self._index, - self._steps_completed) - return False - - def close(self): - self._generator.close() - - -def run_pipeline( - pipelines: Generator[Generator[None, None, None], None, None], - max_concurrent: int, -) -> None: - """Run generator-based pipelines with bounded concurrency. - - Each pipeline is a generator that yields at stage boundaries. - The runtime interleaves pipelines so communication and computation - overlap across chunks. - """ - if max_concurrent <= 0: - raise ValueError(f"max_concurrent must be > 0, got {max_concurrent}") - - have_new = True - task_index = 0 - previous_tasks: list[_Task] = [] - - try: - while have_new or previous_tasks: - running_tasks: list[_Task] = [] - - # Admit one new pipeline per iteration (staggered admission). - # Admitting one at a time ensures that while chunk N does NS - # compute on the default stream, chunk N+1's NCCL all-to-all - # runs concurrently on the NCCL stream — creating real - # communication/computation overlap on the GPU. - if have_new and len(previous_tasks) < max_concurrent: - try: - gen = next(pipelines) - task = _Task(gen, task_index) - task_index += 1 - running_tasks.append(task) - except StopIteration: - have_new = False - - # Advance every previously-yielded task by one step. - for task in previous_tasks: - if task.step(): - running_tasks.append(task) - - previous_tasks = running_tasks - except BaseException: - # Clean up all in-flight generators to release GPU resources. - for task in previous_tasks: - task.close() - raise diff --git a/build/torch29-cxx11-cu128-x86_64-linux/core.py b/build/torch29-cxx11-cu128-x86_64-linux/core.py deleted file mode 100644 index c69d515afef305ad0ed66374095fa2d2468d99cc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/core.py +++ /dev/null @@ -1,219 +0,0 @@ -import logging -import math -from dataclasses import dataclass -from typing import List - -import torch -from torch.distributed import ProcessGroup -from torch.distributed.tensor import DTensor - -# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into -# parameter FQNs. Activation checkpointing similarly inserts -# "_checkpoint_wrapped_module". Strip these so name-based matching (skip_keys, -# expert_keys, QK layer parsing) works regardless of wrapper nesting. -_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) - -logger = logging.getLogger(__name__) - - -def normalize_fqn(name: str) -> str: - """Strip torch.compile / checkpoint wrapper components from a parameter FQN.""" - return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS) - - -@dataclass -class _muon_state: - worker_rank: int - process_group: ProcessGroup - rank_indices: dict[int, tuple] # local_rank -> per-dim indices - rank_numels: dict[int, int] # local_rank -> numel - name: str - qk_clip_state: torch.Tensor | None = None - - -def _batch_momentum( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update (no nesterov).""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - - -def _batch_momentum_nesterov( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update with nesterov correction.""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - nesterov_terms = torch._foreach_mul(momentum_bufs, momentum) - torch._foreach_add_(grads, nesterov_terms) - - -_compiled_momentum: dict[bool, callable] = {} -_use_momentum_compile = True - - -def set_momentum_compile(enabled: bool): - """Toggle torch.compile for batched momentum.""" - global _use_momentum_compile - _use_momentum_compile = enabled - - -def batch_pre_ortho( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, - nesterov: bool, -) -> None: - """Batched momentum update on lists of plain tensors. - - Mirrors dion's ``muon_update_pre_orthogonalize``. - Inputs must be plain CUDA tensors (not DTensor). - Modifies ``momentum_bufs`` and (for nesterov) ``grads`` in-place. - - When compile is enabled, uses separately compiled functions for - nesterov=True/False to avoid graph breaks from the branch. - """ - fn = _batch_momentum_nesterov if nesterov else _batch_momentum - if _use_momentum_compile: - if nesterov not in _compiled_momentum: - _compiled_momentum[nesterov] = torch.compile(fn) - fn = _compiled_momentum[nesterov] - fn(grads, momentum_bufs, momentum) - - -def _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay): - """Weight-decay + update on plain tensors. - - Not compiled: per-param @torch.compile caused ~0.25ms TorchDynamo cache - lookup per call × 256+ params = massive overhead. The pipeline path uses - batched _foreach_* ops instead; this function remains for base() and - distributed_muon(). - """ - p_data.mul_(1 - lr * weight_decay) - p_data.add_(u_data, alpha=-adjusted_lr) - - -def update_p(p, u, lr, adjusted_lr, weight_decay): - """Apply weight decay and orthogonalized update to parameter. - - Args: - p: Parameter (torch.nn.Parameter or DTensor). - u: Orthogonalized update tensor. - lr: Base learning rate. - adjusted_lr: Size-adjusted learning rate. - weight_decay: Weight decay coefficient. - """ - # Unwrap Parameter -> underlying data tensor. - p_data = p.data if isinstance(p, torch.nn.Parameter) else p - # Unwrap DTensor -> local CUDA tensor for compiled kernel. - if isinstance(p_data, DTensor): - p_data = p_data._local_tensor - u_data = u._local_tensor if isinstance(u, DTensor) else u - _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay) - - -def adjust_lr_for_muon(lr, param_shape): - """Scale learning rate based on parameter matrix dimensions. - - Args: - lr: Base learning rate. - param_shape: Shape of the parameter tensor. - - Returns: - Adjusted learning rate. - """ - A, B = param_shape[:2] - # We adjust the learning rate and weight decay based on the size of the parameter matrix - # as described in the paper - adjusted_ratio = 0.2 * math.sqrt(max(A, B)) - adjusted_lr = lr * adjusted_ratio - return adjusted_lr - - -def _match_key(parts, key): - """Check if key matches as contiguous components in parts. - - Single-component keys (e.g. "experts") match any single component. - Multi-component keys (e.g. "experts.w1") match as a contiguous subsequence. - """ - key_parts = key.split(".") - key_len = len(key_parts) - if key_len == 1: - return key in parts - return any(parts[i:i + key_len] == key_parts - for i in range(len(parts) - key_len + 1)) - - -def is_expert_param(name, expert_keys): - """Check if a parameter name matches any expert key (component-level).""" - if not expert_keys: - return False - parts = normalize_fqn(name).split(".") - return any(_match_key(parts, key) for key in expert_keys) - - -def default_is_muon(name, x, expert_keys=None): - normalized = normalize_fqn(name) - parts = normalized.split(".") - skip_keys = [ - "embed_tokens", - "lm_head", - "tok_embeddings", - "output", - "mhc_attn", - "mhc_ffn", - "lambda_proj", - ] - if any(key in parts for key in skip_keys): - logger.info( - "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d", - normalized, name, x.ndim) - return False - effective_ndim = x.ndim - is_expert = is_expert_param(name, expert_keys) - if is_expert: - effective_ndim -= 1 - result = effective_ndim >= 2 - logger.info( - "[is_muon] %s (orig: %s): ndim=%d, expert=%s, effective_ndim=%d → %s", - normalized, name, x.ndim, is_expert, effective_ndim, - "Muon" if result else "AdamW") - return result - - -def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None): - if is_muon_func is None: - is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys) - - muon_params, muon_names = [], [] - non_muon_params, non_muon_names = [], [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - non_muon_names.append(n) - - logger.info("[param_groups] expert_keys=%s, Muon=%d, AdamW=%d", - expert_keys, len(muon_names), len(non_muon_names)) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] diff --git a/build/torch29-cxx11-cu128-x86_64-linux/cpu_offload.py b/build/torch29-cxx11-cu128-x86_64-linux/cpu_offload.py deleted file mode 100644 index fb5e69154a1d4a6c884491413a37a9acf0f66c80..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/cpu_offload.py +++ /dev/null @@ -1,206 +0,0 @@ -"""CPU offloading for optimizer states. - -Manages a pinned CPU memory pool and async CUDA streams to offload -optimizer state tensors (momentum buffers, Adam moments) to CPU between -optimizer steps, freeing GPU memory. - -All tracked tensors are packed into a single flat pinned CPU buffer -(per dtype). D2H and H2D copies are performed per-tensor directly -between individual GPU tensors and their slice of the CPU flat buffer -— no GPU staging buffer is allocated, so there is **no temporary GPU -memory spike** during offload or reload. - -Individual tensor storages are freed after offload via -``untyped_storage().resize_(0)``, preserving tensor identity so -downstream caches remain valid. -""" - -import logging -from collections import defaultdict - -import torch -from torch.distributed.tensor import DTensor - -logger = logging.getLogger(__name__) - - -class CPUOffloadPool: - """Pinned CPU memory pool for async optimizer state offloading. - - Tracked tensors are grouped by dtype. Each group gets a single flat - pinned CPU buffer. D2H / H2D copies are per-tensor (into slices of - the flat buffer) to avoid allocating a GPU staging buffer. - """ - - def __init__(self): - self._managed: list[torch.Tensor] = [] - self._storage_nbytes: dict[int, int] = {} # id(t) → bytes - - # Per-dtype group: populated on first offload. - # dtype → dict with keys: - # "indices" : list[int] managed-list indices - # "offsets" : list[tuple[int,int]] (start, numel) in flat buf - # "total" : int total numel - # "cpu_flat" : Tensor pinned CPU buffer - self._groups: dict[torch.dtype, dict] = {} - - self._offload_stream: torch.cuda.Stream | None = None - self._device: torch.device | None = None - self._initialized: bool = False - self._logged: bool = False - - # ------------------------------------------------------------------ - @staticmethod - def _local(t: torch.Tensor) -> torch.Tensor: - """Unwrap DTensor to its local CUDA tensor.""" - return t._local_tensor if isinstance(t, DTensor) else t - - def _ensure_stream(self): - if self._offload_stream is None: - self._offload_stream = torch.cuda.Stream(device=self._device) - - # ------------------------------------------------------------------ - def track(self, tensor: torch.Tensor): - """Register a GPU tensor for CPU offloading. Idempotent.""" - tid = id(tensor) - if tid in self._storage_nbytes: - return - local = self._local(tensor) - if self._device is None: - self._device = local.device - storage = local.untyped_storage() - # Skip tensors with empty storage (e.g. empty FSDP shards) - if storage.size() == 0: - return - self._storage_nbytes[tid] = storage.size() - self._managed.append(tensor) - - # ------------------------------------------------------------------ - def _init_buffers(self): - """Build per-dtype flat buffers on first offload.""" - # Group managed tensors by dtype. - dtype_map: dict[torch.dtype, list[tuple[int, int]]] = defaultdict(list) - for idx, t in enumerate(self._managed): - local = self._local(t) - dtype_map[local.dtype].append((idx, local.numel())) - - total_cpu_bytes = 0 - for dtype, entries in dtype_map.items(): - offsets: list[tuple[int, int]] = [] - indices: list[int] = [] - off = 0 - for idx, n in entries: - indices.append(idx) - offsets.append((off, n)) - off += n - cpu_flat = torch.empty(off, dtype=dtype, device="cpu", pin_memory=True) - self._groups[dtype] = { - "indices": indices, - "offsets": offsets, - "total": off, - "cpu_flat": cpu_flat, - } - total_cpu_bytes += off * cpu_flat.element_size() - - self._initialized = True - logger.info( - "[CPUOffload] Pool initialized: %d tensors, %d dtype group(s), " - "%.2f MB pinned CPU memory", - len(self._managed), - len(self._groups), - total_cpu_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def offload(self): - """Per-tensor async D2H into CPU flat buffer, then free GPU storage.""" - if not self._managed: - return - if not self._initialized: - self._init_buffers() - self._ensure_stream() - - # Offload stream waits for compute to finish. - compute_event = torch.cuda.current_stream(self._device).record_event() - self._offload_stream.wait_event(compute_event) - - offloaded_bytes = 0 - - # Per-tensor D2H copies directly into CPU flat buffer slices. - # No GPU staging buffer → no temporary GPU memory spike. - with torch.cuda.stream(self._offload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - cpu_flat[off : off + n].copy_(local.reshape(-1), non_blocking=True) - - offloaded_bytes += grp["total"] * cpu_flat.element_size() - - # Wait for all D2H copies to land, then free GPU storage. - self._offload_stream.synchronize() - for t in self._managed: - storage = self._local(t).untyped_storage() - if storage.size() != 0: - storage.resize_(0) - else: - raise RuntimeError( - f"Tensor storage is already freed (size=0) before offload. " - f"This indicates a double-free or external interference. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - - if not self._logged: - logger.info( - "[CPUOffload] Offloaded %.2f MB (GPU → CPU)", - offloaded_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def reload(self): - """Per-tensor H2D from CPU flat buffer on the default stream. - - Runs on the current (default) CUDA stream to avoid stream - interaction issues with the parallel Muon pipeline. Since - pinned CPU memory is the source, the copies overlap with - GPU idle time between steps. - """ - if not self._managed or not self._initialized: - return - - reloaded_bytes = 0 - - # Re-allocate all GPU storages first. - for t in self._managed: - local = self._local(t) - storage = local.untyped_storage() - if storage.size() != 0: - raise RuntimeError( - f"Storage should have been freed (size=0) before reload, " - f"but got size={storage.size()}. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - storage.resize_(self._storage_nbytes[id(t)]) - - # Per-tensor H2D copies from CPU flat buffer slices. - # non_blocking=True with pinned source allows DMA overlap. - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off : off + n], non_blocking=True) - - reloaded_bytes += grp["total"] * cpu_flat.element_size() - - if not self._logged: - logger.info( - "[CPUOffload] Reloaded %.2f MB (CPU → GPU)", reloaded_bytes / (1024**2) - ) diff --git a/build/torch29-cxx11-cu128-x86_64-linux/distributed/utils.py b/build/torch29-cxx11-cu128-x86_64-linux/distributed/utils.py deleted file mode 100644 index 683096bcd50489e7a09959753fdbc29ba0f3b6a6..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/distributed/utils.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import (Placement, Shard, - _StridedShard) - - -def _is_shard(placement: Placement) -> bool: - """Check if a placement is a shard type (Shard or _StridedShard). - - In PyTorch 2.10+, _StridedShard no longer inherits from Shard, so - ``placement.is_shard()`` returns False for _StridedShard. This helper - handles both old and new hierarchies. - """ - return isinstance(placement, (Shard, _StridedShard)) - - -def get_slices_of_dtensor( - target: DTensor | torch.Tensor, - local_rank: int, - shard_mesh: DeviceMesh, - shard_placements: tuple[Placement], -) -> tuple[slice | torch.Tensor, ...]: - """ - Get per-dimension indices for a given rank's shard of the target tensor. - - Uses ``Shard.local_shard_size_and_offset`` and - ``_StridedShard.local_shard_size_and_offset`` for correct handling of - both contiguous and strided (non-contiguous) sharding. - - Args: - target (DTensor | torch.Tensor): The target tensor (for its shape). - local_rank (int): The local rank within the shard group. - shard_mesh (DeviceMesh): The shard mesh (only shard dimensions). - shard_placements (tuple[Placement]): The shard placements. - - Returns: - A tuple of indices (one per tensor dim). Each element is either: - - A ``slice`` (for contiguous or unsharded dims) - - A 1-D ``torch.LongTensor`` of indices (for strided sharding) - """ - - # find the global rank of the local rank in the shard mesh - rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank] - - rank_coords = (shard_mesh.mesh == rank).nonzero() - - assert len(rank_coords) == 1 - rank_coords = tuple(rank_coords[0].tolist()) - - assert len(rank_coords) == len(shard_placements) - - # Track per-shard-dim indices. - # None means "not yet sharded on this dim". - dim_indices: dict[int, torch.Tensor] = {} - - # Caution: Assuming replicate-to-shard of the shard mesh goes with - # left-to-right sharding. This is ensured by the sorting logic of - # construct_shard_mesh function. - for mesh_dim_idx, (rank_coord, placement) in enumerate( - zip(rank_coords, shard_placements)): - assert _is_shard(placement) - - num_chunks = shard_mesh.mesh.shape[mesh_dim_idx] - shard_dim = placement.dim - - # Current effective size on this dim (may already be sub-sharded) - if shard_dim in dim_indices: - curr_size = len(dim_indices[shard_dim]) - else: - curr_size = target.size()[shard_dim] - - # Compute indices for this level of sharding - if isinstance(placement, _StridedShard): - _shard_size, offsets = _StridedShard.local_shard_size_and_offset( - placement, - curr_size, - num_chunks, - rank_coord, - return_first_offset=False) - new_indices = torch.tensor(offsets, dtype=torch.long) - else: - shard_size, offset = Shard.local_shard_size_and_offset( - curr_size, num_chunks, rank_coord) - new_indices = torch.arange(offset, - offset + shard_size, - dtype=torch.long) - - # Compose with previous indices on this dim - if shard_dim in dim_indices: - dim_indices[shard_dim] = dim_indices[shard_dim][new_indices] - else: - dim_indices[shard_dim] = new_indices - - # Build result tuple - result: list[slice | torch.Tensor] = [] - for d in range(len(target.size())): - if d not in dim_indices: - result.append(slice(None)) - else: - indices = dim_indices[d] - # Convert contiguous indices to slice for efficiency - if len(indices) > 0: - start = indices[0].item() - expected = torch.arange(start, - start + len(indices), - dtype=torch.long) - if torch.equal(indices, expected): - result.append(slice(start, start + len(indices))) - else: - result.append(indices) - else: - result.append(slice(0, 0)) - - return tuple(result) - - -_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, - ProcessGroup]] = dict() - - -def construct_shard_mesh( - placements: tuple[Placement], - mesh: DeviceMesh, -) -> tuple[DeviceMesh, ProcessGroup, tuple[Placement, ...]]: - """Construct shard sub-mesh and ProcessGroup for all-to-all communication. - - Given a DTensor's placements and device mesh, extracts the "shard group" - — the set of ranks that together hold all shards of the same replica — - and creates a ProcessGroup for all-to-all among them. - - Steps: - 1. Sort placements: Replicate first, then Shard by (dim, granularity). - 2. Permute the mesh tensor to match the sorted order. - 3. Collapse Replicate dims → list of shard sub-meshes (one per replica). - 4. Create/retrieve a cached ProcessGroup for the current rank's sub-mesh. - - Example — 8 GPUs, mesh shape (2, 2, 2), - placements ``[Shard(0), Replicate, _StridedShard(0)]``:: - - Step 1 — Sort: [Replicate, _StridedShard(0), Shard(0)] - Permutation: [1, 2, 0] - - Step 2 — Permute mesh dims by [1, 2, 0]: - Original: Permuted: - [[[0,1],[2,3]], [[[0,2],[1,3]], - [[4,5],[6,7]]] [[4,6],[5,7]]] - - Step 3 — Unbind replicate dim (dim 0), giving 2 shard sub-meshes: - sub-mesh 0 = [[0,2],[1,3]] (replica group 0) - sub-mesh 1 = [[4,6],[5,7]] (replica group 1) - shard_placements = (_StridedShard(0), Shard(0)) - - Step 4 — Rank 0 → ProcessGroup([0,1,4,5]) - Rank 2 → ProcessGroup([2,3,6,7]) - - Returns: - ``(shard_mesh, process_group, shard_placements)`` - """ - my_rank = dist.get_rank() - assert mesh.mesh.device.type == 'cpu' - - # -- Fast path: 1D all-shard mesh → reuse existing PG. ---------------- - # Reuses the mesh's existing ProcessGroup directly, avoiding the - # overhead of dist.new_group(). The standard path below also handles - # subset calls safely via use_local_synchronization=True, but this - # fast path is still beneficial for the common 1D shard case. - if mesh.ndim == 1 and len(placements) == 1 and _is_shard(placements[0]): - key = (*mesh.mesh.shape, *mesh.mesh.flatten().tolist()) - if key not in _ranks_to_dist_cache: - _ranks_to_dist_cache[key] = (mesh, mesh.get_group()) - return (*_ranks_to_dist_cache[key], tuple(placements)) - - mesh_tensor = mesh.mesh.clone() - - # -- Step 1: Sort placements (Replicate first, then Shard by dim). ------ - # _StridedShard comes BEFORE regular Shard on the same dim so that - # get_slices_of_dtensor applies the outer sharding first, matching - # DTensor's left-to-right (outer-to-inner) composition order. - def _sort_key(item): - index, placement = item - assert not placement.is_partial(), "Partial placement not supported" - if placement.is_replicate(): - return (-1, 0, index) - assert _is_shard(placement), f"Unsupported: {type(placement)}" - split = (-1 / placement.split_factor if isinstance( - placement, _StridedShard) else 0) - return (placement.dim, split, index) - - indexed = sorted(enumerate(placements), key=_sort_key) - perm, sorted_placements = zip(*indexed) - - # -- Step 2: Permute mesh to match sorted placement order. -------------- - sorted_mesh = mesh_tensor.permute(perm) - - # -- Step 3: Collapse replicate dims → list of shard sub-meshes. -------- - # E.g. mesh (2, 3, 4, 4) with [R, R, S(0), S(1)] → 6 sub-meshes of (4, 4) - num_rep = sum(1 for p in sorted_placements if p.is_replicate()) - if num_rep > 0: - if num_rep > 1: - sorted_mesh = sorted_mesh.flatten(0, num_rep - 1) - shard_meshes = list(torch.unbind(sorted_mesh, dim=0)) - else: - shard_meshes = [sorted_mesh] - shard_placements = sorted_placements[num_rep:] - assert len(shard_placements) == len(set(shard_placements)) - - # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. -- - # Each rank only creates the group it belongs to, using - # use_local_synchronization=True so that only group members need to - # coordinate. This avoids deadlocks when different PP stages call - # construct_shard_mesh for different parameters. - def _cache_key(t: torch.Tensor) -> tuple: - return (*t.shape, *t.flatten().tolist()) - - my_key = None - for sm in shard_meshes: - if (my_rank == sm).any().item(): - key = _cache_key(sm) - assert my_key is None, "Rank appears in multiple shard groups" - my_key = key - if key not in _ranks_to_dist_cache: - pg = dist.new_group(sm.flatten().tolist(), - use_local_synchronization=True) - _ranks_to_dist_cache[key] = ( - DeviceMesh(device_type="cuda", mesh=sm), - pg, - ) - - return (*_ranks_to_dist_cache[my_key], shard_placements) diff --git a/build/torch29-cxx11-cu128-x86_64-linux/matmul_transpose_triton.py b/build/torch29-cxx11-cu128-x86_64-linux/matmul_transpose_triton.py deleted file mode 100644 index 792de23d82c3fb45fe33d397ab9b76a0787259d0..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/matmul_transpose_triton.py +++ /dev/null @@ -1,122 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], - restore_value=['y'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -@torch.library.custom_op("muon::matmul_transpose_assign", - mutates_args=("d_out", )) -def matmul_transpose_assign(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """Compute d_out = d_in @ d_in.T using an optimized Triton kernel.""" - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -@matmul_transpose_assign.register_fake -def _(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """FakeTensor impl: d_out is already allocated, mutation is declared.""" - pass diff --git a/build/torch29-cxx11-cu128-x86_64-linux/metadata.json b/build/torch29-cxx11-cu128-x86_64-linux/metadata.json deleted file mode 100644 index c55a35717622f1dd5c8ba376ea3a814cbcc10d78..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/metadata.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch29-cxx11-cu128-x86_64-linux/muon.py b/build/torch29-cxx11-cu128-x86_64-linux/muon.py deleted file mode 100644 index 14c0e22471fa6d47a51ed95e0e0c341dc18d5194..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/muon.py +++ /dev/null @@ -1,1068 +0,0 @@ -import logging -import types -from collections import defaultdict -from typing import Any - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Replicate, Shard -from torch.profiler import record_function - -from .adamw import _placement_cache, _tensor_cache, step_adamw -from .async_utils import run_pipeline -from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho, - get_default_muon_param_groups, is_expert_param, update_p) -from .cpu_offload import CPUOffloadPool -from .distributed.utils import (_is_shard, construct_shard_mesh, - get_slices_of_dtensor) -from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO, - _zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5_batched) -from .pipeline import muon_chunk_pipeline, prelaunch_first_gather -from .qk_clip import compute_scales, get_qk_clip_info, qk_clip - -logger = logging.getLogger(__name__) - - -def _expand_expert_params(names, params, expert_keys): - """Expand expert params by splitting on dim 0 (expert dimension). - - Params whose name matches any key in ``expert_keys`` are treated as - expert-parallel tensors. Their outermost dimension is the expert - dimension: an ``(E, out, in)`` tensor becomes ``E`` separate 2D - ``nn.Parameter`` views so that in-place updates propagate back to - the original storage. - - Non-expert params with ``ndim > 2`` trigger an ``AssertionError`` — - if they are expert params, their key must be added to ``expert_keys``. - - The grad must already be set on each expert param (e.g. after momentum). - - For DTensor expert params, placements that shard on dim 0 (expert dim) - are consumed by the split. Non-dim-0 shard placements (e.g. TP) are - preserved: each 2D slice is wrapped as a DTensor on the corresponding - submesh so the parallel pipeline handles the TP communication. - """ - expanded_names = [] - expanded_params = [] - - for n, p in zip(names, params): - is_expert = is_expert_param(n, expert_keys) - is_dtensor = isinstance(p.data, DTensor) - - if is_expert: - if is_dtensor: - logger.debug( - "[expand_expert] %s: expert DTensor, shape=%s, " - "placements=%s, mesh=%s, local_shape=%s", n, p.shape, - p.placements, p.device_mesh.mesh_dim_names, - p.to_local().shape) - else: - logger.debug( - "[expand_expert] %s: expert plain tensor, shape=%s", n, - p.data.shape) - - if not is_expert: - assert p.data.ndim <= 2, ( - f"Param {n} has ndim={p.data.ndim} but does not match " - f"expert_keys={expert_keys}. If this is an expert param, " - f"add its key to expert_keys.") - expanded_names.append(n) - expanded_params.append(p) - continue - - g = p.grad - assert g is not None, ( - f"Expert param {n} must have grad set before expansion") - - tp_mesh = None - tp_placements_2d = None - - if is_dtensor: - local_data = p.to_local() - local_grad = g.to_local() if isinstance(g, DTensor) else g - - # Find non-dim-0 shard placements (e.g. TP sharding). - # After splitting on dim 0, Shard(k) becomes Shard(k-1). - tp_dim_indices = [] - tp_placements_2d = [] - for i, pl in enumerate(p.placements): - if _is_shard(pl) and pl.dim != 0: - tp_dim_indices.append(i) - tp_placements_2d.append(Shard(pl.dim - 1)) - - if tp_dim_indices: - tp_dim_names = tuple(p.device_mesh.mesh_dim_names[i] - for i in tp_dim_indices) - if len(tp_dim_names) == 1: - tp_mesh = p.device_mesh[tp_dim_names[0]] - else: - tp_mesh = p.device_mesh[tp_dim_names] - else: - local_data = p.data - local_grad = g - - # Expand: split dim 0, reshape each slice to 2D. - num_local_experts = local_data.shape[0] - for i in range(num_local_experts): - slice_data = local_data[i] - slice_grad = local_grad[i] - - if tp_mesh is not None: - # Wrap as DTensor on TP submesh so the pipeline handles - # TP communication (gather/scatter across TP ranks). - dt_data = DTensor.from_local(slice_data, - device_mesh=tp_mesh, - placements=tp_placements_2d) - dt_grad = DTensor.from_local(slice_grad, - device_mesh=tp_mesh, - placements=tp_placements_2d) - expert_param = torch.nn.Parameter(dt_data, requires_grad=False) - expert_param.grad = dt_grad - else: - expert_param = torch.nn.Parameter(slice_data, - requires_grad=False) - expert_param.grad = slice_grad - - expanded_names.append(f"{n}[{i}]") - expanded_params.append(expert_param) - - p.grad = None # allow expert grad storage to be freed after pipeline - - return expanded_names, expanded_params - - -class Muon(torch.optim.Optimizer): - """ - Muon - MomentUm Orthogonalized by Newton-schulz - - Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- - processing step, in which each 2D parameter's update is replaced with the nearest orthogonal - matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has - the advantage that it can be stably run in bfloat16 on the GPU. - - Some warnings: - - We believe this optimizer is unlikely to work well for training with small batch size. - - We believe it may not work well for finetuning pretrained models, but we haven't tested this. - - Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. - lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) - momentum: The momentum used by the internal SGD. (0.95 is a good default) - nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) - ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. - Parameters that are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW instead. - adamw_lr: The learning rate for the internal AdamW. - adamw_betas: The betas for the internal AdamW. - adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - warmup_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher warmup_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. - chunk_size : Batch size of parameters to process in each - all2all gather/compute/scatter step. - Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified. - use_distributed_muon: Use distributed muon by Liu et al. (2024). - For testing purpose only. - expert_keys: List of strings to identify expert-parallel parameters. - If any key appears in a parameter's name, its outermost - dimension is treated as the expert dimension and expanded - into per-expert 2D params for Muon. For example, - ``expert_keys=["experts"]`` matches any param whose name - contains "experts". 3D+ params not matched by any key - will raise an error. - """ - - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config=None, - warmup_step=5, - chunk_size=-1, - use_distributed_muon=False, - expert_keys=None): - defaults = dict( - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - nesterov=nesterov, - ns_steps=ns_steps, - adamw_betas=adamw_betas, - adamw_eps=adamw_eps, - none_grad=none_grad, - use_muon=True, - ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) - super().__init__(params, defaults) - - self.debug = debug - self.clip_config = clip_config if clip_config is not None else { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100, - } - self.warmup_step = warmup_step - self.chunk_size = chunk_size - self.use_distributed_muon = use_distributed_muon - self.expert_keys = expert_keys - self.cpu_offload = False - self._cpu_offload_pool: CPUOffloadPool | None = None - self._offload_initialized = False - self._parallel_cache: dict[tuple[str, ...], dict] = {} - self._expert_expand_cache: dict[tuple[int, ...], dict] = {} - - def _calc_flops(self, G, steps): - assert len(G.shape) == 2 - M, N = G.shape - if M > N: - M, N = N, M - - return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) - - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - p.placements, p.device_mesh) - - return shard_mesh, shard_pg, shard_placements - - def init_state_and_assign_params(self, names, params, group, qk_logits): - param_to_state = {} - param_to_flops = {} - - total_flops = 0 - for p in params: - g = p.grad - if g is None: - continue - assert g.ndim == 2, "Muon only supports 2D parameters." - - flops = self._calc_flops(g, group["ns_steps"]) - param_to_flops[id(p)] = flops - total_flops += flops - - if self.debug: - logger.debug("Total TFLOPs for Muon: %.2f TFLOPs", - total_flops / 1e12) - - paired = list(zip(names, params)) - - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) - - round_robin = 0 - mesh = ordered_params[0].device_mesh - placements = ordered_params[0].placements - - shard_mesh, shard_pg, shard_placements = self.get_shard_mesh( - ordered_params[0]) - shard_mesh_flattened = shard_mesh.mesh.flatten() - num_ranks = dist.get_world_size(group=shard_pg) - - for n, p in zip(ordered_names, ordered_params): - if mesh != p.device_mesh: - raise ValueError("All parameters must be on the same mesh.") - if placements != p.placements: - raise ValueError("All parameters must have same placements.") - - worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks - round_robin = (round_robin + 1) % len(shard_mesh_flattened) - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - # Precompute per-rank indices and numels for all-to-all. - rank_indices: dict[int, tuple] = {} - rank_numels: dict[int, int] = {} - for r in range(num_ranks): - indices = get_slices_of_dtensor(p, r, shard_mesh, - shard_placements) - rank_indices[r] = indices - numel = 1 - for idx, dim_size in zip(indices, p.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) - rank_numels[r] = numel - - param_to_state[id(p)] = _muon_state( - worker_rank=worker_rank, - process_group=shard_pg, - rank_indices=rank_indices, - rank_numels=rank_numels, - name=n, - qk_clip_state=qk_clip_state, - ) - - return param_to_state, ordered_params - - def base(self, names, params, group, lr, weight_decay, qk_logits): - # Momentum is already applied by _step_muon before this method. - for n, p in zip(names, params): - g = p.grad - if g is None: - continue - - u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) - - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - def distributed_muon( - self, - names: list[str], - params: list[torch.nn.Parameter], - group: dict[str, Any], - lr: float, - weight_decay: float, - qk_logits: list[torch.Tensor | DTensor] | None, - ): - """Batched Distributed Muon — for testing/correctness verification only. - - Uses all-gather to reconstruct full tensors, computes Newton-Schulz on - the full grad, then slices back to local shards. This is simpler but - slower than the parallel pipeline (all2all) path, so it serves as a - reference implementation for verifying correctness. - """ - with record_function("distributed_muon"): - # Momentum is already applied by _step_muon before this method. - ns_steps = group["ns_steps"] - - # Separate plain tensors (no communication) from DTensors. - plain_names, plain_params = [], [] - dtensor_names, dtensor_params = [], [] - for n, p in zip(names, params): - if p.grad is None: - continue - if isinstance(p.data, DTensor): - dtensor_names.append(n) - dtensor_params.append(p) - else: - plain_names.append(n) - plain_params.append(p) - - # Process plain tensors per-param (no communication). - for n, p in zip(plain_names, plain_params): - u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE), - steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, - qk_logits) - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - if not dtensor_params: - return - - # Group DTensors by (placements, mesh) for batched all-gather. - placement_groups: dict[tuple, - tuple[list, - list]] = defaultdict(lambda: ([], [])) - for n, p in zip(dtensor_names, dtensor_params): - key = (p.placements, p.device_mesh) - placement_groups[key][0].append(n) - placement_groups[key][1].append(p) - - logger.info( - "distributed_muon: %d placement groups, %d total dtensors", - len(placement_groups), len(dtensor_params)) - - for (placements, mesh), (grp_names, - grp_params) in placement_groups.items(): - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - placements, mesh) - rank = dist.get_rank(shard_pg) - world_size = dist.get_world_size(shard_pg) - - logger.info(" group: %d params, placements=%s, world_size=%d", - len(grp_params), placements, world_size) - - # Separate params that can be batched (all shard dims evenly - # divisible) from those needing per-param full_tensor - # (e.g. MoE gate weights with fewer rows than shard ranks). - # all_gather_into_tensor requires equal buffer sizes across - # ranks, so uneven splits must use DTensor full_tensor(). - batch_names, batch_params = [], [] - single_names, single_params = [], [] - for n, p in zip(grp_names, grp_params): - even = all(p.shape[pl.dim] % - shard_mesh.mesh.shape[dim_idx] == 0 - for dim_idx, pl in enumerate(shard_placements)) - if even: - batch_names.append(n) - batch_params.append(p) - else: - single_names.append(n) - single_params.append(p) - - # Process uneven-split params per-param via full_tensor(). - for n, p in zip(single_names, single_params): - with record_function("distributed_muon::newton_schulz"): - g_full = p.grad.full_tensor().to(COMM_DTYPE) - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - if not batch_params: - continue - - logger.info(" batched=%d, single=%d", len(batch_params), - len(single_params)) - - # Concat all local grad shards into a single flat buffer. - with record_function("distributed_muon::gather"): - grad_locals = [ - p.grad.to_local().to(COMM_DTYPE).flatten() - for p in batch_params - ] - numels = [g.numel() for g in grad_locals] - grad_concat = torch.cat(grad_locals) - del grad_locals - - # Single all-gather (replaces N separate full_tensor). - grad_gathered = torch.empty( - grad_concat.numel() * world_size, - dtype=COMM_DTYPE, - device="cuda", - ) - dist.all_gather_into_tensor(grad_gathered, - grad_concat, - group=shard_pg) - - total_numel = grad_concat.numel() - del grad_concat - - # Precompute per-param offsets within the concat buffer. - offsets = [] - off = 0 - for ne in numels: - offsets.append(off) - off += ne - - # Per-param: reconstruct full grad → NS → local update. - for i, (n, p) in enumerate(zip(batch_names, batch_params)): - with record_function("distributed_muon::newton_schulz"): - g_full = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - for r in range(world_size): - r_start = r * total_numel + offsets[i] - shard = grad_gathered[r_start:r_start + numels[i]] - indices = get_slices_of_dtensor( - p, r, shard_mesh, shard_placements) - g_full[indices] = shard.reshape( - g_full[indices].shape) - - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - def _setup_parallel(self, names, params, group, qk_logits): - """Compute (or retrieve cached) parallel pipeline metadata. - - Returns: - (ordered_params, param_to_state, rank, chunk_size) - """ - cache_key = tuple(names) - - if cache_key not in self._parallel_cache: - # First call: compute metadata and populate cache. - param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - shard_pg = param_to_state[id(ordered_params[0])].process_group - rank = dist.get_rank(group=shard_pg) - - if self.chunk_size == -1: - shard_ranks = dist.get_world_size(shard_pg) - chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO - elif self.chunk_size > 0: - chunk_size = self.chunk_size - else: - raise ValueError( - "chunk_size must be -1 or a positive integer.") - - ordered_names = [ - param_to_state[id(p)].name for p in ordered_params - ] - name_to_state = { - param_to_state[id(p)].name: param_to_state[id(p)] - for p in ordered_params - } - self._parallel_cache[cache_key] = { - 'ordered_names': ordered_names, - 'name_to_state': name_to_state, - 'rank': rank, - 'chunk_size': chunk_size, - } - else: - # Cached path: rebuild param_to_state with current id(p) keys. - cache = self._parallel_cache[cache_key] - rank = cache['rank'] - chunk_size = cache['chunk_size'] - - name_to_param = dict(zip(names, params)) - ordered_params = [name_to_param[n] for n in cache['ordered_names']] - - param_to_state = {} - for p, n in zip(ordered_params, cache['ordered_names']): - cached_state = cache['name_to_state'][n] - param_to_state[id(p)] = _muon_state( - worker_rank=cached_state.worker_rank, - process_group=cached_state.process_group, - rank_indices=cached_state.rank_indices, - rank_numels=cached_state.rank_numels, - name=n, - qk_clip_state=get_qk_clip_info(self.clip_config, n, - qk_logits), - ) - - return ordered_params, param_to_state, rank, chunk_size - - def parallel(self, - names, - params, - group, - lr, - weight_decay, - qk_logits, - prelaunch_gather=None): - """ - Perform a parallel optimization step using Muon. - - Parameters are chunked and each chunk is processed by a - :func:`muon_chunk_pipeline` generator. :func:`run_pipeline` - interleaves multiple chunks so that communication and computation - overlap across chunks (the same overlap previously achieved by the - warmup + main-loop index scheduling). - - If ``prelaunch_gather`` is provided, it is passed to the first - chunk's generator to skip re-launching the already in-flight - A2A gather. - """ - - # Momentum is already applied by _step_muon before this method. - - ordered_params, param_to_state, rank, chunk_size = ( - self._setup_parallel(names, params, group, qk_logits)) - - def pipelines(): - first = True - for start in range(0, len(ordered_params), chunk_size): - chunk = ordered_params[start:start + chunk_size] - if chunk: - kwargs = dict( - params=chunk, - param_to_state=param_to_state, - rank=rank, - ns_steps=group["ns_steps"], - lr=lr, - weight_decay=weight_decay, - none_grad=group["none_grad"], - ) - if first and prelaunch_gather is not None: - kwargs['prelaunch_gather'] = prelaunch_gather - first = False - yield muon_chunk_pipeline(**kwargs) - - with record_function("muon::pipeline"): - run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1) - - def _step_muon(self, group, qk_logits=None): - params = group["params"] - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - # Apply momentum to all params before routing/expansion. - # Batched using _foreach_* ops (compiled, fullgraph=True). - with record_function("muon::momentum"): - active_params = [p for p in params if p.grad is not None] - if active_params: - # Ensure momentum buffers exist (avoid zeros_like when already present). - for p in active_params: - if "momentum_buffer" not in self.state[p]: - self.state[p]["momentum_buffer"] = torch.zeros_like( - p.grad) - - # Extract local tensors for compiled batch function. - local_grads = [ - p.grad._local_tensor - if isinstance(p.grad, DTensor) else p.grad - for p in active_params - ] - local_bufs = [ - self.state[p]["momentum_buffer"]._local_tensor - if isinstance(self.state[p]["momentum_buffer"], DTensor) - else self.state[p]["momentum_buffer"] - for p in active_params - ] - - # Wrap momentum as tensor for torch.compile. - batch_pre_ortho(local_grads, local_bufs, - torch.tensor(momentum), group["nesterov"]) - - # For non-nesterov, the result is the momentum buffer. - if not group["nesterov"]: - for p in active_params: - p.grad = self.state[p]["momentum_buffer"] - - # Identify batched experts for deferred NS. - # Detection is cheap (condition checks only); actual NS compute is - # deferred so it can overlap with the first chunk's A2A gather. - deferred_expert_work = [] - if self.expert_keys: - batched_expert_indices = [] - for i, (n, p) in enumerate(zip(names, params)): - if not (is_expert_param(n, self.expert_keys) - and p.grad is not None): - continue - # Eligible: plain tensor, or DTensor with no non-dim-0 shards. - if isinstance(p.data, DTensor): - has_tp = any( - _is_shard(pl) and pl.dim != 0 for pl in p.placements) - if has_tp: - continue - batched_expert_indices.append(i) - - if batched_expert_indices: - # Save refs for deferred NS; free grads from param list. - for i in batched_expert_indices: - p = params[i] - g = p.grad - local_g = (g._local_tensor - if isinstance(g, DTensor) else g) - local_data = (p.data._local_tensor if isinstance( - p.data, DTensor) else p.data) - deferred_expert_work.append((local_data, local_g)) - p.grad = None - - # Remove batched experts from lists before expansion. - keep = sorted( - set(range(len(params))) - set(batched_expert_indices)) - names = [names[i] for i in keep] - params = [params[i] for i in keep] - - def _run_deferred_expert_ns(): - """Execute deferred batched expert NS.""" - if not deferred_expert_work: - return - with record_function("muon::batched_expert_ns"): - ns_steps = group["ns_steps"] - for local_data, local_g in deferred_expert_work: - u = zeropower_via_newtonschulz5_batched( - local_g.to(COMM_DTYPE), steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:]) - local_data.mul_(1 - lr * weight_decay) - local_data.add_(u, alpha=-adjusted_lr) - - # Expand expert params by splitting on dim 0. - logger.debug("[_step_muon] before expand: %d params, expert_keys=%s", - len(params), self.expert_keys) - if self.expert_keys: - cache_key = tuple(id(p) for p in params) - cache = self._expert_expand_cache.get(cache_key) - - if cache is None: - # Cold path: full expansion + build cache metadata. - exp_names, exp_params = _expand_expert_params( - names, params, self.expert_keys) - - # Build per-expert-group info for hot-path grad updates. - grad_info = [] - exp_idx = 0 - for orig_idx, (n, p) in enumerate(zip(names, params)): - if not is_expert_param(n, self.expert_keys): - exp_idx += 1 - continue - - is_dt = isinstance(p.data, DTensor) - num_experts = (p.to_local() if is_dt else p.data).shape[0] - - # Detect TP mesh from the first expanded expert param. - tp_mesh = None - tp_pls = None - sample = exp_params[exp_idx] - if isinstance(sample.data, DTensor): - tp_mesh = sample.data.device_mesh - tp_pls = list(sample.data.placements) - - grad_info.append((orig_idx, num_experts, exp_idx, is_dt, - tp_mesh, tp_pls)) - exp_idx += num_experts - - self._expert_expand_cache[cache_key] = { - 'names': exp_names, - 'params': exp_params, - 'grad_info': grad_info, - } - names, params = exp_names, exp_params - else: - # Hot path: reuse cached params, only update expert grads. - for (orig_idx, num_experts, exp_start, is_dt, tp_mesh, - tp_pls) in cache['grad_info']: - p = params[orig_idx] - g = p.grad - local_grad = (g.to_local() - if is_dt and isinstance(g, DTensor) else g) - for i in range(num_experts): - expert_p = cache['params'][exp_start + i] - sg = local_grad[i] - if tp_mesh is not None: - expert_p.grad = DTensor.from_local( - sg, device_mesh=tp_mesh, placements=tp_pls) - else: - expert_p.grad = sg - p.grad = None - - names = cache['names'] - params = cache['params'] - else: - names, params = _expand_expert_params(names, params, - self.expert_keys) - logger.debug("[_step_muon] after expand: %d params", len(params)) - - param_dtensors = [] - name_dtensors = [] - - param_tensors = [] - name_tensors = [] - - # distributed_muon is a reference implementation for testing only. - # The parallel pipeline (all2all) path below is the production path. - if self.use_distributed_muon: - _run_deferred_expert_ns() - self.distributed_muon(names=names, - params=params, - group=group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits) - return - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - logger.debug( - "[route] %s → base (DTensor all-Replicate), " - "shape=%s, placements=%s", n, p.shape, p.placements) - param_tensors.append(p) - name_tensors.append(n) - else: - logger.debug( - "[route] %s → parallel (DTensor), shape=%s, " - "placements=%s, mesh=%s", n, p.shape, p.placements, - p.device_mesh.mesh_dim_names) - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - logger.debug("[route] %s → base (plain tensor), shape=%s", n, - p.data.shape) - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError(f"Unsupported parameter type: {type(p.data)}") - - logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, " - f"{len(param_tensors)} Tensors → base") - - def group_dtensors(dtensors, names): - # To support different placements, we group parameters by placements - # and run parallel Muon on each group. - - placement_to_params = defaultdict(lambda: ([], [])) - - assert len(dtensors) == len(names) - for p, n in zip(dtensors, names): - placement_to_params[tuple([p.placements, - p.device_mesh])][0].append(n) - placement_to_params[tuple([p.placements, - p.device_mesh])][1].append(p) - return placement_to_params - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - dtensor_group = group_dtensors(param_dtensors, name_dtensors) - - # Pre-launch the first chunk's A2A gather so that the NCCL - # communication overlaps with the (deferred) batched expert NS - # compute on the default CUDA stream. - prelaunch = None - if deferred_expert_work: - first_names, first_params = next(iter(dtensor_group.values())) - ordered, pts, rnk, csz = self._setup_parallel( - first_names, first_params, group, qk_logits) - first_chunk = ordered[:csz] - if first_chunk: - prelaunch = prelaunch_first_gather(first_chunk, pts, rnk, - group["none_grad"]) - - _run_deferred_expert_ns() - - first_group = True - for _, (names, params) in dtensor_group.items(): - pg = prelaunch if first_group else None - first_group = False - self.parallel( - names, - params, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - prelaunch_gather=pg, - ) - else: - _run_deferred_expert_ns() - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - ) - - def _register_states_for_offload(self): - """Register all optimizer state tensors with the CPU offload pool. - - Called once after the first step when states have been lazily created. - Offloads all param states (momentum buffers for Muon, moment1/moment2 - for AdamW) to free GPU memory between steps. - """ - pool = self._cpu_offload_pool - tracked = 0 - for group in self.param_groups: - for p in group["params"]: - if p not in self.state: - continue - state = self.state[p] - if group.get("use_muon", False): - if "momentum_buffer" in state: - pool.track(state["momentum_buffer"]) - tracked += 1 - else: - if "moment1" in state: - pool.track(state["moment1"]) - if "moment2" in state: - pool.track(state["moment2"]) - tracked += 1 - logger.info("[CPUOffload] Registered %d param states for offload", - tracked) - - @torch.no_grad - def step(self, closure=None, qk_logits=None): - """Perform a single optimization step. - - Args: - closure (Callable, optional): A closure that reevaluates the model - and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). - """ - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # H2D: reload optimizer states from CPU before computation. - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload() - - logger.debug("[Muon.step] expert_keys=%s, %d param groups", - self.expert_keys, len(self.param_groups)) - - for i, group in enumerate(self.param_groups): - if group["use_muon"]: - logger.debug("[Muon.step] group %d: use_muon=True, %d params", - i, len(group["params"])) - self._step_muon(group, qk_logits=qk_logits) - else: - logger.debug( - "[Muon.step] group %d: use_muon=False (AdamW), %d params", - i, len(group["params"])) - step_adamw(self.state, group) - - # D2H: offload optimizer states to CPU after computation. - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - return loss - - # ------------------------------------------------------------------ - # CPU offload public helpers - # ------------------------------------------------------------------ - - def turn_on_cpu_offload(self): - """Enable CPU offload for optimizer states.""" - if self.cpu_offload: - return - logger.info("[Muon] turn_on_cpu_offload") - self.cpu_offload = True - if not self.state: - return - self._cpu_offload_pool = CPUOffloadPool() - self._offload_initialized = False - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_off_cpu_offload(self): - """Disable CPU offload and keep optimizer states resident on GPU.""" - if not self.cpu_offload: - return - logger.info("[Muon] turn_off_cpu_offload") - if self._offload_initialized: - self._cpu_offload_pool.reload() - torch.cuda.current_stream().synchronize() - self._cpu_offload_pool = None - self._offload_initialized = False - self.cpu_offload = False - - # ------------------------------------------------------------------ - # Checkpoint support for cpu_offload - # ------------------------------------------------------------------ - - def state_dict(self) -> dict: - if self.cpu_offload: - raise RuntimeError( - "Muon.state_dict() requires turn_off_cpu_offload() before checkpoint save." - ) - return super().state_dict() - - def load_state_dict(self, state_dict: dict) -> None: - if self.cpu_offload: - raise RuntimeError( - "Muon.load_state_dict() requires turn_off_cpu_offload() before checkpoint load." - ) - super().load_state_dict(state_dict) - - # Invalidate adamw.py's module-level tensor caches so that - # the next step rebuilds them with the newly loaded state tensors. - _placement_cache.clear() - _tensor_cache.clear() diff --git a/build/torch29-cxx11-cu128-x86_64-linux/newton_schulz.py b/build/torch29-cxx11-cu128-x86_64-linux/newton_schulz.py deleted file mode 100644 index d939264b69a34e7a3fa78859f34dc265a1159d59..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/newton_schulz.py +++ /dev/null @@ -1,240 +0,0 @@ -from itertools import repeat -from math import inf, sqrt - -import numpy as np -import torch - -from .matmul_transpose_triton import matmul_transpose_assign - -COMM_DTYPE = torch.bfloat16 -DEFAULT_CHUNK_SIZE_RATIO = 4 - - -def _optimal_quintic(l, u, max_iter=1000): - """ - Use the simplified Remez algorithm to find the optimal odd quintic approximant - to the constant function x -> 1 over the interval [l, u]. - - Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum - approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the - two interior equioscillation nodes q, r until convergence. Returns the - closed-form equioscillating solution when l ≈ u. - - Raises ValueError if any intermediate value (a, b, c, E, q, r) is non-finite - (NaN or inf). Raises RuntimeError if convergence is not reached within - max_iter iterations. - """ - assert 0 <= l <= u - if 1 - 5e-6 <= l / u: - return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5) - q = (3 * l + u) / 4 - r = (l + 3 * u) / 4 - E = inf - for _ in range(max_iter): - old_E = E - LHS = np.array( - [ - [l, l**3, l**5, 1], - [q, q**3, q**5, -1], - [r, r**3, r**5, 1], - [u, u**3, u**5, -1], - ] - ) - a, b, c, E = np.linalg.solve(LHS, np.ones(4)) - if not np.all(np.isfinite([a, b, c, E])): - raise ValueError( - f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}" - ) - q, r = np.sqrt( - (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / (10 * c) - ) - if not np.all(np.isfinite([q, r])): - raise ValueError(f"_optimal_quintic: non-finite node update q={q}, r={r}") - if abs(old_E - E) <= 1e-15: - break - else: - raise RuntimeError( - f"_optimal_quintic: did not converge after {max_iter} iterations" - ) - return float(a), float(b), float(c) - - -def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0): - """ - Compute the Polar Express coefficient series for `num_iters` quintic iterations. - - Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that - compose to map singular values from [l, 1] toward 1. At each step: - 1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion` - prevents near-zero singular values from stalling by raising the effective - lower bound; if it is active (cushion*u > l), the coefficients are - rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u]. - 2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the - last iteration, providing numerical headroom at the cost of a slightly slower - final convergence step. - 3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1). - - Returns a list of (a, b, c) tuples, one per iteration. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - u = 1 - assert 0 <= l <= u - safety_factor = 1 + safety_factor_eps - coefficients = [] - for iter in range(num_iters): - a, b, c = _optimal_quintic(max(l, cushion * u), u) - if cushion * u > l: - pl = a * l + b * l**3 + c * l**5 - pu = a * u + b * u**3 + c * u**5 - rescaler = 2 / (pl + pu) - a *= rescaler - b *= rescaler - c *= rescaler - if iter < num_iters - 1: - a /= safety_factor - b /= safety_factor**3 - c /= safety_factor**5 - coefficients.append((a, b, c)) - l = a * l + b * l**3 + c * l**5 - u = 2 - l - return coefficients - - -# Precomputed Polar Express coefficients (a, b, c) for 10 quintic Newton-Schulz -# iterations. Each tuple is the minimax-optimal (Remez/equioscillation) odd quintic -# approximant to x->1 over the current singular-value interval, computed once at -# import time and reused across all optimizer steps. -# -# Contrast with the former hardcoded NS coefficients (5 fixed tuples): -# - Former: empirically tuned to maximize slope at zero; did not converge -# singular values to 1, yielding US'V^T with S' ~ Uniform(0.5, 1.5) instead -# of the true polar factor UV^T. -# - Polar Express: analytically optimal per step, adapting to the shrinking -# singular-value interval [l, u] as iterations progress; converges all -# singular values to 1, producing the exact polar factor UV^T. -_coeffs_list = _optimal_composition( - l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02 -) - - -# This code is adapted from: -# KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py) -# NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress) -# matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon) -@torch.no_grad() -def _zeropower_via_newtonschulz5(G, steps): - """ - Compute the polar factor of G via the Polar Express method. - - Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c) - are the Polar Express coefficients from `_coeffs_list`. Each step is the - optimal odd quintic approximant to x -> 1 over the current singular-value - interval, minimizing the maximum approximation error (Remez / minimax criterion). - The composition maps singular values from [l, 1] to near 1, producing the - polar factor (orthogonal factor in the polar decomposition G = UP). - - `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2, - cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE - X = G # no manual typecast - - if G.size(0) > G.size(1): - X = X.T - - X = X / (X.norm() + 1e-7) - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - # Perform the NS iterations - for a, b, c in hs: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(0) > G.size(1): - X = X.T - - return X - - -@torch.no_grad() -def _zeropower_via_newtonschulz5_batched(G, steps): - """Batched polar factor computation for 3D (E, out, in) tensors. - - Same algorithm as ``_zeropower_via_newtonschulz5`` but uses - ``torch.bmm`` / ``torch.baddbmm`` instead of the 2D Triton kernel, - processing all E expert matrices in a single batched call. - """ - assert len(G.shape) == 3 - assert G.dtype == COMM_DTYPE - X = G - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - # Per-expert Frobenius norm. - X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) - - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - for a, b, c in hs: - buf1 = torch.bmm(X, X.transpose(-2, -1)) - buf2 = torch.bmm(buf1, buf1.transpose(-2, -1)) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.baddbmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - return X - - -_ns_per_shape: dict[tuple[int, ...], callable] = {} -_use_compile = True - - -def set_ns_compile(enabled: bool): - """Toggle torch.compile for Newton-Schulz iteration.""" - global _use_compile - _use_compile = enabled - - -def zeropower_via_newtonschulz5(G, steps=5): - if not _use_compile: - return _zeropower_via_newtonschulz5(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() - - -def zeropower_via_newtonschulz5_batched(G, steps=5): - """Compile-cached batched Newton-Schulz for 3D expert tensors.""" - if not _use_compile: - return _zeropower_via_newtonschulz5_batched(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile( - _zeropower_via_newtonschulz5_batched, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() diff --git a/build/torch29-cxx11-cu128-x86_64-linux/optimizer/__init__.py b/build/torch29-cxx11-cu128-x86_64-linux/optimizer/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/optimizer/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu128-x86_64-linux/pipeline.py b/build/torch29-cxx11-cu128-x86_64-linux/pipeline.py deleted file mode 100644 index c0c2d515856182d8d15ad27dd4e4e093b29397d6..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/pipeline.py +++ /dev/null @@ -1,468 +0,0 @@ -import logging -from typing import Generator - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -from .core import _muon_state, adjust_lr_for_muon -from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5 -from .qk_clip import compute_scales - -logger = logging.getLogger(__name__) - -# ====================================================================== -# Stage helpers -# ====================================================================== - - -def _launch_gather( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Allocate gather buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_gather``). - gathered_grads: ``{id(p): empty_tensor}`` for owned params, - ``None`` for non-owned. - recv_counts: Per-source-rank element counts. - """ - # Allocate gathered-grad buffers - gathered_grads: dict[int, torch.Tensor | None] = {} - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - gathered_grads[id(p)] = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - else: - gathered_grads[id(p)] = None - - # Build send buffer – batch grad copies via torch.cat - # (1-2 fused kernels vs N individual narrow().copy_() calls). - send_counts = [0] * num_ranks - for p in params: - state = param_to_state[id(p)] - send_counts[state.worker_rank] += state.rank_numels[rank] - - total_send = sum(send_counts) - if total_send > 0: - # Group grad slices by destination rank in a single pass. - dst_to_grads = [[] for _ in range(num_ranks)] - for p in params: - state = param_to_state[id(p)] - n = state.rank_numels[rank] - if n > 0: - g = p.grad.to_local() - dst_to_grads[state.worker_rank].append(g.reshape(-1)) - - # Flatten in dst order and cat once. - all_slices = [] - for dst in range(num_ranks): - all_slices.extend(dst_to_grads[dst]) - send_buf = torch.cat(all_slices) - if send_buf.dtype != COMM_DTYPE: - send_buf = send_buf.to(COMM_DTYPE) - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += state.rank_numels[src] - recv_counts[src] = total - - recv_buf = torch.empty(sum(recv_counts), dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - logger.debug(f"send_buf size: {send_buf.numel()}, " - f"recv_buf size: {recv_buf.numel()}, " - f"recv_counts: {recv_counts}, " - f"send_counts: {send_counts}, " - f"process_group: {str(process_group)}") - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, gathered_grads, recv_counts - - -def _complete_gather( - recv_buf: torch.Tensor, - recv_counts: list[int], - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - param_to_state: dict[int, _muon_state], - rank: int, -) -> None: - """Reconstruct gathered grads from the recv buffer (in-place).""" - off = 0 - for src in range(len(recv_counts)): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - - indices = state.rank_indices[src] - - shard_view = gathered_grads[id(p)][indices] - n = shard_view.numel() - if n == 0: - continue - - sg = recv_buf.narrow(0, off + inner_off, n) - sg = sg.reshape(shard_view.shape) - gathered_grads[id(p)][indices] = sg - - inner_off += n - assert inner_off == block - off += block - - -def _compute_ns( - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - ns_steps: int, -) -> dict[int, torch.Tensor | None]: - """Run Newton-Schulz orthogonalization on owned parameters. - - Returns: - computed_us: ``{id(p): orthogonalized_update}`` for owned params. - """ - computed_us: dict[int, torch.Tensor | None] = {} - for p in owned_params: - u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps) - gathered_grads[id(p)] = None # free gathered grad - computed_us[id(p)] = u - return computed_us - - -def _launch_scatter( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, - computed_us: dict[int, torch.Tensor | None], -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor], list[int]]: - """Allocate scatter buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_scatter``). - scattered_us: Empty dict, populated by ``_complete_scatter`` with - zero-copy views into ``recv_buf``. - recv_counts: Per-source-rank element counts. - """ - # scattered_us is populated by _complete_scatter with zero-copy views - # into recv_buf, avoiding N empty_like allocations + N copy_ calls. - # Pre-seed entries for params whose local shard is empty (rank_numels == 0) - # so _update_params can iterate all params without KeyError. - scattered_us: dict[int, torch.Tensor] = {} - for p in params: - if param_to_state[id(p)].rank_numels[rank] == 0: - scattered_us[id(p)] = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - # Build send buffer – batch via torch.cat - # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls). - send_counts = [0] * num_ranks - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - for dst_rank in range(num_ranks): - send_counts[dst_rank] += state.rank_numels[dst_rank] - - total_send = sum(send_counts) - if total_send > 0: - # Cache u_full conversions to avoid redundant .to() per dst_rank. - u_fulls = {} - for p in owned_params: - u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous() - - # Collect slices in dst order (matches all-to-all send layout). - all_slices = [] - for dst_rank in range(num_ranks): - for p in owned_params: - state = param_to_state[id(p)] - su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten() - if su.numel() > 0: - all_slices.append(su) - - send_buf = torch.cat(all_slices) if all_slices else torch.empty( - 0, dtype=COMM_DTYPE, device="cuda") - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += state.rank_numels[rank] - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, scattered_us, recv_counts - - -def _complete_scatter( - recv_buf: torch.Tensor, - recv_counts: list[int], - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], -) -> None: - """Populate scattered_us with zero-copy views into recv_buf. - - Instead of pre-allocating tensors and copying, we assign views directly - from ``recv_buf``. This eliminates N ``empty_like`` + N ``copy_`` calls. - The underlying storage of ``recv_buf`` is kept alive through the views - until ``scattered_us`` is cleared after ``_update_params``. - """ - off = 0 - for src in range(len(recv_counts)): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = state.rank_numels[rank] - if n == 0: - continue - - scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - - inner_off += n - - assert inner_off == block - off += block - - -def _update_params( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], - lr: float, - weight_decay: float, -) -> None: - """Apply weight decay, Muon update, and optional QK clipping. - - Uses batched ``_foreach_mul_`` for weight decay and batched - ``_foreach_add_`` for the Muon update, grouping parameters by - adjusted_lr to minimize kernel launches while preserving float32 - precision for the alpha scaling. - """ - if not params: - return - - # Batched weight decay: p *= (1 - lr * wd) — single fused kernel. - p_locals = [p._local_tensor for p in params] - torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay) - - # Group params by adjusted_lr so _foreach_add_ can use a single - # alpha per group (preserves float32 precision for alpha scaling). - lr_groups: dict[float, tuple[list, list]] = {} - for p in params: - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - if adjusted_lr not in lr_groups: - lr_groups[adjusted_lr] = ([], []) - lr_groups[adjusted_lr][0].append(p._local_tensor) - lr_groups[adjusted_lr][1].append(scattered_us[id(p)]) - - for adjusted_lr, (p_group, u_group) in lr_groups.items(): - torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr) - - # QK clipping – applied directly on the local tensor to - # avoid DTensor sharding-propagation issues with _StridedShard. - for p in params: - state = param_to_state[id(p)] - if state.qk_clip_state is None: - continue - scales_full = compute_scales(p, state.qk_clip_state) - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = state.rank_indices[rank][0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - -# ====================================================================== -# Pre-launch helper for overlapping first chunk's gather with other work. -# ====================================================================== - - -@torch.no_grad() -def prelaunch_first_gather( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - none_grad: bool, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Launch the first chunk's A2A gather early for overlap with other compute. - - Call this *before* expensive GPU work (e.g. batched expert NS) so that - the NCCL all-to-all runs concurrently on the NCCL stream while the - default stream executes compute. - - Returns the same 4-tuple that ``_launch_gather`` produces, which should - be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - with record_function("muon::prelaunch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - return work, recv_buf, gathered_grads, recv_counts - - -# ====================================================================== -# Main generator – thin orchestrator that wires stages together. -# ====================================================================== - - -@torch.no_grad() -def muon_chunk_pipeline( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - ns_steps: int, - lr: float, - weight_decay: float, - none_grad: bool, - prelaunch_gather: tuple | None = None, -) -> Generator[None, None, None]: - """Process one chunk of parameters through the full Muon pipeline. - - Stages: gather -> compute (Newton-Schulz) -> scatter -> update. - - Each ``yield`` lets :func:`run_pipeline` interleave other chunks so - that communication and computation overlap across chunks. Async - communication is launched via ``async_op=True`` and completed after - the yield with ``work.wait()``. - - Overlap happens because :func:`run_pipeline` admits one new chunk - per iteration (staggered admission). While chunk *N* does NS - compute on the default CUDA stream, chunk *N+1*'s async all-to-all - runs concurrently on the NCCL stream — no separate ``comm_stream`` - is required. - - If ``prelaunch_gather`` is provided, the gather was already launched - by :func:`prelaunch_first_gather` and we skip launching it again. - - Yields exactly **2** times: - - 1. After launching async all-to-all gather (or immediately if pre-launched). - 2. After launching async all-to-all scatter. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - if prelaunch_gather is not None: - # Gather was pre-launched; none_grad already handled by caller. - work, recv_buf, gathered_grads, recv_counts = prelaunch_gather - else: - # Normal path: launch async gather. - with record_function("muon::launch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - yield # --- YIELD 1: other chunks can launch their gather --- - - with record_function("muon::wait_gather"): - work.wait() - _complete_gather(recv_buf, recv_counts, owned_params, gathered_grads, - param_to_state, rank) - del recv_buf - - # Stage 3: Newton-Schulz orthogonalization. - with record_function("muon::newton_schulz"): - computed_us = _compute_ns(owned_params, gathered_grads, ns_steps) - gathered_grads.clear() - - # Stages 4-5: launch async scatter. - with record_function("muon::launch_scatter"): - work, recv_buf, scattered_us, recv_counts = _launch_scatter( - params, owned_params, param_to_state, rank, num_ranks, - process_group, computed_us) - computed_us.clear() - - yield # --- YIELD 2: other chunks can launch their scatter --- - - with record_function("muon::wait_scatter"): - work.wait() - _complete_scatter(recv_buf, recv_counts, params, param_to_state, rank, - scattered_us) - del recv_buf - - # Stage 6: apply parameter updates. - with record_function("muon::update_params"): - _update_params(params, param_to_state, rank, scattered_us, lr, - weight_decay) - scattered_us.clear() diff --git a/build/torch29-cxx11-cu128-x86_64-linux/qk_clip.py b/build/torch29-cxx11-cu128-x86_64-linux/qk_clip.py deleted file mode 100644 index 2aba711b3004b7f09e7141da7ef834bd61cc2430..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/qk_clip.py +++ /dev/null @@ -1,198 +0,0 @@ -import logging -import math -from dataclasses import dataclass - -import torch -from torch.distributed.tensor import DTensor - -from .core import normalize_fqn - -logger = logging.getLogger(__name__) - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - and return (kind, layer_index). - - Supported kinds: - MHA/GQA: 'wq', 'wk', 'q_proj', 'k_proj' - MLA: 'wq_b' (Q up-proj), 'wkv_b' (KV up-proj) - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.1.attn.wq_b.weight' -> ('wq_b', 1) - 'model.0.attn.wkv_b.weight' -> ('wkv_b', 0) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = normalize_fqn(name).split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: str | None # 'wq'/'q_proj'/'wq_b' or 'wk'/'k_proj'/'wkv_b' or None - indices: list[int] # which heads to consider for clipping - head_dim: int # from config (qk_head_dim for MLA wq_b) - threshold: float # from config - logit: torch.Tensor | None - - # MLA-specific fields - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 - - -def get_qk_clip_info(clip_config, n, qk_logits): - """Extract QK clipping info for a named parameter. - - Args: - clip_config: QK clipping configuration dict (or None). - MHA/GQA keys: head_dim, threshold, q_indices, k_indices - MLA extra keys: is_mla=True, qk_nope_head_dim, qk_rope_head_dim, v_head_dim - n: Parameter name string. - qk_logits: Dict mapping layer indices to logit tensors (or None). - - Returns: - QKClipInfo instance with clipping configuration for this parameter. - """ - if clip_config is None: - return None - - head_dim = clip_config.get('head_dim') - threshold = clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - # In TP settings, qk_logits may be DTensor - # We convert it to full tensor here for simplicity - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - -def compute_scales(p, qk_clip_state): - """Compute per-head scaling factors for QK clipping. - - Returns scales tensor (√γ per head) if any head exceeds threshold, else None. - For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim. - """ - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - # Check if any head exceeds threshold before allocating. - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - - if not head_scales: - return None - - # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows - if qk_clip_state.is_mla and kind == 'wkv_b': - effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim - else: - effective_head_dim = head_dim - - H_global = p.shape[0] // effective_head_dim - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale - return scales_full - - -def qk_clip(p, scales, info): - """Apply per-head scaling to a Q/K projection weight matrix. - - Args: - p: Parameter (nn.Parameter or raw tensor). - scales: [n_heads] tensor, each element = √γ_h. - info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions. - - MLA sub-region scaling per Algorithm 1 (MuonClip): - wq_b: q_nope rows → √γ, q_pe rows → γ - wkv_b: k_nope rows → √γ, v rows → unchanged - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not info.is_mla: - # MHA/GQA: uniform √γ applied to all rows in each head - W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: vectorized sub-region scaling within each head - if info.kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_head_dim = qk_nope + info.qk_rope_head_dim - W_3d = W.view(-1, qk_head_dim, W.shape[1]) # [H, qk_head_dim, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # q_nope → √γ - W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1, - 1)) # q_pe → γ - - elif info.kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - W_3d = W.view(-1, kv_stride, W.shape[1]) # [H, kv_stride, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # k_nope → √γ - # v rows: not touched (k_R shared rotary unchanged) diff --git a/build/torch29-cxx11-cu130-x86_64-linux/__init__.py b/build/torch29-cxx11-cu130-x86_64-linux/__init__.py deleted file mode 100644 index 239c7a65f8293e7d0df28f05fce645af56d628c0..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .muon import Muon - -__all__ = [ - "Muon", -] diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_ops.py b/build/torch29-cxx11-cu130-x86_64-linux/_ops.py deleted file mode 100644 index 034bff088659b5df6f6d401feb18c89dc5f33b29..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _optimizer_8d53b78_dirty -ops = torch.ops._optimizer_8d53b78_dirty - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_optimizer_8d53b78_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so b/build/torch29-cxx11-cu130-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so deleted file mode 100755 index c3a9a8d14613666ccf77a88386d75632e689b3bd..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b73946a6e5c0366cfb776d7d553e6343a354392da029564aff8ad0d961ffa25b -size 2000456 diff --git a/build/torch29-cxx11-cu130-x86_64-linux/adamw.py b/build/torch29-cxx11-cu130-x86_64-linux/adamw.py deleted file mode 100644 index b5a95816a9f5b9e1889eaadae65373bfbced809a..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/adamw.py +++ /dev/null @@ -1,271 +0,0 @@ -import logging -from collections import defaultdict -from typing import cast - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -def fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: float | torch.Tensor, - weight_decay: float, - eps: float, - maximize: bool, -) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: dict | None = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) - - -def _to_local(t): - """Unwrap DTensor to local tensor for fused ops.""" - return t._local_tensor if isinstance(t, DTensor) else t - - -# --------------------------------------------------------------------------- -# Caches for eliminating per-step Python overhead. -# -# Placement grouping and tensor list assembly are identical every step -# (params don't change placement, moment/step tensors are the same objects -# after initialisation). We cache them keyed by id() of the param list -# stored in param_groups (stable across steps). -# -# Only gradients change each step and must be collected fresh. -# --------------------------------------------------------------------------- - -# id(group["params"]) → dict[placement_key, list[param]] -_placement_cache: dict[int, dict[tuple, list]] = {} - -# id(placement_group_list) → (params_local, moment1, moment2, state_steps) -_tensor_cache: dict[int, tuple[list, list, list, list]] = {} - - -def _step_adamw_params_slow(optimizer_state, params, group): - """Uncached fallback for the rare case where some params lack grads.""" - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - g = p.grad - if g is None: - continue - state = optimizer_state[p] - params_with_grads.append(_to_local(p)) - grads.append(_to_local(g)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - if not params_with_grads: - return - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw_params(optimizer_state, params, group): - """Run fused AdamW on a list of parameters sharing the same placement. - - After the first call, cached tensor lists (params_local, moment1, - moment2, state_steps) are reused — only gradients are collected fresh. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - params: List of parameters to update. - group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay. - """ - # Collect grads — the only thing that changes each step. - with record_function("adamw::collect_grads"): - grads = [] - for p in params: - g = p.grad - if g is None: - # Rare: fall back to slow path that filters per-param. - _step_adamw_params_slow(optimizer_state, params, group) - return - grads.append(_to_local(g)) - - tensor_key = id(params) - if tensor_key not in _tensor_cache: - with record_function("adamw::init_tensor_cache"): - params_local = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - state = optimizer_state[p] - params_local.append(_to_local(p)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(p.grad) - state["moment2"] = torch.zeros_like(p.grad) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - _tensor_cache[tensor_key] = (params_local, moment1, moment2, - state_steps) - - params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key] - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - with record_function("adamw::fused_adamw"): - fused_adamw( - params_local, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw(optimizer_state, group): - """Dispatch AdamW step, grouping parameters by type and placement. - - Placement grouping is cached after the first call since params never - change their placement between steps. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - group: Parameter group dict. - """ - params = group["params"] - placement_key = id(params) - - if placement_key not in _placement_cache: - with record_function("adamw::group_by_placement"): - placement_to_params: dict[tuple, - list[torch.Tensor]] = defaultdict(list) - for p in params: - match p: - case DTensor(): - logger.debug( - "[AdamW] DTensor param: shape=%s, placements=%s, " - "mesh=%s, grad=%s", p.shape, p.placements, - p.device_mesh.mesh_dim_names, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple( - [p.placements, p.device_mesh])].append(p) - case torch.Tensor(): - logger.debug( - "[AdamW] plain param: shape=%s, grad=%s", p.shape, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple([torch.Tensor, - None])].append(p) - - logger.debug("[AdamW] %d placement groups, %d total params", - len(placement_to_params), len(params)) - - _placement_cache[placement_key] = dict(placement_to_params) - - for group_params in _placement_cache[placement_key].values(): - step_adamw_params(optimizer_state, group_params, group) diff --git a/build/torch29-cxx11-cu130-x86_64-linux/async_utils.py b/build/torch29-cxx11-cu130-x86_64-linux/async_utils.py deleted file mode 100644 index a45c530ac9cad88e3555ec1047a6aa59f225347e..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/async_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import logging -from typing import Generator - -logger = logging.getLogger(__name__) - - -class _Task: - """Internal: wraps a generator, advances one yield at a time.""" - - def __init__(self, generator: Generator[None, None, None], index: int): - self._generator = generator - self._index = index - self._steps_completed = 0 - self.step() # run to first yield - - def step(self) -> bool: - try: - next(self._generator) - self._steps_completed += 1 - logger.debug("pipeline[%d] completed stage %d", self._index, - self._steps_completed) - return True - except StopIteration: - logger.debug("pipeline[%d] finished after %d stages", self._index, - self._steps_completed) - return False - - def close(self): - self._generator.close() - - -def run_pipeline( - pipelines: Generator[Generator[None, None, None], None, None], - max_concurrent: int, -) -> None: - """Run generator-based pipelines with bounded concurrency. - - Each pipeline is a generator that yields at stage boundaries. - The runtime interleaves pipelines so communication and computation - overlap across chunks. - """ - if max_concurrent <= 0: - raise ValueError(f"max_concurrent must be > 0, got {max_concurrent}") - - have_new = True - task_index = 0 - previous_tasks: list[_Task] = [] - - try: - while have_new or previous_tasks: - running_tasks: list[_Task] = [] - - # Admit one new pipeline per iteration (staggered admission). - # Admitting one at a time ensures that while chunk N does NS - # compute on the default stream, chunk N+1's NCCL all-to-all - # runs concurrently on the NCCL stream — creating real - # communication/computation overlap on the GPU. - if have_new and len(previous_tasks) < max_concurrent: - try: - gen = next(pipelines) - task = _Task(gen, task_index) - task_index += 1 - running_tasks.append(task) - except StopIteration: - have_new = False - - # Advance every previously-yielded task by one step. - for task in previous_tasks: - if task.step(): - running_tasks.append(task) - - previous_tasks = running_tasks - except BaseException: - # Clean up all in-flight generators to release GPU resources. - for task in previous_tasks: - task.close() - raise diff --git a/build/torch29-cxx11-cu130-x86_64-linux/core.py b/build/torch29-cxx11-cu130-x86_64-linux/core.py deleted file mode 100644 index c69d515afef305ad0ed66374095fa2d2468d99cc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/core.py +++ /dev/null @@ -1,219 +0,0 @@ -import logging -import math -from dataclasses import dataclass -from typing import List - -import torch -from torch.distributed import ProcessGroup -from torch.distributed.tensor import DTensor - -# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into -# parameter FQNs. Activation checkpointing similarly inserts -# "_checkpoint_wrapped_module". Strip these so name-based matching (skip_keys, -# expert_keys, QK layer parsing) works regardless of wrapper nesting. -_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) - -logger = logging.getLogger(__name__) - - -def normalize_fqn(name: str) -> str: - """Strip torch.compile / checkpoint wrapper components from a parameter FQN.""" - return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS) - - -@dataclass -class _muon_state: - worker_rank: int - process_group: ProcessGroup - rank_indices: dict[int, tuple] # local_rank -> per-dim indices - rank_numels: dict[int, int] # local_rank -> numel - name: str - qk_clip_state: torch.Tensor | None = None - - -def _batch_momentum( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update (no nesterov).""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - - -def _batch_momentum_nesterov( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update with nesterov correction.""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - nesterov_terms = torch._foreach_mul(momentum_bufs, momentum) - torch._foreach_add_(grads, nesterov_terms) - - -_compiled_momentum: dict[bool, callable] = {} -_use_momentum_compile = True - - -def set_momentum_compile(enabled: bool): - """Toggle torch.compile for batched momentum.""" - global _use_momentum_compile - _use_momentum_compile = enabled - - -def batch_pre_ortho( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, - nesterov: bool, -) -> None: - """Batched momentum update on lists of plain tensors. - - Mirrors dion's ``muon_update_pre_orthogonalize``. - Inputs must be plain CUDA tensors (not DTensor). - Modifies ``momentum_bufs`` and (for nesterov) ``grads`` in-place. - - When compile is enabled, uses separately compiled functions for - nesterov=True/False to avoid graph breaks from the branch. - """ - fn = _batch_momentum_nesterov if nesterov else _batch_momentum - if _use_momentum_compile: - if nesterov not in _compiled_momentum: - _compiled_momentum[nesterov] = torch.compile(fn) - fn = _compiled_momentum[nesterov] - fn(grads, momentum_bufs, momentum) - - -def _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay): - """Weight-decay + update on plain tensors. - - Not compiled: per-param @torch.compile caused ~0.25ms TorchDynamo cache - lookup per call × 256+ params = massive overhead. The pipeline path uses - batched _foreach_* ops instead; this function remains for base() and - distributed_muon(). - """ - p_data.mul_(1 - lr * weight_decay) - p_data.add_(u_data, alpha=-adjusted_lr) - - -def update_p(p, u, lr, adjusted_lr, weight_decay): - """Apply weight decay and orthogonalized update to parameter. - - Args: - p: Parameter (torch.nn.Parameter or DTensor). - u: Orthogonalized update tensor. - lr: Base learning rate. - adjusted_lr: Size-adjusted learning rate. - weight_decay: Weight decay coefficient. - """ - # Unwrap Parameter -> underlying data tensor. - p_data = p.data if isinstance(p, torch.nn.Parameter) else p - # Unwrap DTensor -> local CUDA tensor for compiled kernel. - if isinstance(p_data, DTensor): - p_data = p_data._local_tensor - u_data = u._local_tensor if isinstance(u, DTensor) else u - _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay) - - -def adjust_lr_for_muon(lr, param_shape): - """Scale learning rate based on parameter matrix dimensions. - - Args: - lr: Base learning rate. - param_shape: Shape of the parameter tensor. - - Returns: - Adjusted learning rate. - """ - A, B = param_shape[:2] - # We adjust the learning rate and weight decay based on the size of the parameter matrix - # as described in the paper - adjusted_ratio = 0.2 * math.sqrt(max(A, B)) - adjusted_lr = lr * adjusted_ratio - return adjusted_lr - - -def _match_key(parts, key): - """Check if key matches as contiguous components in parts. - - Single-component keys (e.g. "experts") match any single component. - Multi-component keys (e.g. "experts.w1") match as a contiguous subsequence. - """ - key_parts = key.split(".") - key_len = len(key_parts) - if key_len == 1: - return key in parts - return any(parts[i:i + key_len] == key_parts - for i in range(len(parts) - key_len + 1)) - - -def is_expert_param(name, expert_keys): - """Check if a parameter name matches any expert key (component-level).""" - if not expert_keys: - return False - parts = normalize_fqn(name).split(".") - return any(_match_key(parts, key) for key in expert_keys) - - -def default_is_muon(name, x, expert_keys=None): - normalized = normalize_fqn(name) - parts = normalized.split(".") - skip_keys = [ - "embed_tokens", - "lm_head", - "tok_embeddings", - "output", - "mhc_attn", - "mhc_ffn", - "lambda_proj", - ] - if any(key in parts for key in skip_keys): - logger.info( - "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d", - normalized, name, x.ndim) - return False - effective_ndim = x.ndim - is_expert = is_expert_param(name, expert_keys) - if is_expert: - effective_ndim -= 1 - result = effective_ndim >= 2 - logger.info( - "[is_muon] %s (orig: %s): ndim=%d, expert=%s, effective_ndim=%d → %s", - normalized, name, x.ndim, is_expert, effective_ndim, - "Muon" if result else "AdamW") - return result - - -def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None): - if is_muon_func is None: - is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys) - - muon_params, muon_names = [], [] - non_muon_params, non_muon_names = [], [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - non_muon_names.append(n) - - logger.info("[param_groups] expert_keys=%s, Muon=%d, AdamW=%d", - expert_keys, len(muon_names), len(non_muon_names)) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] diff --git a/build/torch29-cxx11-cu130-x86_64-linux/cpu_offload.py b/build/torch29-cxx11-cu130-x86_64-linux/cpu_offload.py deleted file mode 100644 index fb5e69154a1d4a6c884491413a37a9acf0f66c80..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/cpu_offload.py +++ /dev/null @@ -1,206 +0,0 @@ -"""CPU offloading for optimizer states. - -Manages a pinned CPU memory pool and async CUDA streams to offload -optimizer state tensors (momentum buffers, Adam moments) to CPU between -optimizer steps, freeing GPU memory. - -All tracked tensors are packed into a single flat pinned CPU buffer -(per dtype). D2H and H2D copies are performed per-tensor directly -between individual GPU tensors and their slice of the CPU flat buffer -— no GPU staging buffer is allocated, so there is **no temporary GPU -memory spike** during offload or reload. - -Individual tensor storages are freed after offload via -``untyped_storage().resize_(0)``, preserving tensor identity so -downstream caches remain valid. -""" - -import logging -from collections import defaultdict - -import torch -from torch.distributed.tensor import DTensor - -logger = logging.getLogger(__name__) - - -class CPUOffloadPool: - """Pinned CPU memory pool for async optimizer state offloading. - - Tracked tensors are grouped by dtype. Each group gets a single flat - pinned CPU buffer. D2H / H2D copies are per-tensor (into slices of - the flat buffer) to avoid allocating a GPU staging buffer. - """ - - def __init__(self): - self._managed: list[torch.Tensor] = [] - self._storage_nbytes: dict[int, int] = {} # id(t) → bytes - - # Per-dtype group: populated on first offload. - # dtype → dict with keys: - # "indices" : list[int] managed-list indices - # "offsets" : list[tuple[int,int]] (start, numel) in flat buf - # "total" : int total numel - # "cpu_flat" : Tensor pinned CPU buffer - self._groups: dict[torch.dtype, dict] = {} - - self._offload_stream: torch.cuda.Stream | None = None - self._device: torch.device | None = None - self._initialized: bool = False - self._logged: bool = False - - # ------------------------------------------------------------------ - @staticmethod - def _local(t: torch.Tensor) -> torch.Tensor: - """Unwrap DTensor to its local CUDA tensor.""" - return t._local_tensor if isinstance(t, DTensor) else t - - def _ensure_stream(self): - if self._offload_stream is None: - self._offload_stream = torch.cuda.Stream(device=self._device) - - # ------------------------------------------------------------------ - def track(self, tensor: torch.Tensor): - """Register a GPU tensor for CPU offloading. Idempotent.""" - tid = id(tensor) - if tid in self._storage_nbytes: - return - local = self._local(tensor) - if self._device is None: - self._device = local.device - storage = local.untyped_storage() - # Skip tensors with empty storage (e.g. empty FSDP shards) - if storage.size() == 0: - return - self._storage_nbytes[tid] = storage.size() - self._managed.append(tensor) - - # ------------------------------------------------------------------ - def _init_buffers(self): - """Build per-dtype flat buffers on first offload.""" - # Group managed tensors by dtype. - dtype_map: dict[torch.dtype, list[tuple[int, int]]] = defaultdict(list) - for idx, t in enumerate(self._managed): - local = self._local(t) - dtype_map[local.dtype].append((idx, local.numel())) - - total_cpu_bytes = 0 - for dtype, entries in dtype_map.items(): - offsets: list[tuple[int, int]] = [] - indices: list[int] = [] - off = 0 - for idx, n in entries: - indices.append(idx) - offsets.append((off, n)) - off += n - cpu_flat = torch.empty(off, dtype=dtype, device="cpu", pin_memory=True) - self._groups[dtype] = { - "indices": indices, - "offsets": offsets, - "total": off, - "cpu_flat": cpu_flat, - } - total_cpu_bytes += off * cpu_flat.element_size() - - self._initialized = True - logger.info( - "[CPUOffload] Pool initialized: %d tensors, %d dtype group(s), " - "%.2f MB pinned CPU memory", - len(self._managed), - len(self._groups), - total_cpu_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def offload(self): - """Per-tensor async D2H into CPU flat buffer, then free GPU storage.""" - if not self._managed: - return - if not self._initialized: - self._init_buffers() - self._ensure_stream() - - # Offload stream waits for compute to finish. - compute_event = torch.cuda.current_stream(self._device).record_event() - self._offload_stream.wait_event(compute_event) - - offloaded_bytes = 0 - - # Per-tensor D2H copies directly into CPU flat buffer slices. - # No GPU staging buffer → no temporary GPU memory spike. - with torch.cuda.stream(self._offload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - cpu_flat[off : off + n].copy_(local.reshape(-1), non_blocking=True) - - offloaded_bytes += grp["total"] * cpu_flat.element_size() - - # Wait for all D2H copies to land, then free GPU storage. - self._offload_stream.synchronize() - for t in self._managed: - storage = self._local(t).untyped_storage() - if storage.size() != 0: - storage.resize_(0) - else: - raise RuntimeError( - f"Tensor storage is already freed (size=0) before offload. " - f"This indicates a double-free or external interference. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - - if not self._logged: - logger.info( - "[CPUOffload] Offloaded %.2f MB (GPU → CPU)", - offloaded_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def reload(self): - """Per-tensor H2D from CPU flat buffer on the default stream. - - Runs on the current (default) CUDA stream to avoid stream - interaction issues with the parallel Muon pipeline. Since - pinned CPU memory is the source, the copies overlap with - GPU idle time between steps. - """ - if not self._managed or not self._initialized: - return - - reloaded_bytes = 0 - - # Re-allocate all GPU storages first. - for t in self._managed: - local = self._local(t) - storage = local.untyped_storage() - if storage.size() != 0: - raise RuntimeError( - f"Storage should have been freed (size=0) before reload, " - f"but got size={storage.size()}. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - storage.resize_(self._storage_nbytes[id(t)]) - - # Per-tensor H2D copies from CPU flat buffer slices. - # non_blocking=True with pinned source allows DMA overlap. - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off : off + n], non_blocking=True) - - reloaded_bytes += grp["total"] * cpu_flat.element_size() - - if not self._logged: - logger.info( - "[CPUOffload] Reloaded %.2f MB (CPU → GPU)", reloaded_bytes / (1024**2) - ) diff --git a/build/torch29-cxx11-cu130-x86_64-linux/distributed/utils.py b/build/torch29-cxx11-cu130-x86_64-linux/distributed/utils.py deleted file mode 100644 index 683096bcd50489e7a09959753fdbc29ba0f3b6a6..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/distributed/utils.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import (Placement, Shard, - _StridedShard) - - -def _is_shard(placement: Placement) -> bool: - """Check if a placement is a shard type (Shard or _StridedShard). - - In PyTorch 2.10+, _StridedShard no longer inherits from Shard, so - ``placement.is_shard()`` returns False for _StridedShard. This helper - handles both old and new hierarchies. - """ - return isinstance(placement, (Shard, _StridedShard)) - - -def get_slices_of_dtensor( - target: DTensor | torch.Tensor, - local_rank: int, - shard_mesh: DeviceMesh, - shard_placements: tuple[Placement], -) -> tuple[slice | torch.Tensor, ...]: - """ - Get per-dimension indices for a given rank's shard of the target tensor. - - Uses ``Shard.local_shard_size_and_offset`` and - ``_StridedShard.local_shard_size_and_offset`` for correct handling of - both contiguous and strided (non-contiguous) sharding. - - Args: - target (DTensor | torch.Tensor): The target tensor (for its shape). - local_rank (int): The local rank within the shard group. - shard_mesh (DeviceMesh): The shard mesh (only shard dimensions). - shard_placements (tuple[Placement]): The shard placements. - - Returns: - A tuple of indices (one per tensor dim). Each element is either: - - A ``slice`` (for contiguous or unsharded dims) - - A 1-D ``torch.LongTensor`` of indices (for strided sharding) - """ - - # find the global rank of the local rank in the shard mesh - rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank] - - rank_coords = (shard_mesh.mesh == rank).nonzero() - - assert len(rank_coords) == 1 - rank_coords = tuple(rank_coords[0].tolist()) - - assert len(rank_coords) == len(shard_placements) - - # Track per-shard-dim indices. - # None means "not yet sharded on this dim". - dim_indices: dict[int, torch.Tensor] = {} - - # Caution: Assuming replicate-to-shard of the shard mesh goes with - # left-to-right sharding. This is ensured by the sorting logic of - # construct_shard_mesh function. - for mesh_dim_idx, (rank_coord, placement) in enumerate( - zip(rank_coords, shard_placements)): - assert _is_shard(placement) - - num_chunks = shard_mesh.mesh.shape[mesh_dim_idx] - shard_dim = placement.dim - - # Current effective size on this dim (may already be sub-sharded) - if shard_dim in dim_indices: - curr_size = len(dim_indices[shard_dim]) - else: - curr_size = target.size()[shard_dim] - - # Compute indices for this level of sharding - if isinstance(placement, _StridedShard): - _shard_size, offsets = _StridedShard.local_shard_size_and_offset( - placement, - curr_size, - num_chunks, - rank_coord, - return_first_offset=False) - new_indices = torch.tensor(offsets, dtype=torch.long) - else: - shard_size, offset = Shard.local_shard_size_and_offset( - curr_size, num_chunks, rank_coord) - new_indices = torch.arange(offset, - offset + shard_size, - dtype=torch.long) - - # Compose with previous indices on this dim - if shard_dim in dim_indices: - dim_indices[shard_dim] = dim_indices[shard_dim][new_indices] - else: - dim_indices[shard_dim] = new_indices - - # Build result tuple - result: list[slice | torch.Tensor] = [] - for d in range(len(target.size())): - if d not in dim_indices: - result.append(slice(None)) - else: - indices = dim_indices[d] - # Convert contiguous indices to slice for efficiency - if len(indices) > 0: - start = indices[0].item() - expected = torch.arange(start, - start + len(indices), - dtype=torch.long) - if torch.equal(indices, expected): - result.append(slice(start, start + len(indices))) - else: - result.append(indices) - else: - result.append(slice(0, 0)) - - return tuple(result) - - -_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, - ProcessGroup]] = dict() - - -def construct_shard_mesh( - placements: tuple[Placement], - mesh: DeviceMesh, -) -> tuple[DeviceMesh, ProcessGroup, tuple[Placement, ...]]: - """Construct shard sub-mesh and ProcessGroup for all-to-all communication. - - Given a DTensor's placements and device mesh, extracts the "shard group" - — the set of ranks that together hold all shards of the same replica — - and creates a ProcessGroup for all-to-all among them. - - Steps: - 1. Sort placements: Replicate first, then Shard by (dim, granularity). - 2. Permute the mesh tensor to match the sorted order. - 3. Collapse Replicate dims → list of shard sub-meshes (one per replica). - 4. Create/retrieve a cached ProcessGroup for the current rank's sub-mesh. - - Example — 8 GPUs, mesh shape (2, 2, 2), - placements ``[Shard(0), Replicate, _StridedShard(0)]``:: - - Step 1 — Sort: [Replicate, _StridedShard(0), Shard(0)] - Permutation: [1, 2, 0] - - Step 2 — Permute mesh dims by [1, 2, 0]: - Original: Permuted: - [[[0,1],[2,3]], [[[0,2],[1,3]], - [[4,5],[6,7]]] [[4,6],[5,7]]] - - Step 3 — Unbind replicate dim (dim 0), giving 2 shard sub-meshes: - sub-mesh 0 = [[0,2],[1,3]] (replica group 0) - sub-mesh 1 = [[4,6],[5,7]] (replica group 1) - shard_placements = (_StridedShard(0), Shard(0)) - - Step 4 — Rank 0 → ProcessGroup([0,1,4,5]) - Rank 2 → ProcessGroup([2,3,6,7]) - - Returns: - ``(shard_mesh, process_group, shard_placements)`` - """ - my_rank = dist.get_rank() - assert mesh.mesh.device.type == 'cpu' - - # -- Fast path: 1D all-shard mesh → reuse existing PG. ---------------- - # Reuses the mesh's existing ProcessGroup directly, avoiding the - # overhead of dist.new_group(). The standard path below also handles - # subset calls safely via use_local_synchronization=True, but this - # fast path is still beneficial for the common 1D shard case. - if mesh.ndim == 1 and len(placements) == 1 and _is_shard(placements[0]): - key = (*mesh.mesh.shape, *mesh.mesh.flatten().tolist()) - if key not in _ranks_to_dist_cache: - _ranks_to_dist_cache[key] = (mesh, mesh.get_group()) - return (*_ranks_to_dist_cache[key], tuple(placements)) - - mesh_tensor = mesh.mesh.clone() - - # -- Step 1: Sort placements (Replicate first, then Shard by dim). ------ - # _StridedShard comes BEFORE regular Shard on the same dim so that - # get_slices_of_dtensor applies the outer sharding first, matching - # DTensor's left-to-right (outer-to-inner) composition order. - def _sort_key(item): - index, placement = item - assert not placement.is_partial(), "Partial placement not supported" - if placement.is_replicate(): - return (-1, 0, index) - assert _is_shard(placement), f"Unsupported: {type(placement)}" - split = (-1 / placement.split_factor if isinstance( - placement, _StridedShard) else 0) - return (placement.dim, split, index) - - indexed = sorted(enumerate(placements), key=_sort_key) - perm, sorted_placements = zip(*indexed) - - # -- Step 2: Permute mesh to match sorted placement order. -------------- - sorted_mesh = mesh_tensor.permute(perm) - - # -- Step 3: Collapse replicate dims → list of shard sub-meshes. -------- - # E.g. mesh (2, 3, 4, 4) with [R, R, S(0), S(1)] → 6 sub-meshes of (4, 4) - num_rep = sum(1 for p in sorted_placements if p.is_replicate()) - if num_rep > 0: - if num_rep > 1: - sorted_mesh = sorted_mesh.flatten(0, num_rep - 1) - shard_meshes = list(torch.unbind(sorted_mesh, dim=0)) - else: - shard_meshes = [sorted_mesh] - shard_placements = sorted_placements[num_rep:] - assert len(shard_placements) == len(set(shard_placements)) - - # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. -- - # Each rank only creates the group it belongs to, using - # use_local_synchronization=True so that only group members need to - # coordinate. This avoids deadlocks when different PP stages call - # construct_shard_mesh for different parameters. - def _cache_key(t: torch.Tensor) -> tuple: - return (*t.shape, *t.flatten().tolist()) - - my_key = None - for sm in shard_meshes: - if (my_rank == sm).any().item(): - key = _cache_key(sm) - assert my_key is None, "Rank appears in multiple shard groups" - my_key = key - if key not in _ranks_to_dist_cache: - pg = dist.new_group(sm.flatten().tolist(), - use_local_synchronization=True) - _ranks_to_dist_cache[key] = ( - DeviceMesh(device_type="cuda", mesh=sm), - pg, - ) - - return (*_ranks_to_dist_cache[my_key], shard_placements) diff --git a/build/torch29-cxx11-cu130-x86_64-linux/matmul_transpose_triton.py b/build/torch29-cxx11-cu130-x86_64-linux/matmul_transpose_triton.py deleted file mode 100644 index 792de23d82c3fb45fe33d397ab9b76a0787259d0..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/matmul_transpose_triton.py +++ /dev/null @@ -1,122 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], - restore_value=['y'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -@torch.library.custom_op("muon::matmul_transpose_assign", - mutates_args=("d_out", )) -def matmul_transpose_assign(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """Compute d_out = d_in @ d_in.T using an optimized Triton kernel.""" - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -@matmul_transpose_assign.register_fake -def _(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """FakeTensor impl: d_out is already allocated, mutation is declared.""" - pass diff --git a/build/torch29-cxx11-cu130-x86_64-linux/metadata.json b/build/torch29-cxx11-cu130-x86_64-linux/metadata.json deleted file mode 100644 index c55a35717622f1dd5c8ba376ea3a814cbcc10d78..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/metadata.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch29-cxx11-cu130-x86_64-linux/muon.py b/build/torch29-cxx11-cu130-x86_64-linux/muon.py deleted file mode 100644 index 14c0e22471fa6d47a51ed95e0e0c341dc18d5194..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/muon.py +++ /dev/null @@ -1,1068 +0,0 @@ -import logging -import types -from collections import defaultdict -from typing import Any - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Replicate, Shard -from torch.profiler import record_function - -from .adamw import _placement_cache, _tensor_cache, step_adamw -from .async_utils import run_pipeline -from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho, - get_default_muon_param_groups, is_expert_param, update_p) -from .cpu_offload import CPUOffloadPool -from .distributed.utils import (_is_shard, construct_shard_mesh, - get_slices_of_dtensor) -from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO, - _zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5_batched) -from .pipeline import muon_chunk_pipeline, prelaunch_first_gather -from .qk_clip import compute_scales, get_qk_clip_info, qk_clip - -logger = logging.getLogger(__name__) - - -def _expand_expert_params(names, params, expert_keys): - """Expand expert params by splitting on dim 0 (expert dimension). - - Params whose name matches any key in ``expert_keys`` are treated as - expert-parallel tensors. Their outermost dimension is the expert - dimension: an ``(E, out, in)`` tensor becomes ``E`` separate 2D - ``nn.Parameter`` views so that in-place updates propagate back to - the original storage. - - Non-expert params with ``ndim > 2`` trigger an ``AssertionError`` — - if they are expert params, their key must be added to ``expert_keys``. - - The grad must already be set on each expert param (e.g. after momentum). - - For DTensor expert params, placements that shard on dim 0 (expert dim) - are consumed by the split. Non-dim-0 shard placements (e.g. TP) are - preserved: each 2D slice is wrapped as a DTensor on the corresponding - submesh so the parallel pipeline handles the TP communication. - """ - expanded_names = [] - expanded_params = [] - - for n, p in zip(names, params): - is_expert = is_expert_param(n, expert_keys) - is_dtensor = isinstance(p.data, DTensor) - - if is_expert: - if is_dtensor: - logger.debug( - "[expand_expert] %s: expert DTensor, shape=%s, " - "placements=%s, mesh=%s, local_shape=%s", n, p.shape, - p.placements, p.device_mesh.mesh_dim_names, - p.to_local().shape) - else: - logger.debug( - "[expand_expert] %s: expert plain tensor, shape=%s", n, - p.data.shape) - - if not is_expert: - assert p.data.ndim <= 2, ( - f"Param {n} has ndim={p.data.ndim} but does not match " - f"expert_keys={expert_keys}. If this is an expert param, " - f"add its key to expert_keys.") - expanded_names.append(n) - expanded_params.append(p) - continue - - g = p.grad - assert g is not None, ( - f"Expert param {n} must have grad set before expansion") - - tp_mesh = None - tp_placements_2d = None - - if is_dtensor: - local_data = p.to_local() - local_grad = g.to_local() if isinstance(g, DTensor) else g - - # Find non-dim-0 shard placements (e.g. TP sharding). - # After splitting on dim 0, Shard(k) becomes Shard(k-1). - tp_dim_indices = [] - tp_placements_2d = [] - for i, pl in enumerate(p.placements): - if _is_shard(pl) and pl.dim != 0: - tp_dim_indices.append(i) - tp_placements_2d.append(Shard(pl.dim - 1)) - - if tp_dim_indices: - tp_dim_names = tuple(p.device_mesh.mesh_dim_names[i] - for i in tp_dim_indices) - if len(tp_dim_names) == 1: - tp_mesh = p.device_mesh[tp_dim_names[0]] - else: - tp_mesh = p.device_mesh[tp_dim_names] - else: - local_data = p.data - local_grad = g - - # Expand: split dim 0, reshape each slice to 2D. - num_local_experts = local_data.shape[0] - for i in range(num_local_experts): - slice_data = local_data[i] - slice_grad = local_grad[i] - - if tp_mesh is not None: - # Wrap as DTensor on TP submesh so the pipeline handles - # TP communication (gather/scatter across TP ranks). - dt_data = DTensor.from_local(slice_data, - device_mesh=tp_mesh, - placements=tp_placements_2d) - dt_grad = DTensor.from_local(slice_grad, - device_mesh=tp_mesh, - placements=tp_placements_2d) - expert_param = torch.nn.Parameter(dt_data, requires_grad=False) - expert_param.grad = dt_grad - else: - expert_param = torch.nn.Parameter(slice_data, - requires_grad=False) - expert_param.grad = slice_grad - - expanded_names.append(f"{n}[{i}]") - expanded_params.append(expert_param) - - p.grad = None # allow expert grad storage to be freed after pipeline - - return expanded_names, expanded_params - - -class Muon(torch.optim.Optimizer): - """ - Muon - MomentUm Orthogonalized by Newton-schulz - - Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- - processing step, in which each 2D parameter's update is replaced with the nearest orthogonal - matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has - the advantage that it can be stably run in bfloat16 on the GPU. - - Some warnings: - - We believe this optimizer is unlikely to work well for training with small batch size. - - We believe it may not work well for finetuning pretrained models, but we haven't tested this. - - Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. - lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) - momentum: The momentum used by the internal SGD. (0.95 is a good default) - nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) - ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. - Parameters that are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW instead. - adamw_lr: The learning rate for the internal AdamW. - adamw_betas: The betas for the internal AdamW. - adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - warmup_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher warmup_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. - chunk_size : Batch size of parameters to process in each - all2all gather/compute/scatter step. - Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified. - use_distributed_muon: Use distributed muon by Liu et al. (2024). - For testing purpose only. - expert_keys: List of strings to identify expert-parallel parameters. - If any key appears in a parameter's name, its outermost - dimension is treated as the expert dimension and expanded - into per-expert 2D params for Muon. For example, - ``expert_keys=["experts"]`` matches any param whose name - contains "experts". 3D+ params not matched by any key - will raise an error. - """ - - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config=None, - warmup_step=5, - chunk_size=-1, - use_distributed_muon=False, - expert_keys=None): - defaults = dict( - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - nesterov=nesterov, - ns_steps=ns_steps, - adamw_betas=adamw_betas, - adamw_eps=adamw_eps, - none_grad=none_grad, - use_muon=True, - ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) - super().__init__(params, defaults) - - self.debug = debug - self.clip_config = clip_config if clip_config is not None else { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100, - } - self.warmup_step = warmup_step - self.chunk_size = chunk_size - self.use_distributed_muon = use_distributed_muon - self.expert_keys = expert_keys - self.cpu_offload = False - self._cpu_offload_pool: CPUOffloadPool | None = None - self._offload_initialized = False - self._parallel_cache: dict[tuple[str, ...], dict] = {} - self._expert_expand_cache: dict[tuple[int, ...], dict] = {} - - def _calc_flops(self, G, steps): - assert len(G.shape) == 2 - M, N = G.shape - if M > N: - M, N = N, M - - return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) - - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - p.placements, p.device_mesh) - - return shard_mesh, shard_pg, shard_placements - - def init_state_and_assign_params(self, names, params, group, qk_logits): - param_to_state = {} - param_to_flops = {} - - total_flops = 0 - for p in params: - g = p.grad - if g is None: - continue - assert g.ndim == 2, "Muon only supports 2D parameters." - - flops = self._calc_flops(g, group["ns_steps"]) - param_to_flops[id(p)] = flops - total_flops += flops - - if self.debug: - logger.debug("Total TFLOPs for Muon: %.2f TFLOPs", - total_flops / 1e12) - - paired = list(zip(names, params)) - - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) - - round_robin = 0 - mesh = ordered_params[0].device_mesh - placements = ordered_params[0].placements - - shard_mesh, shard_pg, shard_placements = self.get_shard_mesh( - ordered_params[0]) - shard_mesh_flattened = shard_mesh.mesh.flatten() - num_ranks = dist.get_world_size(group=shard_pg) - - for n, p in zip(ordered_names, ordered_params): - if mesh != p.device_mesh: - raise ValueError("All parameters must be on the same mesh.") - if placements != p.placements: - raise ValueError("All parameters must have same placements.") - - worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks - round_robin = (round_robin + 1) % len(shard_mesh_flattened) - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - # Precompute per-rank indices and numels for all-to-all. - rank_indices: dict[int, tuple] = {} - rank_numels: dict[int, int] = {} - for r in range(num_ranks): - indices = get_slices_of_dtensor(p, r, shard_mesh, - shard_placements) - rank_indices[r] = indices - numel = 1 - for idx, dim_size in zip(indices, p.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) - rank_numels[r] = numel - - param_to_state[id(p)] = _muon_state( - worker_rank=worker_rank, - process_group=shard_pg, - rank_indices=rank_indices, - rank_numels=rank_numels, - name=n, - qk_clip_state=qk_clip_state, - ) - - return param_to_state, ordered_params - - def base(self, names, params, group, lr, weight_decay, qk_logits): - # Momentum is already applied by _step_muon before this method. - for n, p in zip(names, params): - g = p.grad - if g is None: - continue - - u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) - - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - def distributed_muon( - self, - names: list[str], - params: list[torch.nn.Parameter], - group: dict[str, Any], - lr: float, - weight_decay: float, - qk_logits: list[torch.Tensor | DTensor] | None, - ): - """Batched Distributed Muon — for testing/correctness verification only. - - Uses all-gather to reconstruct full tensors, computes Newton-Schulz on - the full grad, then slices back to local shards. This is simpler but - slower than the parallel pipeline (all2all) path, so it serves as a - reference implementation for verifying correctness. - """ - with record_function("distributed_muon"): - # Momentum is already applied by _step_muon before this method. - ns_steps = group["ns_steps"] - - # Separate plain tensors (no communication) from DTensors. - plain_names, plain_params = [], [] - dtensor_names, dtensor_params = [], [] - for n, p in zip(names, params): - if p.grad is None: - continue - if isinstance(p.data, DTensor): - dtensor_names.append(n) - dtensor_params.append(p) - else: - plain_names.append(n) - plain_params.append(p) - - # Process plain tensors per-param (no communication). - for n, p in zip(plain_names, plain_params): - u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE), - steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, - qk_logits) - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - if not dtensor_params: - return - - # Group DTensors by (placements, mesh) for batched all-gather. - placement_groups: dict[tuple, - tuple[list, - list]] = defaultdict(lambda: ([], [])) - for n, p in zip(dtensor_names, dtensor_params): - key = (p.placements, p.device_mesh) - placement_groups[key][0].append(n) - placement_groups[key][1].append(p) - - logger.info( - "distributed_muon: %d placement groups, %d total dtensors", - len(placement_groups), len(dtensor_params)) - - for (placements, mesh), (grp_names, - grp_params) in placement_groups.items(): - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - placements, mesh) - rank = dist.get_rank(shard_pg) - world_size = dist.get_world_size(shard_pg) - - logger.info(" group: %d params, placements=%s, world_size=%d", - len(grp_params), placements, world_size) - - # Separate params that can be batched (all shard dims evenly - # divisible) from those needing per-param full_tensor - # (e.g. MoE gate weights with fewer rows than shard ranks). - # all_gather_into_tensor requires equal buffer sizes across - # ranks, so uneven splits must use DTensor full_tensor(). - batch_names, batch_params = [], [] - single_names, single_params = [], [] - for n, p in zip(grp_names, grp_params): - even = all(p.shape[pl.dim] % - shard_mesh.mesh.shape[dim_idx] == 0 - for dim_idx, pl in enumerate(shard_placements)) - if even: - batch_names.append(n) - batch_params.append(p) - else: - single_names.append(n) - single_params.append(p) - - # Process uneven-split params per-param via full_tensor(). - for n, p in zip(single_names, single_params): - with record_function("distributed_muon::newton_schulz"): - g_full = p.grad.full_tensor().to(COMM_DTYPE) - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - if not batch_params: - continue - - logger.info(" batched=%d, single=%d", len(batch_params), - len(single_params)) - - # Concat all local grad shards into a single flat buffer. - with record_function("distributed_muon::gather"): - grad_locals = [ - p.grad.to_local().to(COMM_DTYPE).flatten() - for p in batch_params - ] - numels = [g.numel() for g in grad_locals] - grad_concat = torch.cat(grad_locals) - del grad_locals - - # Single all-gather (replaces N separate full_tensor). - grad_gathered = torch.empty( - grad_concat.numel() * world_size, - dtype=COMM_DTYPE, - device="cuda", - ) - dist.all_gather_into_tensor(grad_gathered, - grad_concat, - group=shard_pg) - - total_numel = grad_concat.numel() - del grad_concat - - # Precompute per-param offsets within the concat buffer. - offsets = [] - off = 0 - for ne in numels: - offsets.append(off) - off += ne - - # Per-param: reconstruct full grad → NS → local update. - for i, (n, p) in enumerate(zip(batch_names, batch_params)): - with record_function("distributed_muon::newton_schulz"): - g_full = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - for r in range(world_size): - r_start = r * total_numel + offsets[i] - shard = grad_gathered[r_start:r_start + numels[i]] - indices = get_slices_of_dtensor( - p, r, shard_mesh, shard_placements) - g_full[indices] = shard.reshape( - g_full[indices].shape) - - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - def _setup_parallel(self, names, params, group, qk_logits): - """Compute (or retrieve cached) parallel pipeline metadata. - - Returns: - (ordered_params, param_to_state, rank, chunk_size) - """ - cache_key = tuple(names) - - if cache_key not in self._parallel_cache: - # First call: compute metadata and populate cache. - param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - shard_pg = param_to_state[id(ordered_params[0])].process_group - rank = dist.get_rank(group=shard_pg) - - if self.chunk_size == -1: - shard_ranks = dist.get_world_size(shard_pg) - chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO - elif self.chunk_size > 0: - chunk_size = self.chunk_size - else: - raise ValueError( - "chunk_size must be -1 or a positive integer.") - - ordered_names = [ - param_to_state[id(p)].name for p in ordered_params - ] - name_to_state = { - param_to_state[id(p)].name: param_to_state[id(p)] - for p in ordered_params - } - self._parallel_cache[cache_key] = { - 'ordered_names': ordered_names, - 'name_to_state': name_to_state, - 'rank': rank, - 'chunk_size': chunk_size, - } - else: - # Cached path: rebuild param_to_state with current id(p) keys. - cache = self._parallel_cache[cache_key] - rank = cache['rank'] - chunk_size = cache['chunk_size'] - - name_to_param = dict(zip(names, params)) - ordered_params = [name_to_param[n] for n in cache['ordered_names']] - - param_to_state = {} - for p, n in zip(ordered_params, cache['ordered_names']): - cached_state = cache['name_to_state'][n] - param_to_state[id(p)] = _muon_state( - worker_rank=cached_state.worker_rank, - process_group=cached_state.process_group, - rank_indices=cached_state.rank_indices, - rank_numels=cached_state.rank_numels, - name=n, - qk_clip_state=get_qk_clip_info(self.clip_config, n, - qk_logits), - ) - - return ordered_params, param_to_state, rank, chunk_size - - def parallel(self, - names, - params, - group, - lr, - weight_decay, - qk_logits, - prelaunch_gather=None): - """ - Perform a parallel optimization step using Muon. - - Parameters are chunked and each chunk is processed by a - :func:`muon_chunk_pipeline` generator. :func:`run_pipeline` - interleaves multiple chunks so that communication and computation - overlap across chunks (the same overlap previously achieved by the - warmup + main-loop index scheduling). - - If ``prelaunch_gather`` is provided, it is passed to the first - chunk's generator to skip re-launching the already in-flight - A2A gather. - """ - - # Momentum is already applied by _step_muon before this method. - - ordered_params, param_to_state, rank, chunk_size = ( - self._setup_parallel(names, params, group, qk_logits)) - - def pipelines(): - first = True - for start in range(0, len(ordered_params), chunk_size): - chunk = ordered_params[start:start + chunk_size] - if chunk: - kwargs = dict( - params=chunk, - param_to_state=param_to_state, - rank=rank, - ns_steps=group["ns_steps"], - lr=lr, - weight_decay=weight_decay, - none_grad=group["none_grad"], - ) - if first and prelaunch_gather is not None: - kwargs['prelaunch_gather'] = prelaunch_gather - first = False - yield muon_chunk_pipeline(**kwargs) - - with record_function("muon::pipeline"): - run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1) - - def _step_muon(self, group, qk_logits=None): - params = group["params"] - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - # Apply momentum to all params before routing/expansion. - # Batched using _foreach_* ops (compiled, fullgraph=True). - with record_function("muon::momentum"): - active_params = [p for p in params if p.grad is not None] - if active_params: - # Ensure momentum buffers exist (avoid zeros_like when already present). - for p in active_params: - if "momentum_buffer" not in self.state[p]: - self.state[p]["momentum_buffer"] = torch.zeros_like( - p.grad) - - # Extract local tensors for compiled batch function. - local_grads = [ - p.grad._local_tensor - if isinstance(p.grad, DTensor) else p.grad - for p in active_params - ] - local_bufs = [ - self.state[p]["momentum_buffer"]._local_tensor - if isinstance(self.state[p]["momentum_buffer"], DTensor) - else self.state[p]["momentum_buffer"] - for p in active_params - ] - - # Wrap momentum as tensor for torch.compile. - batch_pre_ortho(local_grads, local_bufs, - torch.tensor(momentum), group["nesterov"]) - - # For non-nesterov, the result is the momentum buffer. - if not group["nesterov"]: - for p in active_params: - p.grad = self.state[p]["momentum_buffer"] - - # Identify batched experts for deferred NS. - # Detection is cheap (condition checks only); actual NS compute is - # deferred so it can overlap with the first chunk's A2A gather. - deferred_expert_work = [] - if self.expert_keys: - batched_expert_indices = [] - for i, (n, p) in enumerate(zip(names, params)): - if not (is_expert_param(n, self.expert_keys) - and p.grad is not None): - continue - # Eligible: plain tensor, or DTensor with no non-dim-0 shards. - if isinstance(p.data, DTensor): - has_tp = any( - _is_shard(pl) and pl.dim != 0 for pl in p.placements) - if has_tp: - continue - batched_expert_indices.append(i) - - if batched_expert_indices: - # Save refs for deferred NS; free grads from param list. - for i in batched_expert_indices: - p = params[i] - g = p.grad - local_g = (g._local_tensor - if isinstance(g, DTensor) else g) - local_data = (p.data._local_tensor if isinstance( - p.data, DTensor) else p.data) - deferred_expert_work.append((local_data, local_g)) - p.grad = None - - # Remove batched experts from lists before expansion. - keep = sorted( - set(range(len(params))) - set(batched_expert_indices)) - names = [names[i] for i in keep] - params = [params[i] for i in keep] - - def _run_deferred_expert_ns(): - """Execute deferred batched expert NS.""" - if not deferred_expert_work: - return - with record_function("muon::batched_expert_ns"): - ns_steps = group["ns_steps"] - for local_data, local_g in deferred_expert_work: - u = zeropower_via_newtonschulz5_batched( - local_g.to(COMM_DTYPE), steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:]) - local_data.mul_(1 - lr * weight_decay) - local_data.add_(u, alpha=-adjusted_lr) - - # Expand expert params by splitting on dim 0. - logger.debug("[_step_muon] before expand: %d params, expert_keys=%s", - len(params), self.expert_keys) - if self.expert_keys: - cache_key = tuple(id(p) for p in params) - cache = self._expert_expand_cache.get(cache_key) - - if cache is None: - # Cold path: full expansion + build cache metadata. - exp_names, exp_params = _expand_expert_params( - names, params, self.expert_keys) - - # Build per-expert-group info for hot-path grad updates. - grad_info = [] - exp_idx = 0 - for orig_idx, (n, p) in enumerate(zip(names, params)): - if not is_expert_param(n, self.expert_keys): - exp_idx += 1 - continue - - is_dt = isinstance(p.data, DTensor) - num_experts = (p.to_local() if is_dt else p.data).shape[0] - - # Detect TP mesh from the first expanded expert param. - tp_mesh = None - tp_pls = None - sample = exp_params[exp_idx] - if isinstance(sample.data, DTensor): - tp_mesh = sample.data.device_mesh - tp_pls = list(sample.data.placements) - - grad_info.append((orig_idx, num_experts, exp_idx, is_dt, - tp_mesh, tp_pls)) - exp_idx += num_experts - - self._expert_expand_cache[cache_key] = { - 'names': exp_names, - 'params': exp_params, - 'grad_info': grad_info, - } - names, params = exp_names, exp_params - else: - # Hot path: reuse cached params, only update expert grads. - for (orig_idx, num_experts, exp_start, is_dt, tp_mesh, - tp_pls) in cache['grad_info']: - p = params[orig_idx] - g = p.grad - local_grad = (g.to_local() - if is_dt and isinstance(g, DTensor) else g) - for i in range(num_experts): - expert_p = cache['params'][exp_start + i] - sg = local_grad[i] - if tp_mesh is not None: - expert_p.grad = DTensor.from_local( - sg, device_mesh=tp_mesh, placements=tp_pls) - else: - expert_p.grad = sg - p.grad = None - - names = cache['names'] - params = cache['params'] - else: - names, params = _expand_expert_params(names, params, - self.expert_keys) - logger.debug("[_step_muon] after expand: %d params", len(params)) - - param_dtensors = [] - name_dtensors = [] - - param_tensors = [] - name_tensors = [] - - # distributed_muon is a reference implementation for testing only. - # The parallel pipeline (all2all) path below is the production path. - if self.use_distributed_muon: - _run_deferred_expert_ns() - self.distributed_muon(names=names, - params=params, - group=group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits) - return - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - logger.debug( - "[route] %s → base (DTensor all-Replicate), " - "shape=%s, placements=%s", n, p.shape, p.placements) - param_tensors.append(p) - name_tensors.append(n) - else: - logger.debug( - "[route] %s → parallel (DTensor), shape=%s, " - "placements=%s, mesh=%s", n, p.shape, p.placements, - p.device_mesh.mesh_dim_names) - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - logger.debug("[route] %s → base (plain tensor), shape=%s", n, - p.data.shape) - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError(f"Unsupported parameter type: {type(p.data)}") - - logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, " - f"{len(param_tensors)} Tensors → base") - - def group_dtensors(dtensors, names): - # To support different placements, we group parameters by placements - # and run parallel Muon on each group. - - placement_to_params = defaultdict(lambda: ([], [])) - - assert len(dtensors) == len(names) - for p, n in zip(dtensors, names): - placement_to_params[tuple([p.placements, - p.device_mesh])][0].append(n) - placement_to_params[tuple([p.placements, - p.device_mesh])][1].append(p) - return placement_to_params - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - dtensor_group = group_dtensors(param_dtensors, name_dtensors) - - # Pre-launch the first chunk's A2A gather so that the NCCL - # communication overlaps with the (deferred) batched expert NS - # compute on the default CUDA stream. - prelaunch = None - if deferred_expert_work: - first_names, first_params = next(iter(dtensor_group.values())) - ordered, pts, rnk, csz = self._setup_parallel( - first_names, first_params, group, qk_logits) - first_chunk = ordered[:csz] - if first_chunk: - prelaunch = prelaunch_first_gather(first_chunk, pts, rnk, - group["none_grad"]) - - _run_deferred_expert_ns() - - first_group = True - for _, (names, params) in dtensor_group.items(): - pg = prelaunch if first_group else None - first_group = False - self.parallel( - names, - params, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - prelaunch_gather=pg, - ) - else: - _run_deferred_expert_ns() - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - ) - - def _register_states_for_offload(self): - """Register all optimizer state tensors with the CPU offload pool. - - Called once after the first step when states have been lazily created. - Offloads all param states (momentum buffers for Muon, moment1/moment2 - for AdamW) to free GPU memory between steps. - """ - pool = self._cpu_offload_pool - tracked = 0 - for group in self.param_groups: - for p in group["params"]: - if p not in self.state: - continue - state = self.state[p] - if group.get("use_muon", False): - if "momentum_buffer" in state: - pool.track(state["momentum_buffer"]) - tracked += 1 - else: - if "moment1" in state: - pool.track(state["moment1"]) - if "moment2" in state: - pool.track(state["moment2"]) - tracked += 1 - logger.info("[CPUOffload] Registered %d param states for offload", - tracked) - - @torch.no_grad - def step(self, closure=None, qk_logits=None): - """Perform a single optimization step. - - Args: - closure (Callable, optional): A closure that reevaluates the model - and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). - """ - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # H2D: reload optimizer states from CPU before computation. - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload() - - logger.debug("[Muon.step] expert_keys=%s, %d param groups", - self.expert_keys, len(self.param_groups)) - - for i, group in enumerate(self.param_groups): - if group["use_muon"]: - logger.debug("[Muon.step] group %d: use_muon=True, %d params", - i, len(group["params"])) - self._step_muon(group, qk_logits=qk_logits) - else: - logger.debug( - "[Muon.step] group %d: use_muon=False (AdamW), %d params", - i, len(group["params"])) - step_adamw(self.state, group) - - # D2H: offload optimizer states to CPU after computation. - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - return loss - - # ------------------------------------------------------------------ - # CPU offload public helpers - # ------------------------------------------------------------------ - - def turn_on_cpu_offload(self): - """Enable CPU offload for optimizer states.""" - if self.cpu_offload: - return - logger.info("[Muon] turn_on_cpu_offload") - self.cpu_offload = True - if not self.state: - return - self._cpu_offload_pool = CPUOffloadPool() - self._offload_initialized = False - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_off_cpu_offload(self): - """Disable CPU offload and keep optimizer states resident on GPU.""" - if not self.cpu_offload: - return - logger.info("[Muon] turn_off_cpu_offload") - if self._offload_initialized: - self._cpu_offload_pool.reload() - torch.cuda.current_stream().synchronize() - self._cpu_offload_pool = None - self._offload_initialized = False - self.cpu_offload = False - - # ------------------------------------------------------------------ - # Checkpoint support for cpu_offload - # ------------------------------------------------------------------ - - def state_dict(self) -> dict: - if self.cpu_offload: - raise RuntimeError( - "Muon.state_dict() requires turn_off_cpu_offload() before checkpoint save." - ) - return super().state_dict() - - def load_state_dict(self, state_dict: dict) -> None: - if self.cpu_offload: - raise RuntimeError( - "Muon.load_state_dict() requires turn_off_cpu_offload() before checkpoint load." - ) - super().load_state_dict(state_dict) - - # Invalidate adamw.py's module-level tensor caches so that - # the next step rebuilds them with the newly loaded state tensors. - _placement_cache.clear() - _tensor_cache.clear() diff --git a/build/torch29-cxx11-cu130-x86_64-linux/newton_schulz.py b/build/torch29-cxx11-cu130-x86_64-linux/newton_schulz.py deleted file mode 100644 index d939264b69a34e7a3fa78859f34dc265a1159d59..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/newton_schulz.py +++ /dev/null @@ -1,240 +0,0 @@ -from itertools import repeat -from math import inf, sqrt - -import numpy as np -import torch - -from .matmul_transpose_triton import matmul_transpose_assign - -COMM_DTYPE = torch.bfloat16 -DEFAULT_CHUNK_SIZE_RATIO = 4 - - -def _optimal_quintic(l, u, max_iter=1000): - """ - Use the simplified Remez algorithm to find the optimal odd quintic approximant - to the constant function x -> 1 over the interval [l, u]. - - Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum - approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the - two interior equioscillation nodes q, r until convergence. Returns the - closed-form equioscillating solution when l ≈ u. - - Raises ValueError if any intermediate value (a, b, c, E, q, r) is non-finite - (NaN or inf). Raises RuntimeError if convergence is not reached within - max_iter iterations. - """ - assert 0 <= l <= u - if 1 - 5e-6 <= l / u: - return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5) - q = (3 * l + u) / 4 - r = (l + 3 * u) / 4 - E = inf - for _ in range(max_iter): - old_E = E - LHS = np.array( - [ - [l, l**3, l**5, 1], - [q, q**3, q**5, -1], - [r, r**3, r**5, 1], - [u, u**3, u**5, -1], - ] - ) - a, b, c, E = np.linalg.solve(LHS, np.ones(4)) - if not np.all(np.isfinite([a, b, c, E])): - raise ValueError( - f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}" - ) - q, r = np.sqrt( - (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / (10 * c) - ) - if not np.all(np.isfinite([q, r])): - raise ValueError(f"_optimal_quintic: non-finite node update q={q}, r={r}") - if abs(old_E - E) <= 1e-15: - break - else: - raise RuntimeError( - f"_optimal_quintic: did not converge after {max_iter} iterations" - ) - return float(a), float(b), float(c) - - -def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0): - """ - Compute the Polar Express coefficient series for `num_iters` quintic iterations. - - Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that - compose to map singular values from [l, 1] toward 1. At each step: - 1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion` - prevents near-zero singular values from stalling by raising the effective - lower bound; if it is active (cushion*u > l), the coefficients are - rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u]. - 2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the - last iteration, providing numerical headroom at the cost of a slightly slower - final convergence step. - 3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1). - - Returns a list of (a, b, c) tuples, one per iteration. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - u = 1 - assert 0 <= l <= u - safety_factor = 1 + safety_factor_eps - coefficients = [] - for iter in range(num_iters): - a, b, c = _optimal_quintic(max(l, cushion * u), u) - if cushion * u > l: - pl = a * l + b * l**3 + c * l**5 - pu = a * u + b * u**3 + c * u**5 - rescaler = 2 / (pl + pu) - a *= rescaler - b *= rescaler - c *= rescaler - if iter < num_iters - 1: - a /= safety_factor - b /= safety_factor**3 - c /= safety_factor**5 - coefficients.append((a, b, c)) - l = a * l + b * l**3 + c * l**5 - u = 2 - l - return coefficients - - -# Precomputed Polar Express coefficients (a, b, c) for 10 quintic Newton-Schulz -# iterations. Each tuple is the minimax-optimal (Remez/equioscillation) odd quintic -# approximant to x->1 over the current singular-value interval, computed once at -# import time and reused across all optimizer steps. -# -# Contrast with the former hardcoded NS coefficients (5 fixed tuples): -# - Former: empirically tuned to maximize slope at zero; did not converge -# singular values to 1, yielding US'V^T with S' ~ Uniform(0.5, 1.5) instead -# of the true polar factor UV^T. -# - Polar Express: analytically optimal per step, adapting to the shrinking -# singular-value interval [l, u] as iterations progress; converges all -# singular values to 1, producing the exact polar factor UV^T. -_coeffs_list = _optimal_composition( - l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02 -) - - -# This code is adapted from: -# KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py) -# NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress) -# matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon) -@torch.no_grad() -def _zeropower_via_newtonschulz5(G, steps): - """ - Compute the polar factor of G via the Polar Express method. - - Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c) - are the Polar Express coefficients from `_coeffs_list`. Each step is the - optimal odd quintic approximant to x -> 1 over the current singular-value - interval, minimizing the maximum approximation error (Remez / minimax criterion). - The composition maps singular values from [l, 1] to near 1, producing the - polar factor (orthogonal factor in the polar decomposition G = UP). - - `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2, - cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE - X = G # no manual typecast - - if G.size(0) > G.size(1): - X = X.T - - X = X / (X.norm() + 1e-7) - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - # Perform the NS iterations - for a, b, c in hs: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(0) > G.size(1): - X = X.T - - return X - - -@torch.no_grad() -def _zeropower_via_newtonschulz5_batched(G, steps): - """Batched polar factor computation for 3D (E, out, in) tensors. - - Same algorithm as ``_zeropower_via_newtonschulz5`` but uses - ``torch.bmm`` / ``torch.baddbmm`` instead of the 2D Triton kernel, - processing all E expert matrices in a single batched call. - """ - assert len(G.shape) == 3 - assert G.dtype == COMM_DTYPE - X = G - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - # Per-expert Frobenius norm. - X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) - - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - for a, b, c in hs: - buf1 = torch.bmm(X, X.transpose(-2, -1)) - buf2 = torch.bmm(buf1, buf1.transpose(-2, -1)) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.baddbmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - return X - - -_ns_per_shape: dict[tuple[int, ...], callable] = {} -_use_compile = True - - -def set_ns_compile(enabled: bool): - """Toggle torch.compile for Newton-Schulz iteration.""" - global _use_compile - _use_compile = enabled - - -def zeropower_via_newtonschulz5(G, steps=5): - if not _use_compile: - return _zeropower_via_newtonschulz5(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() - - -def zeropower_via_newtonschulz5_batched(G, steps=5): - """Compile-cached batched Newton-Schulz for 3D expert tensors.""" - if not _use_compile: - return _zeropower_via_newtonschulz5_batched(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile( - _zeropower_via_newtonschulz5_batched, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() diff --git a/build/torch29-cxx11-cu130-x86_64-linux/optimizer/__init__.py b/build/torch29-cxx11-cu130-x86_64-linux/optimizer/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/optimizer/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu130-x86_64-linux/pipeline.py b/build/torch29-cxx11-cu130-x86_64-linux/pipeline.py deleted file mode 100644 index c0c2d515856182d8d15ad27dd4e4e093b29397d6..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/pipeline.py +++ /dev/null @@ -1,468 +0,0 @@ -import logging -from typing import Generator - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -from .core import _muon_state, adjust_lr_for_muon -from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5 -from .qk_clip import compute_scales - -logger = logging.getLogger(__name__) - -# ====================================================================== -# Stage helpers -# ====================================================================== - - -def _launch_gather( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Allocate gather buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_gather``). - gathered_grads: ``{id(p): empty_tensor}`` for owned params, - ``None`` for non-owned. - recv_counts: Per-source-rank element counts. - """ - # Allocate gathered-grad buffers - gathered_grads: dict[int, torch.Tensor | None] = {} - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - gathered_grads[id(p)] = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - else: - gathered_grads[id(p)] = None - - # Build send buffer – batch grad copies via torch.cat - # (1-2 fused kernels vs N individual narrow().copy_() calls). - send_counts = [0] * num_ranks - for p in params: - state = param_to_state[id(p)] - send_counts[state.worker_rank] += state.rank_numels[rank] - - total_send = sum(send_counts) - if total_send > 0: - # Group grad slices by destination rank in a single pass. - dst_to_grads = [[] for _ in range(num_ranks)] - for p in params: - state = param_to_state[id(p)] - n = state.rank_numels[rank] - if n > 0: - g = p.grad.to_local() - dst_to_grads[state.worker_rank].append(g.reshape(-1)) - - # Flatten in dst order and cat once. - all_slices = [] - for dst in range(num_ranks): - all_slices.extend(dst_to_grads[dst]) - send_buf = torch.cat(all_slices) - if send_buf.dtype != COMM_DTYPE: - send_buf = send_buf.to(COMM_DTYPE) - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += state.rank_numels[src] - recv_counts[src] = total - - recv_buf = torch.empty(sum(recv_counts), dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - logger.debug(f"send_buf size: {send_buf.numel()}, " - f"recv_buf size: {recv_buf.numel()}, " - f"recv_counts: {recv_counts}, " - f"send_counts: {send_counts}, " - f"process_group: {str(process_group)}") - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, gathered_grads, recv_counts - - -def _complete_gather( - recv_buf: torch.Tensor, - recv_counts: list[int], - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - param_to_state: dict[int, _muon_state], - rank: int, -) -> None: - """Reconstruct gathered grads from the recv buffer (in-place).""" - off = 0 - for src in range(len(recv_counts)): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - - indices = state.rank_indices[src] - - shard_view = gathered_grads[id(p)][indices] - n = shard_view.numel() - if n == 0: - continue - - sg = recv_buf.narrow(0, off + inner_off, n) - sg = sg.reshape(shard_view.shape) - gathered_grads[id(p)][indices] = sg - - inner_off += n - assert inner_off == block - off += block - - -def _compute_ns( - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - ns_steps: int, -) -> dict[int, torch.Tensor | None]: - """Run Newton-Schulz orthogonalization on owned parameters. - - Returns: - computed_us: ``{id(p): orthogonalized_update}`` for owned params. - """ - computed_us: dict[int, torch.Tensor | None] = {} - for p in owned_params: - u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps) - gathered_grads[id(p)] = None # free gathered grad - computed_us[id(p)] = u - return computed_us - - -def _launch_scatter( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, - computed_us: dict[int, torch.Tensor | None], -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor], list[int]]: - """Allocate scatter buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_scatter``). - scattered_us: Empty dict, populated by ``_complete_scatter`` with - zero-copy views into ``recv_buf``. - recv_counts: Per-source-rank element counts. - """ - # scattered_us is populated by _complete_scatter with zero-copy views - # into recv_buf, avoiding N empty_like allocations + N copy_ calls. - # Pre-seed entries for params whose local shard is empty (rank_numels == 0) - # so _update_params can iterate all params without KeyError. - scattered_us: dict[int, torch.Tensor] = {} - for p in params: - if param_to_state[id(p)].rank_numels[rank] == 0: - scattered_us[id(p)] = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - # Build send buffer – batch via torch.cat - # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls). - send_counts = [0] * num_ranks - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - for dst_rank in range(num_ranks): - send_counts[dst_rank] += state.rank_numels[dst_rank] - - total_send = sum(send_counts) - if total_send > 0: - # Cache u_full conversions to avoid redundant .to() per dst_rank. - u_fulls = {} - for p in owned_params: - u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous() - - # Collect slices in dst order (matches all-to-all send layout). - all_slices = [] - for dst_rank in range(num_ranks): - for p in owned_params: - state = param_to_state[id(p)] - su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten() - if su.numel() > 0: - all_slices.append(su) - - send_buf = torch.cat(all_slices) if all_slices else torch.empty( - 0, dtype=COMM_DTYPE, device="cuda") - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += state.rank_numels[rank] - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, scattered_us, recv_counts - - -def _complete_scatter( - recv_buf: torch.Tensor, - recv_counts: list[int], - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], -) -> None: - """Populate scattered_us with zero-copy views into recv_buf. - - Instead of pre-allocating tensors and copying, we assign views directly - from ``recv_buf``. This eliminates N ``empty_like`` + N ``copy_`` calls. - The underlying storage of ``recv_buf`` is kept alive through the views - until ``scattered_us`` is cleared after ``_update_params``. - """ - off = 0 - for src in range(len(recv_counts)): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = state.rank_numels[rank] - if n == 0: - continue - - scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - - inner_off += n - - assert inner_off == block - off += block - - -def _update_params( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], - lr: float, - weight_decay: float, -) -> None: - """Apply weight decay, Muon update, and optional QK clipping. - - Uses batched ``_foreach_mul_`` for weight decay and batched - ``_foreach_add_`` for the Muon update, grouping parameters by - adjusted_lr to minimize kernel launches while preserving float32 - precision for the alpha scaling. - """ - if not params: - return - - # Batched weight decay: p *= (1 - lr * wd) — single fused kernel. - p_locals = [p._local_tensor for p in params] - torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay) - - # Group params by adjusted_lr so _foreach_add_ can use a single - # alpha per group (preserves float32 precision for alpha scaling). - lr_groups: dict[float, tuple[list, list]] = {} - for p in params: - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - if adjusted_lr not in lr_groups: - lr_groups[adjusted_lr] = ([], []) - lr_groups[adjusted_lr][0].append(p._local_tensor) - lr_groups[adjusted_lr][1].append(scattered_us[id(p)]) - - for adjusted_lr, (p_group, u_group) in lr_groups.items(): - torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr) - - # QK clipping – applied directly on the local tensor to - # avoid DTensor sharding-propagation issues with _StridedShard. - for p in params: - state = param_to_state[id(p)] - if state.qk_clip_state is None: - continue - scales_full = compute_scales(p, state.qk_clip_state) - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = state.rank_indices[rank][0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - -# ====================================================================== -# Pre-launch helper for overlapping first chunk's gather with other work. -# ====================================================================== - - -@torch.no_grad() -def prelaunch_first_gather( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - none_grad: bool, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Launch the first chunk's A2A gather early for overlap with other compute. - - Call this *before* expensive GPU work (e.g. batched expert NS) so that - the NCCL all-to-all runs concurrently on the NCCL stream while the - default stream executes compute. - - Returns the same 4-tuple that ``_launch_gather`` produces, which should - be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - with record_function("muon::prelaunch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - return work, recv_buf, gathered_grads, recv_counts - - -# ====================================================================== -# Main generator – thin orchestrator that wires stages together. -# ====================================================================== - - -@torch.no_grad() -def muon_chunk_pipeline( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - ns_steps: int, - lr: float, - weight_decay: float, - none_grad: bool, - prelaunch_gather: tuple | None = None, -) -> Generator[None, None, None]: - """Process one chunk of parameters through the full Muon pipeline. - - Stages: gather -> compute (Newton-Schulz) -> scatter -> update. - - Each ``yield`` lets :func:`run_pipeline` interleave other chunks so - that communication and computation overlap across chunks. Async - communication is launched via ``async_op=True`` and completed after - the yield with ``work.wait()``. - - Overlap happens because :func:`run_pipeline` admits one new chunk - per iteration (staggered admission). While chunk *N* does NS - compute on the default CUDA stream, chunk *N+1*'s async all-to-all - runs concurrently on the NCCL stream — no separate ``comm_stream`` - is required. - - If ``prelaunch_gather`` is provided, the gather was already launched - by :func:`prelaunch_first_gather` and we skip launching it again. - - Yields exactly **2** times: - - 1. After launching async all-to-all gather (or immediately if pre-launched). - 2. After launching async all-to-all scatter. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - if prelaunch_gather is not None: - # Gather was pre-launched; none_grad already handled by caller. - work, recv_buf, gathered_grads, recv_counts = prelaunch_gather - else: - # Normal path: launch async gather. - with record_function("muon::launch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - yield # --- YIELD 1: other chunks can launch their gather --- - - with record_function("muon::wait_gather"): - work.wait() - _complete_gather(recv_buf, recv_counts, owned_params, gathered_grads, - param_to_state, rank) - del recv_buf - - # Stage 3: Newton-Schulz orthogonalization. - with record_function("muon::newton_schulz"): - computed_us = _compute_ns(owned_params, gathered_grads, ns_steps) - gathered_grads.clear() - - # Stages 4-5: launch async scatter. - with record_function("muon::launch_scatter"): - work, recv_buf, scattered_us, recv_counts = _launch_scatter( - params, owned_params, param_to_state, rank, num_ranks, - process_group, computed_us) - computed_us.clear() - - yield # --- YIELD 2: other chunks can launch their scatter --- - - with record_function("muon::wait_scatter"): - work.wait() - _complete_scatter(recv_buf, recv_counts, params, param_to_state, rank, - scattered_us) - del recv_buf - - # Stage 6: apply parameter updates. - with record_function("muon::update_params"): - _update_params(params, param_to_state, rank, scattered_us, lr, - weight_decay) - scattered_us.clear() diff --git a/build/torch29-cxx11-cu130-x86_64-linux/qk_clip.py b/build/torch29-cxx11-cu130-x86_64-linux/qk_clip.py deleted file mode 100644 index 2aba711b3004b7f09e7141da7ef834bd61cc2430..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/qk_clip.py +++ /dev/null @@ -1,198 +0,0 @@ -import logging -import math -from dataclasses import dataclass - -import torch -from torch.distributed.tensor import DTensor - -from .core import normalize_fqn - -logger = logging.getLogger(__name__) - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - and return (kind, layer_index). - - Supported kinds: - MHA/GQA: 'wq', 'wk', 'q_proj', 'k_proj' - MLA: 'wq_b' (Q up-proj), 'wkv_b' (KV up-proj) - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.1.attn.wq_b.weight' -> ('wq_b', 1) - 'model.0.attn.wkv_b.weight' -> ('wkv_b', 0) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = normalize_fqn(name).split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: str | None # 'wq'/'q_proj'/'wq_b' or 'wk'/'k_proj'/'wkv_b' or None - indices: list[int] # which heads to consider for clipping - head_dim: int # from config (qk_head_dim for MLA wq_b) - threshold: float # from config - logit: torch.Tensor | None - - # MLA-specific fields - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 - - -def get_qk_clip_info(clip_config, n, qk_logits): - """Extract QK clipping info for a named parameter. - - Args: - clip_config: QK clipping configuration dict (or None). - MHA/GQA keys: head_dim, threshold, q_indices, k_indices - MLA extra keys: is_mla=True, qk_nope_head_dim, qk_rope_head_dim, v_head_dim - n: Parameter name string. - qk_logits: Dict mapping layer indices to logit tensors (or None). - - Returns: - QKClipInfo instance with clipping configuration for this parameter. - """ - if clip_config is None: - return None - - head_dim = clip_config.get('head_dim') - threshold = clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - # In TP settings, qk_logits may be DTensor - # We convert it to full tensor here for simplicity - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - -def compute_scales(p, qk_clip_state): - """Compute per-head scaling factors for QK clipping. - - Returns scales tensor (√γ per head) if any head exceeds threshold, else None. - For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim. - """ - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - # Check if any head exceeds threshold before allocating. - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - - if not head_scales: - return None - - # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows - if qk_clip_state.is_mla and kind == 'wkv_b': - effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim - else: - effective_head_dim = head_dim - - H_global = p.shape[0] // effective_head_dim - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale - return scales_full - - -def qk_clip(p, scales, info): - """Apply per-head scaling to a Q/K projection weight matrix. - - Args: - p: Parameter (nn.Parameter or raw tensor). - scales: [n_heads] tensor, each element = √γ_h. - info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions. - - MLA sub-region scaling per Algorithm 1 (MuonClip): - wq_b: q_nope rows → √γ, q_pe rows → γ - wkv_b: k_nope rows → √γ, v rows → unchanged - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not info.is_mla: - # MHA/GQA: uniform √γ applied to all rows in each head - W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: vectorized sub-region scaling within each head - if info.kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_head_dim = qk_nope + info.qk_rope_head_dim - W_3d = W.view(-1, qk_head_dim, W.shape[1]) # [H, qk_head_dim, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # q_nope → √γ - W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1, - 1)) # q_pe → γ - - elif info.kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - W_3d = W.view(-1, kv_stride, W.shape[1]) # [H, kv_stride, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # k_nope → √γ - # v rows: not touched (k_R shared rotary unchanged) diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/__init__.py b/build/torch29-cxx11-rocm63-x86_64-linux/__init__.py deleted file mode 100644 index 239c7a65f8293e7d0df28f05fce645af56d628c0..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .muon import Muon - -__all__ = [ - "Muon", -] diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/_ops.py b/build/torch29-cxx11-rocm63-x86_64-linux/_ops.py deleted file mode 100644 index 034bff088659b5df6f6d401feb18c89dc5f33b29..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _optimizer_8d53b78_dirty -ops = torch.ops._optimizer_8d53b78_dirty - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_optimizer_8d53b78_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so b/build/torch29-cxx11-rocm63-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so deleted file mode 100755 index 1c0c2fe89693c7112bf2ec2e0ea203ba1a8292bb..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:9bdbe93877a276dbfdd4c596b788d70ed61804c4c5bc555259c6a3be0e9ec8fe -size 1865112 diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/adamw.py b/build/torch29-cxx11-rocm63-x86_64-linux/adamw.py deleted file mode 100644 index b5a95816a9f5b9e1889eaadae65373bfbced809a..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/adamw.py +++ /dev/null @@ -1,271 +0,0 @@ -import logging -from collections import defaultdict -from typing import cast - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -def fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: float | torch.Tensor, - weight_decay: float, - eps: float, - maximize: bool, -) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: dict | None = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) - - -def _to_local(t): - """Unwrap DTensor to local tensor for fused ops.""" - return t._local_tensor if isinstance(t, DTensor) else t - - -# --------------------------------------------------------------------------- -# Caches for eliminating per-step Python overhead. -# -# Placement grouping and tensor list assembly are identical every step -# (params don't change placement, moment/step tensors are the same objects -# after initialisation). We cache them keyed by id() of the param list -# stored in param_groups (stable across steps). -# -# Only gradients change each step and must be collected fresh. -# --------------------------------------------------------------------------- - -# id(group["params"]) → dict[placement_key, list[param]] -_placement_cache: dict[int, dict[tuple, list]] = {} - -# id(placement_group_list) → (params_local, moment1, moment2, state_steps) -_tensor_cache: dict[int, tuple[list, list, list, list]] = {} - - -def _step_adamw_params_slow(optimizer_state, params, group): - """Uncached fallback for the rare case where some params lack grads.""" - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - g = p.grad - if g is None: - continue - state = optimizer_state[p] - params_with_grads.append(_to_local(p)) - grads.append(_to_local(g)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - if not params_with_grads: - return - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw_params(optimizer_state, params, group): - """Run fused AdamW on a list of parameters sharing the same placement. - - After the first call, cached tensor lists (params_local, moment1, - moment2, state_steps) are reused — only gradients are collected fresh. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - params: List of parameters to update. - group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay. - """ - # Collect grads — the only thing that changes each step. - with record_function("adamw::collect_grads"): - grads = [] - for p in params: - g = p.grad - if g is None: - # Rare: fall back to slow path that filters per-param. - _step_adamw_params_slow(optimizer_state, params, group) - return - grads.append(_to_local(g)) - - tensor_key = id(params) - if tensor_key not in _tensor_cache: - with record_function("adamw::init_tensor_cache"): - params_local = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - state = optimizer_state[p] - params_local.append(_to_local(p)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(p.grad) - state["moment2"] = torch.zeros_like(p.grad) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - _tensor_cache[tensor_key] = (params_local, moment1, moment2, - state_steps) - - params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key] - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - with record_function("adamw::fused_adamw"): - fused_adamw( - params_local, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw(optimizer_state, group): - """Dispatch AdamW step, grouping parameters by type and placement. - - Placement grouping is cached after the first call since params never - change their placement between steps. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - group: Parameter group dict. - """ - params = group["params"] - placement_key = id(params) - - if placement_key not in _placement_cache: - with record_function("adamw::group_by_placement"): - placement_to_params: dict[tuple, - list[torch.Tensor]] = defaultdict(list) - for p in params: - match p: - case DTensor(): - logger.debug( - "[AdamW] DTensor param: shape=%s, placements=%s, " - "mesh=%s, grad=%s", p.shape, p.placements, - p.device_mesh.mesh_dim_names, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple( - [p.placements, p.device_mesh])].append(p) - case torch.Tensor(): - logger.debug( - "[AdamW] plain param: shape=%s, grad=%s", p.shape, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple([torch.Tensor, - None])].append(p) - - logger.debug("[AdamW] %d placement groups, %d total params", - len(placement_to_params), len(params)) - - _placement_cache[placement_key] = dict(placement_to_params) - - for group_params in _placement_cache[placement_key].values(): - step_adamw_params(optimizer_state, group_params, group) diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/async_utils.py b/build/torch29-cxx11-rocm63-x86_64-linux/async_utils.py deleted file mode 100644 index a45c530ac9cad88e3555ec1047a6aa59f225347e..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/async_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import logging -from typing import Generator - -logger = logging.getLogger(__name__) - - -class _Task: - """Internal: wraps a generator, advances one yield at a time.""" - - def __init__(self, generator: Generator[None, None, None], index: int): - self._generator = generator - self._index = index - self._steps_completed = 0 - self.step() # run to first yield - - def step(self) -> bool: - try: - next(self._generator) - self._steps_completed += 1 - logger.debug("pipeline[%d] completed stage %d", self._index, - self._steps_completed) - return True - except StopIteration: - logger.debug("pipeline[%d] finished after %d stages", self._index, - self._steps_completed) - return False - - def close(self): - self._generator.close() - - -def run_pipeline( - pipelines: Generator[Generator[None, None, None], None, None], - max_concurrent: int, -) -> None: - """Run generator-based pipelines with bounded concurrency. - - Each pipeline is a generator that yields at stage boundaries. - The runtime interleaves pipelines so communication and computation - overlap across chunks. - """ - if max_concurrent <= 0: - raise ValueError(f"max_concurrent must be > 0, got {max_concurrent}") - - have_new = True - task_index = 0 - previous_tasks: list[_Task] = [] - - try: - while have_new or previous_tasks: - running_tasks: list[_Task] = [] - - # Admit one new pipeline per iteration (staggered admission). - # Admitting one at a time ensures that while chunk N does NS - # compute on the default stream, chunk N+1's NCCL all-to-all - # runs concurrently on the NCCL stream — creating real - # communication/computation overlap on the GPU. - if have_new and len(previous_tasks) < max_concurrent: - try: - gen = next(pipelines) - task = _Task(gen, task_index) - task_index += 1 - running_tasks.append(task) - except StopIteration: - have_new = False - - # Advance every previously-yielded task by one step. - for task in previous_tasks: - if task.step(): - running_tasks.append(task) - - previous_tasks = running_tasks - except BaseException: - # Clean up all in-flight generators to release GPU resources. - for task in previous_tasks: - task.close() - raise diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/core.py b/build/torch29-cxx11-rocm63-x86_64-linux/core.py deleted file mode 100644 index c69d515afef305ad0ed66374095fa2d2468d99cc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/core.py +++ /dev/null @@ -1,219 +0,0 @@ -import logging -import math -from dataclasses import dataclass -from typing import List - -import torch -from torch.distributed import ProcessGroup -from torch.distributed.tensor import DTensor - -# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into -# parameter FQNs. Activation checkpointing similarly inserts -# "_checkpoint_wrapped_module". Strip these so name-based matching (skip_keys, -# expert_keys, QK layer parsing) works regardless of wrapper nesting. -_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) - -logger = logging.getLogger(__name__) - - -def normalize_fqn(name: str) -> str: - """Strip torch.compile / checkpoint wrapper components from a parameter FQN.""" - return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS) - - -@dataclass -class _muon_state: - worker_rank: int - process_group: ProcessGroup - rank_indices: dict[int, tuple] # local_rank -> per-dim indices - rank_numels: dict[int, int] # local_rank -> numel - name: str - qk_clip_state: torch.Tensor | None = None - - -def _batch_momentum( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update (no nesterov).""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - - -def _batch_momentum_nesterov( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update with nesterov correction.""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - nesterov_terms = torch._foreach_mul(momentum_bufs, momentum) - torch._foreach_add_(grads, nesterov_terms) - - -_compiled_momentum: dict[bool, callable] = {} -_use_momentum_compile = True - - -def set_momentum_compile(enabled: bool): - """Toggle torch.compile for batched momentum.""" - global _use_momentum_compile - _use_momentum_compile = enabled - - -def batch_pre_ortho( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, - nesterov: bool, -) -> None: - """Batched momentum update on lists of plain tensors. - - Mirrors dion's ``muon_update_pre_orthogonalize``. - Inputs must be plain CUDA tensors (not DTensor). - Modifies ``momentum_bufs`` and (for nesterov) ``grads`` in-place. - - When compile is enabled, uses separately compiled functions for - nesterov=True/False to avoid graph breaks from the branch. - """ - fn = _batch_momentum_nesterov if nesterov else _batch_momentum - if _use_momentum_compile: - if nesterov not in _compiled_momentum: - _compiled_momentum[nesterov] = torch.compile(fn) - fn = _compiled_momentum[nesterov] - fn(grads, momentum_bufs, momentum) - - -def _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay): - """Weight-decay + update on plain tensors. - - Not compiled: per-param @torch.compile caused ~0.25ms TorchDynamo cache - lookup per call × 256+ params = massive overhead. The pipeline path uses - batched _foreach_* ops instead; this function remains for base() and - distributed_muon(). - """ - p_data.mul_(1 - lr * weight_decay) - p_data.add_(u_data, alpha=-adjusted_lr) - - -def update_p(p, u, lr, adjusted_lr, weight_decay): - """Apply weight decay and orthogonalized update to parameter. - - Args: - p: Parameter (torch.nn.Parameter or DTensor). - u: Orthogonalized update tensor. - lr: Base learning rate. - adjusted_lr: Size-adjusted learning rate. - weight_decay: Weight decay coefficient. - """ - # Unwrap Parameter -> underlying data tensor. - p_data = p.data if isinstance(p, torch.nn.Parameter) else p - # Unwrap DTensor -> local CUDA tensor for compiled kernel. - if isinstance(p_data, DTensor): - p_data = p_data._local_tensor - u_data = u._local_tensor if isinstance(u, DTensor) else u - _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay) - - -def adjust_lr_for_muon(lr, param_shape): - """Scale learning rate based on parameter matrix dimensions. - - Args: - lr: Base learning rate. - param_shape: Shape of the parameter tensor. - - Returns: - Adjusted learning rate. - """ - A, B = param_shape[:2] - # We adjust the learning rate and weight decay based on the size of the parameter matrix - # as described in the paper - adjusted_ratio = 0.2 * math.sqrt(max(A, B)) - adjusted_lr = lr * adjusted_ratio - return adjusted_lr - - -def _match_key(parts, key): - """Check if key matches as contiguous components in parts. - - Single-component keys (e.g. "experts") match any single component. - Multi-component keys (e.g. "experts.w1") match as a contiguous subsequence. - """ - key_parts = key.split(".") - key_len = len(key_parts) - if key_len == 1: - return key in parts - return any(parts[i:i + key_len] == key_parts - for i in range(len(parts) - key_len + 1)) - - -def is_expert_param(name, expert_keys): - """Check if a parameter name matches any expert key (component-level).""" - if not expert_keys: - return False - parts = normalize_fqn(name).split(".") - return any(_match_key(parts, key) for key in expert_keys) - - -def default_is_muon(name, x, expert_keys=None): - normalized = normalize_fqn(name) - parts = normalized.split(".") - skip_keys = [ - "embed_tokens", - "lm_head", - "tok_embeddings", - "output", - "mhc_attn", - "mhc_ffn", - "lambda_proj", - ] - if any(key in parts for key in skip_keys): - logger.info( - "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d", - normalized, name, x.ndim) - return False - effective_ndim = x.ndim - is_expert = is_expert_param(name, expert_keys) - if is_expert: - effective_ndim -= 1 - result = effective_ndim >= 2 - logger.info( - "[is_muon] %s (orig: %s): ndim=%d, expert=%s, effective_ndim=%d → %s", - normalized, name, x.ndim, is_expert, effective_ndim, - "Muon" if result else "AdamW") - return result - - -def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None): - if is_muon_func is None: - is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys) - - muon_params, muon_names = [], [] - non_muon_params, non_muon_names = [], [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - non_muon_names.append(n) - - logger.info("[param_groups] expert_keys=%s, Muon=%d, AdamW=%d", - expert_keys, len(muon_names), len(non_muon_names)) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/cpu_offload.py b/build/torch29-cxx11-rocm63-x86_64-linux/cpu_offload.py deleted file mode 100644 index fb5e69154a1d4a6c884491413a37a9acf0f66c80..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/cpu_offload.py +++ /dev/null @@ -1,206 +0,0 @@ -"""CPU offloading for optimizer states. - -Manages a pinned CPU memory pool and async CUDA streams to offload -optimizer state tensors (momentum buffers, Adam moments) to CPU between -optimizer steps, freeing GPU memory. - -All tracked tensors are packed into a single flat pinned CPU buffer -(per dtype). D2H and H2D copies are performed per-tensor directly -between individual GPU tensors and their slice of the CPU flat buffer -— no GPU staging buffer is allocated, so there is **no temporary GPU -memory spike** during offload or reload. - -Individual tensor storages are freed after offload via -``untyped_storage().resize_(0)``, preserving tensor identity so -downstream caches remain valid. -""" - -import logging -from collections import defaultdict - -import torch -from torch.distributed.tensor import DTensor - -logger = logging.getLogger(__name__) - - -class CPUOffloadPool: - """Pinned CPU memory pool for async optimizer state offloading. - - Tracked tensors are grouped by dtype. Each group gets a single flat - pinned CPU buffer. D2H / H2D copies are per-tensor (into slices of - the flat buffer) to avoid allocating a GPU staging buffer. - """ - - def __init__(self): - self._managed: list[torch.Tensor] = [] - self._storage_nbytes: dict[int, int] = {} # id(t) → bytes - - # Per-dtype group: populated on first offload. - # dtype → dict with keys: - # "indices" : list[int] managed-list indices - # "offsets" : list[tuple[int,int]] (start, numel) in flat buf - # "total" : int total numel - # "cpu_flat" : Tensor pinned CPU buffer - self._groups: dict[torch.dtype, dict] = {} - - self._offload_stream: torch.cuda.Stream | None = None - self._device: torch.device | None = None - self._initialized: bool = False - self._logged: bool = False - - # ------------------------------------------------------------------ - @staticmethod - def _local(t: torch.Tensor) -> torch.Tensor: - """Unwrap DTensor to its local CUDA tensor.""" - return t._local_tensor if isinstance(t, DTensor) else t - - def _ensure_stream(self): - if self._offload_stream is None: - self._offload_stream = torch.cuda.Stream(device=self._device) - - # ------------------------------------------------------------------ - def track(self, tensor: torch.Tensor): - """Register a GPU tensor for CPU offloading. Idempotent.""" - tid = id(tensor) - if tid in self._storage_nbytes: - return - local = self._local(tensor) - if self._device is None: - self._device = local.device - storage = local.untyped_storage() - # Skip tensors with empty storage (e.g. empty FSDP shards) - if storage.size() == 0: - return - self._storage_nbytes[tid] = storage.size() - self._managed.append(tensor) - - # ------------------------------------------------------------------ - def _init_buffers(self): - """Build per-dtype flat buffers on first offload.""" - # Group managed tensors by dtype. - dtype_map: dict[torch.dtype, list[tuple[int, int]]] = defaultdict(list) - for idx, t in enumerate(self._managed): - local = self._local(t) - dtype_map[local.dtype].append((idx, local.numel())) - - total_cpu_bytes = 0 - for dtype, entries in dtype_map.items(): - offsets: list[tuple[int, int]] = [] - indices: list[int] = [] - off = 0 - for idx, n in entries: - indices.append(idx) - offsets.append((off, n)) - off += n - cpu_flat = torch.empty(off, dtype=dtype, device="cpu", pin_memory=True) - self._groups[dtype] = { - "indices": indices, - "offsets": offsets, - "total": off, - "cpu_flat": cpu_flat, - } - total_cpu_bytes += off * cpu_flat.element_size() - - self._initialized = True - logger.info( - "[CPUOffload] Pool initialized: %d tensors, %d dtype group(s), " - "%.2f MB pinned CPU memory", - len(self._managed), - len(self._groups), - total_cpu_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def offload(self): - """Per-tensor async D2H into CPU flat buffer, then free GPU storage.""" - if not self._managed: - return - if not self._initialized: - self._init_buffers() - self._ensure_stream() - - # Offload stream waits for compute to finish. - compute_event = torch.cuda.current_stream(self._device).record_event() - self._offload_stream.wait_event(compute_event) - - offloaded_bytes = 0 - - # Per-tensor D2H copies directly into CPU flat buffer slices. - # No GPU staging buffer → no temporary GPU memory spike. - with torch.cuda.stream(self._offload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - cpu_flat[off : off + n].copy_(local.reshape(-1), non_blocking=True) - - offloaded_bytes += grp["total"] * cpu_flat.element_size() - - # Wait for all D2H copies to land, then free GPU storage. - self._offload_stream.synchronize() - for t in self._managed: - storage = self._local(t).untyped_storage() - if storage.size() != 0: - storage.resize_(0) - else: - raise RuntimeError( - f"Tensor storage is already freed (size=0) before offload. " - f"This indicates a double-free or external interference. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - - if not self._logged: - logger.info( - "[CPUOffload] Offloaded %.2f MB (GPU → CPU)", - offloaded_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def reload(self): - """Per-tensor H2D from CPU flat buffer on the default stream. - - Runs on the current (default) CUDA stream to avoid stream - interaction issues with the parallel Muon pipeline. Since - pinned CPU memory is the source, the copies overlap with - GPU idle time between steps. - """ - if not self._managed or not self._initialized: - return - - reloaded_bytes = 0 - - # Re-allocate all GPU storages first. - for t in self._managed: - local = self._local(t) - storage = local.untyped_storage() - if storage.size() != 0: - raise RuntimeError( - f"Storage should have been freed (size=0) before reload, " - f"but got size={storage.size()}. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - storage.resize_(self._storage_nbytes[id(t)]) - - # Per-tensor H2D copies from CPU flat buffer slices. - # non_blocking=True with pinned source allows DMA overlap. - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off : off + n], non_blocking=True) - - reloaded_bytes += grp["total"] * cpu_flat.element_size() - - if not self._logged: - logger.info( - "[CPUOffload] Reloaded %.2f MB (CPU → GPU)", reloaded_bytes / (1024**2) - ) diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/distributed/utils.py b/build/torch29-cxx11-rocm63-x86_64-linux/distributed/utils.py deleted file mode 100644 index 683096bcd50489e7a09959753fdbc29ba0f3b6a6..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/distributed/utils.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import (Placement, Shard, - _StridedShard) - - -def _is_shard(placement: Placement) -> bool: - """Check if a placement is a shard type (Shard or _StridedShard). - - In PyTorch 2.10+, _StridedShard no longer inherits from Shard, so - ``placement.is_shard()`` returns False for _StridedShard. This helper - handles both old and new hierarchies. - """ - return isinstance(placement, (Shard, _StridedShard)) - - -def get_slices_of_dtensor( - target: DTensor | torch.Tensor, - local_rank: int, - shard_mesh: DeviceMesh, - shard_placements: tuple[Placement], -) -> tuple[slice | torch.Tensor, ...]: - """ - Get per-dimension indices for a given rank's shard of the target tensor. - - Uses ``Shard.local_shard_size_and_offset`` and - ``_StridedShard.local_shard_size_and_offset`` for correct handling of - both contiguous and strided (non-contiguous) sharding. - - Args: - target (DTensor | torch.Tensor): The target tensor (for its shape). - local_rank (int): The local rank within the shard group. - shard_mesh (DeviceMesh): The shard mesh (only shard dimensions). - shard_placements (tuple[Placement]): The shard placements. - - Returns: - A tuple of indices (one per tensor dim). Each element is either: - - A ``slice`` (for contiguous or unsharded dims) - - A 1-D ``torch.LongTensor`` of indices (for strided sharding) - """ - - # find the global rank of the local rank in the shard mesh - rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank] - - rank_coords = (shard_mesh.mesh == rank).nonzero() - - assert len(rank_coords) == 1 - rank_coords = tuple(rank_coords[0].tolist()) - - assert len(rank_coords) == len(shard_placements) - - # Track per-shard-dim indices. - # None means "not yet sharded on this dim". - dim_indices: dict[int, torch.Tensor] = {} - - # Caution: Assuming replicate-to-shard of the shard mesh goes with - # left-to-right sharding. This is ensured by the sorting logic of - # construct_shard_mesh function. - for mesh_dim_idx, (rank_coord, placement) in enumerate( - zip(rank_coords, shard_placements)): - assert _is_shard(placement) - - num_chunks = shard_mesh.mesh.shape[mesh_dim_idx] - shard_dim = placement.dim - - # Current effective size on this dim (may already be sub-sharded) - if shard_dim in dim_indices: - curr_size = len(dim_indices[shard_dim]) - else: - curr_size = target.size()[shard_dim] - - # Compute indices for this level of sharding - if isinstance(placement, _StridedShard): - _shard_size, offsets = _StridedShard.local_shard_size_and_offset( - placement, - curr_size, - num_chunks, - rank_coord, - return_first_offset=False) - new_indices = torch.tensor(offsets, dtype=torch.long) - else: - shard_size, offset = Shard.local_shard_size_and_offset( - curr_size, num_chunks, rank_coord) - new_indices = torch.arange(offset, - offset + shard_size, - dtype=torch.long) - - # Compose with previous indices on this dim - if shard_dim in dim_indices: - dim_indices[shard_dim] = dim_indices[shard_dim][new_indices] - else: - dim_indices[shard_dim] = new_indices - - # Build result tuple - result: list[slice | torch.Tensor] = [] - for d in range(len(target.size())): - if d not in dim_indices: - result.append(slice(None)) - else: - indices = dim_indices[d] - # Convert contiguous indices to slice for efficiency - if len(indices) > 0: - start = indices[0].item() - expected = torch.arange(start, - start + len(indices), - dtype=torch.long) - if torch.equal(indices, expected): - result.append(slice(start, start + len(indices))) - else: - result.append(indices) - else: - result.append(slice(0, 0)) - - return tuple(result) - - -_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, - ProcessGroup]] = dict() - - -def construct_shard_mesh( - placements: tuple[Placement], - mesh: DeviceMesh, -) -> tuple[DeviceMesh, ProcessGroup, tuple[Placement, ...]]: - """Construct shard sub-mesh and ProcessGroup for all-to-all communication. - - Given a DTensor's placements and device mesh, extracts the "shard group" - — the set of ranks that together hold all shards of the same replica — - and creates a ProcessGroup for all-to-all among them. - - Steps: - 1. Sort placements: Replicate first, then Shard by (dim, granularity). - 2. Permute the mesh tensor to match the sorted order. - 3. Collapse Replicate dims → list of shard sub-meshes (one per replica). - 4. Create/retrieve a cached ProcessGroup for the current rank's sub-mesh. - - Example — 8 GPUs, mesh shape (2, 2, 2), - placements ``[Shard(0), Replicate, _StridedShard(0)]``:: - - Step 1 — Sort: [Replicate, _StridedShard(0), Shard(0)] - Permutation: [1, 2, 0] - - Step 2 — Permute mesh dims by [1, 2, 0]: - Original: Permuted: - [[[0,1],[2,3]], [[[0,2],[1,3]], - [[4,5],[6,7]]] [[4,6],[5,7]]] - - Step 3 — Unbind replicate dim (dim 0), giving 2 shard sub-meshes: - sub-mesh 0 = [[0,2],[1,3]] (replica group 0) - sub-mesh 1 = [[4,6],[5,7]] (replica group 1) - shard_placements = (_StridedShard(0), Shard(0)) - - Step 4 — Rank 0 → ProcessGroup([0,1,4,5]) - Rank 2 → ProcessGroup([2,3,6,7]) - - Returns: - ``(shard_mesh, process_group, shard_placements)`` - """ - my_rank = dist.get_rank() - assert mesh.mesh.device.type == 'cpu' - - # -- Fast path: 1D all-shard mesh → reuse existing PG. ---------------- - # Reuses the mesh's existing ProcessGroup directly, avoiding the - # overhead of dist.new_group(). The standard path below also handles - # subset calls safely via use_local_synchronization=True, but this - # fast path is still beneficial for the common 1D shard case. - if mesh.ndim == 1 and len(placements) == 1 and _is_shard(placements[0]): - key = (*mesh.mesh.shape, *mesh.mesh.flatten().tolist()) - if key not in _ranks_to_dist_cache: - _ranks_to_dist_cache[key] = (mesh, mesh.get_group()) - return (*_ranks_to_dist_cache[key], tuple(placements)) - - mesh_tensor = mesh.mesh.clone() - - # -- Step 1: Sort placements (Replicate first, then Shard by dim). ------ - # _StridedShard comes BEFORE regular Shard on the same dim so that - # get_slices_of_dtensor applies the outer sharding first, matching - # DTensor's left-to-right (outer-to-inner) composition order. - def _sort_key(item): - index, placement = item - assert not placement.is_partial(), "Partial placement not supported" - if placement.is_replicate(): - return (-1, 0, index) - assert _is_shard(placement), f"Unsupported: {type(placement)}" - split = (-1 / placement.split_factor if isinstance( - placement, _StridedShard) else 0) - return (placement.dim, split, index) - - indexed = sorted(enumerate(placements), key=_sort_key) - perm, sorted_placements = zip(*indexed) - - # -- Step 2: Permute mesh to match sorted placement order. -------------- - sorted_mesh = mesh_tensor.permute(perm) - - # -- Step 3: Collapse replicate dims → list of shard sub-meshes. -------- - # E.g. mesh (2, 3, 4, 4) with [R, R, S(0), S(1)] → 6 sub-meshes of (4, 4) - num_rep = sum(1 for p in sorted_placements if p.is_replicate()) - if num_rep > 0: - if num_rep > 1: - sorted_mesh = sorted_mesh.flatten(0, num_rep - 1) - shard_meshes = list(torch.unbind(sorted_mesh, dim=0)) - else: - shard_meshes = [sorted_mesh] - shard_placements = sorted_placements[num_rep:] - assert len(shard_placements) == len(set(shard_placements)) - - # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. -- - # Each rank only creates the group it belongs to, using - # use_local_synchronization=True so that only group members need to - # coordinate. This avoids deadlocks when different PP stages call - # construct_shard_mesh for different parameters. - def _cache_key(t: torch.Tensor) -> tuple: - return (*t.shape, *t.flatten().tolist()) - - my_key = None - for sm in shard_meshes: - if (my_rank == sm).any().item(): - key = _cache_key(sm) - assert my_key is None, "Rank appears in multiple shard groups" - my_key = key - if key not in _ranks_to_dist_cache: - pg = dist.new_group(sm.flatten().tolist(), - use_local_synchronization=True) - _ranks_to_dist_cache[key] = ( - DeviceMesh(device_type="cuda", mesh=sm), - pg, - ) - - return (*_ranks_to_dist_cache[my_key], shard_placements) diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/matmul_transpose_triton.py b/build/torch29-cxx11-rocm63-x86_64-linux/matmul_transpose_triton.py deleted file mode 100644 index 792de23d82c3fb45fe33d397ab9b76a0787259d0..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/matmul_transpose_triton.py +++ /dev/null @@ -1,122 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], - restore_value=['y'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -@torch.library.custom_op("muon::matmul_transpose_assign", - mutates_args=("d_out", )) -def matmul_transpose_assign(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """Compute d_out = d_in @ d_in.T using an optimized Triton kernel.""" - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -@matmul_transpose_assign.register_fake -def _(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """FakeTensor impl: d_out is already allocated, mutation is declared.""" - pass diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/metadata.json b/build/torch29-cxx11-rocm63-x86_64-linux/metadata.json deleted file mode 100644 index c55a35717622f1dd5c8ba376ea3a814cbcc10d78..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/metadata.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/muon.py b/build/torch29-cxx11-rocm63-x86_64-linux/muon.py deleted file mode 100644 index 14c0e22471fa6d47a51ed95e0e0c341dc18d5194..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/muon.py +++ /dev/null @@ -1,1068 +0,0 @@ -import logging -import types -from collections import defaultdict -from typing import Any - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Replicate, Shard -from torch.profiler import record_function - -from .adamw import _placement_cache, _tensor_cache, step_adamw -from .async_utils import run_pipeline -from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho, - get_default_muon_param_groups, is_expert_param, update_p) -from .cpu_offload import CPUOffloadPool -from .distributed.utils import (_is_shard, construct_shard_mesh, - get_slices_of_dtensor) -from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO, - _zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5_batched) -from .pipeline import muon_chunk_pipeline, prelaunch_first_gather -from .qk_clip import compute_scales, get_qk_clip_info, qk_clip - -logger = logging.getLogger(__name__) - - -def _expand_expert_params(names, params, expert_keys): - """Expand expert params by splitting on dim 0 (expert dimension). - - Params whose name matches any key in ``expert_keys`` are treated as - expert-parallel tensors. Their outermost dimension is the expert - dimension: an ``(E, out, in)`` tensor becomes ``E`` separate 2D - ``nn.Parameter`` views so that in-place updates propagate back to - the original storage. - - Non-expert params with ``ndim > 2`` trigger an ``AssertionError`` — - if they are expert params, their key must be added to ``expert_keys``. - - The grad must already be set on each expert param (e.g. after momentum). - - For DTensor expert params, placements that shard on dim 0 (expert dim) - are consumed by the split. Non-dim-0 shard placements (e.g. TP) are - preserved: each 2D slice is wrapped as a DTensor on the corresponding - submesh so the parallel pipeline handles the TP communication. - """ - expanded_names = [] - expanded_params = [] - - for n, p in zip(names, params): - is_expert = is_expert_param(n, expert_keys) - is_dtensor = isinstance(p.data, DTensor) - - if is_expert: - if is_dtensor: - logger.debug( - "[expand_expert] %s: expert DTensor, shape=%s, " - "placements=%s, mesh=%s, local_shape=%s", n, p.shape, - p.placements, p.device_mesh.mesh_dim_names, - p.to_local().shape) - else: - logger.debug( - "[expand_expert] %s: expert plain tensor, shape=%s", n, - p.data.shape) - - if not is_expert: - assert p.data.ndim <= 2, ( - f"Param {n} has ndim={p.data.ndim} but does not match " - f"expert_keys={expert_keys}. If this is an expert param, " - f"add its key to expert_keys.") - expanded_names.append(n) - expanded_params.append(p) - continue - - g = p.grad - assert g is not None, ( - f"Expert param {n} must have grad set before expansion") - - tp_mesh = None - tp_placements_2d = None - - if is_dtensor: - local_data = p.to_local() - local_grad = g.to_local() if isinstance(g, DTensor) else g - - # Find non-dim-0 shard placements (e.g. TP sharding). - # After splitting on dim 0, Shard(k) becomes Shard(k-1). - tp_dim_indices = [] - tp_placements_2d = [] - for i, pl in enumerate(p.placements): - if _is_shard(pl) and pl.dim != 0: - tp_dim_indices.append(i) - tp_placements_2d.append(Shard(pl.dim - 1)) - - if tp_dim_indices: - tp_dim_names = tuple(p.device_mesh.mesh_dim_names[i] - for i in tp_dim_indices) - if len(tp_dim_names) == 1: - tp_mesh = p.device_mesh[tp_dim_names[0]] - else: - tp_mesh = p.device_mesh[tp_dim_names] - else: - local_data = p.data - local_grad = g - - # Expand: split dim 0, reshape each slice to 2D. - num_local_experts = local_data.shape[0] - for i in range(num_local_experts): - slice_data = local_data[i] - slice_grad = local_grad[i] - - if tp_mesh is not None: - # Wrap as DTensor on TP submesh so the pipeline handles - # TP communication (gather/scatter across TP ranks). - dt_data = DTensor.from_local(slice_data, - device_mesh=tp_mesh, - placements=tp_placements_2d) - dt_grad = DTensor.from_local(slice_grad, - device_mesh=tp_mesh, - placements=tp_placements_2d) - expert_param = torch.nn.Parameter(dt_data, requires_grad=False) - expert_param.grad = dt_grad - else: - expert_param = torch.nn.Parameter(slice_data, - requires_grad=False) - expert_param.grad = slice_grad - - expanded_names.append(f"{n}[{i}]") - expanded_params.append(expert_param) - - p.grad = None # allow expert grad storage to be freed after pipeline - - return expanded_names, expanded_params - - -class Muon(torch.optim.Optimizer): - """ - Muon - MomentUm Orthogonalized by Newton-schulz - - Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- - processing step, in which each 2D parameter's update is replaced with the nearest orthogonal - matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has - the advantage that it can be stably run in bfloat16 on the GPU. - - Some warnings: - - We believe this optimizer is unlikely to work well for training with small batch size. - - We believe it may not work well for finetuning pretrained models, but we haven't tested this. - - Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. - lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) - momentum: The momentum used by the internal SGD. (0.95 is a good default) - nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) - ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. - Parameters that are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW instead. - adamw_lr: The learning rate for the internal AdamW. - adamw_betas: The betas for the internal AdamW. - adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - warmup_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher warmup_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. - chunk_size : Batch size of parameters to process in each - all2all gather/compute/scatter step. - Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified. - use_distributed_muon: Use distributed muon by Liu et al. (2024). - For testing purpose only. - expert_keys: List of strings to identify expert-parallel parameters. - If any key appears in a parameter's name, its outermost - dimension is treated as the expert dimension and expanded - into per-expert 2D params for Muon. For example, - ``expert_keys=["experts"]`` matches any param whose name - contains "experts". 3D+ params not matched by any key - will raise an error. - """ - - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config=None, - warmup_step=5, - chunk_size=-1, - use_distributed_muon=False, - expert_keys=None): - defaults = dict( - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - nesterov=nesterov, - ns_steps=ns_steps, - adamw_betas=adamw_betas, - adamw_eps=adamw_eps, - none_grad=none_grad, - use_muon=True, - ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) - super().__init__(params, defaults) - - self.debug = debug - self.clip_config = clip_config if clip_config is not None else { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100, - } - self.warmup_step = warmup_step - self.chunk_size = chunk_size - self.use_distributed_muon = use_distributed_muon - self.expert_keys = expert_keys - self.cpu_offload = False - self._cpu_offload_pool: CPUOffloadPool | None = None - self._offload_initialized = False - self._parallel_cache: dict[tuple[str, ...], dict] = {} - self._expert_expand_cache: dict[tuple[int, ...], dict] = {} - - def _calc_flops(self, G, steps): - assert len(G.shape) == 2 - M, N = G.shape - if M > N: - M, N = N, M - - return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) - - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - p.placements, p.device_mesh) - - return shard_mesh, shard_pg, shard_placements - - def init_state_and_assign_params(self, names, params, group, qk_logits): - param_to_state = {} - param_to_flops = {} - - total_flops = 0 - for p in params: - g = p.grad - if g is None: - continue - assert g.ndim == 2, "Muon only supports 2D parameters." - - flops = self._calc_flops(g, group["ns_steps"]) - param_to_flops[id(p)] = flops - total_flops += flops - - if self.debug: - logger.debug("Total TFLOPs for Muon: %.2f TFLOPs", - total_flops / 1e12) - - paired = list(zip(names, params)) - - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) - - round_robin = 0 - mesh = ordered_params[0].device_mesh - placements = ordered_params[0].placements - - shard_mesh, shard_pg, shard_placements = self.get_shard_mesh( - ordered_params[0]) - shard_mesh_flattened = shard_mesh.mesh.flatten() - num_ranks = dist.get_world_size(group=shard_pg) - - for n, p in zip(ordered_names, ordered_params): - if mesh != p.device_mesh: - raise ValueError("All parameters must be on the same mesh.") - if placements != p.placements: - raise ValueError("All parameters must have same placements.") - - worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks - round_robin = (round_robin + 1) % len(shard_mesh_flattened) - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - # Precompute per-rank indices and numels for all-to-all. - rank_indices: dict[int, tuple] = {} - rank_numels: dict[int, int] = {} - for r in range(num_ranks): - indices = get_slices_of_dtensor(p, r, shard_mesh, - shard_placements) - rank_indices[r] = indices - numel = 1 - for idx, dim_size in zip(indices, p.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) - rank_numels[r] = numel - - param_to_state[id(p)] = _muon_state( - worker_rank=worker_rank, - process_group=shard_pg, - rank_indices=rank_indices, - rank_numels=rank_numels, - name=n, - qk_clip_state=qk_clip_state, - ) - - return param_to_state, ordered_params - - def base(self, names, params, group, lr, weight_decay, qk_logits): - # Momentum is already applied by _step_muon before this method. - for n, p in zip(names, params): - g = p.grad - if g is None: - continue - - u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) - - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - def distributed_muon( - self, - names: list[str], - params: list[torch.nn.Parameter], - group: dict[str, Any], - lr: float, - weight_decay: float, - qk_logits: list[torch.Tensor | DTensor] | None, - ): - """Batched Distributed Muon — for testing/correctness verification only. - - Uses all-gather to reconstruct full tensors, computes Newton-Schulz on - the full grad, then slices back to local shards. This is simpler but - slower than the parallel pipeline (all2all) path, so it serves as a - reference implementation for verifying correctness. - """ - with record_function("distributed_muon"): - # Momentum is already applied by _step_muon before this method. - ns_steps = group["ns_steps"] - - # Separate plain tensors (no communication) from DTensors. - plain_names, plain_params = [], [] - dtensor_names, dtensor_params = [], [] - for n, p in zip(names, params): - if p.grad is None: - continue - if isinstance(p.data, DTensor): - dtensor_names.append(n) - dtensor_params.append(p) - else: - plain_names.append(n) - plain_params.append(p) - - # Process plain tensors per-param (no communication). - for n, p in zip(plain_names, plain_params): - u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE), - steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, - qk_logits) - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - if not dtensor_params: - return - - # Group DTensors by (placements, mesh) for batched all-gather. - placement_groups: dict[tuple, - tuple[list, - list]] = defaultdict(lambda: ([], [])) - for n, p in zip(dtensor_names, dtensor_params): - key = (p.placements, p.device_mesh) - placement_groups[key][0].append(n) - placement_groups[key][1].append(p) - - logger.info( - "distributed_muon: %d placement groups, %d total dtensors", - len(placement_groups), len(dtensor_params)) - - for (placements, mesh), (grp_names, - grp_params) in placement_groups.items(): - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - placements, mesh) - rank = dist.get_rank(shard_pg) - world_size = dist.get_world_size(shard_pg) - - logger.info(" group: %d params, placements=%s, world_size=%d", - len(grp_params), placements, world_size) - - # Separate params that can be batched (all shard dims evenly - # divisible) from those needing per-param full_tensor - # (e.g. MoE gate weights with fewer rows than shard ranks). - # all_gather_into_tensor requires equal buffer sizes across - # ranks, so uneven splits must use DTensor full_tensor(). - batch_names, batch_params = [], [] - single_names, single_params = [], [] - for n, p in zip(grp_names, grp_params): - even = all(p.shape[pl.dim] % - shard_mesh.mesh.shape[dim_idx] == 0 - for dim_idx, pl in enumerate(shard_placements)) - if even: - batch_names.append(n) - batch_params.append(p) - else: - single_names.append(n) - single_params.append(p) - - # Process uneven-split params per-param via full_tensor(). - for n, p in zip(single_names, single_params): - with record_function("distributed_muon::newton_schulz"): - g_full = p.grad.full_tensor().to(COMM_DTYPE) - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - if not batch_params: - continue - - logger.info(" batched=%d, single=%d", len(batch_params), - len(single_params)) - - # Concat all local grad shards into a single flat buffer. - with record_function("distributed_muon::gather"): - grad_locals = [ - p.grad.to_local().to(COMM_DTYPE).flatten() - for p in batch_params - ] - numels = [g.numel() for g in grad_locals] - grad_concat = torch.cat(grad_locals) - del grad_locals - - # Single all-gather (replaces N separate full_tensor). - grad_gathered = torch.empty( - grad_concat.numel() * world_size, - dtype=COMM_DTYPE, - device="cuda", - ) - dist.all_gather_into_tensor(grad_gathered, - grad_concat, - group=shard_pg) - - total_numel = grad_concat.numel() - del grad_concat - - # Precompute per-param offsets within the concat buffer. - offsets = [] - off = 0 - for ne in numels: - offsets.append(off) - off += ne - - # Per-param: reconstruct full grad → NS → local update. - for i, (n, p) in enumerate(zip(batch_names, batch_params)): - with record_function("distributed_muon::newton_schulz"): - g_full = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - for r in range(world_size): - r_start = r * total_numel + offsets[i] - shard = grad_gathered[r_start:r_start + numels[i]] - indices = get_slices_of_dtensor( - p, r, shard_mesh, shard_placements) - g_full[indices] = shard.reshape( - g_full[indices].shape) - - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - def _setup_parallel(self, names, params, group, qk_logits): - """Compute (or retrieve cached) parallel pipeline metadata. - - Returns: - (ordered_params, param_to_state, rank, chunk_size) - """ - cache_key = tuple(names) - - if cache_key not in self._parallel_cache: - # First call: compute metadata and populate cache. - param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - shard_pg = param_to_state[id(ordered_params[0])].process_group - rank = dist.get_rank(group=shard_pg) - - if self.chunk_size == -1: - shard_ranks = dist.get_world_size(shard_pg) - chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO - elif self.chunk_size > 0: - chunk_size = self.chunk_size - else: - raise ValueError( - "chunk_size must be -1 or a positive integer.") - - ordered_names = [ - param_to_state[id(p)].name for p in ordered_params - ] - name_to_state = { - param_to_state[id(p)].name: param_to_state[id(p)] - for p in ordered_params - } - self._parallel_cache[cache_key] = { - 'ordered_names': ordered_names, - 'name_to_state': name_to_state, - 'rank': rank, - 'chunk_size': chunk_size, - } - else: - # Cached path: rebuild param_to_state with current id(p) keys. - cache = self._parallel_cache[cache_key] - rank = cache['rank'] - chunk_size = cache['chunk_size'] - - name_to_param = dict(zip(names, params)) - ordered_params = [name_to_param[n] for n in cache['ordered_names']] - - param_to_state = {} - for p, n in zip(ordered_params, cache['ordered_names']): - cached_state = cache['name_to_state'][n] - param_to_state[id(p)] = _muon_state( - worker_rank=cached_state.worker_rank, - process_group=cached_state.process_group, - rank_indices=cached_state.rank_indices, - rank_numels=cached_state.rank_numels, - name=n, - qk_clip_state=get_qk_clip_info(self.clip_config, n, - qk_logits), - ) - - return ordered_params, param_to_state, rank, chunk_size - - def parallel(self, - names, - params, - group, - lr, - weight_decay, - qk_logits, - prelaunch_gather=None): - """ - Perform a parallel optimization step using Muon. - - Parameters are chunked and each chunk is processed by a - :func:`muon_chunk_pipeline` generator. :func:`run_pipeline` - interleaves multiple chunks so that communication and computation - overlap across chunks (the same overlap previously achieved by the - warmup + main-loop index scheduling). - - If ``prelaunch_gather`` is provided, it is passed to the first - chunk's generator to skip re-launching the already in-flight - A2A gather. - """ - - # Momentum is already applied by _step_muon before this method. - - ordered_params, param_to_state, rank, chunk_size = ( - self._setup_parallel(names, params, group, qk_logits)) - - def pipelines(): - first = True - for start in range(0, len(ordered_params), chunk_size): - chunk = ordered_params[start:start + chunk_size] - if chunk: - kwargs = dict( - params=chunk, - param_to_state=param_to_state, - rank=rank, - ns_steps=group["ns_steps"], - lr=lr, - weight_decay=weight_decay, - none_grad=group["none_grad"], - ) - if first and prelaunch_gather is not None: - kwargs['prelaunch_gather'] = prelaunch_gather - first = False - yield muon_chunk_pipeline(**kwargs) - - with record_function("muon::pipeline"): - run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1) - - def _step_muon(self, group, qk_logits=None): - params = group["params"] - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - # Apply momentum to all params before routing/expansion. - # Batched using _foreach_* ops (compiled, fullgraph=True). - with record_function("muon::momentum"): - active_params = [p for p in params if p.grad is not None] - if active_params: - # Ensure momentum buffers exist (avoid zeros_like when already present). - for p in active_params: - if "momentum_buffer" not in self.state[p]: - self.state[p]["momentum_buffer"] = torch.zeros_like( - p.grad) - - # Extract local tensors for compiled batch function. - local_grads = [ - p.grad._local_tensor - if isinstance(p.grad, DTensor) else p.grad - for p in active_params - ] - local_bufs = [ - self.state[p]["momentum_buffer"]._local_tensor - if isinstance(self.state[p]["momentum_buffer"], DTensor) - else self.state[p]["momentum_buffer"] - for p in active_params - ] - - # Wrap momentum as tensor for torch.compile. - batch_pre_ortho(local_grads, local_bufs, - torch.tensor(momentum), group["nesterov"]) - - # For non-nesterov, the result is the momentum buffer. - if not group["nesterov"]: - for p in active_params: - p.grad = self.state[p]["momentum_buffer"] - - # Identify batched experts for deferred NS. - # Detection is cheap (condition checks only); actual NS compute is - # deferred so it can overlap with the first chunk's A2A gather. - deferred_expert_work = [] - if self.expert_keys: - batched_expert_indices = [] - for i, (n, p) in enumerate(zip(names, params)): - if not (is_expert_param(n, self.expert_keys) - and p.grad is not None): - continue - # Eligible: plain tensor, or DTensor with no non-dim-0 shards. - if isinstance(p.data, DTensor): - has_tp = any( - _is_shard(pl) and pl.dim != 0 for pl in p.placements) - if has_tp: - continue - batched_expert_indices.append(i) - - if batched_expert_indices: - # Save refs for deferred NS; free grads from param list. - for i in batched_expert_indices: - p = params[i] - g = p.grad - local_g = (g._local_tensor - if isinstance(g, DTensor) else g) - local_data = (p.data._local_tensor if isinstance( - p.data, DTensor) else p.data) - deferred_expert_work.append((local_data, local_g)) - p.grad = None - - # Remove batched experts from lists before expansion. - keep = sorted( - set(range(len(params))) - set(batched_expert_indices)) - names = [names[i] for i in keep] - params = [params[i] for i in keep] - - def _run_deferred_expert_ns(): - """Execute deferred batched expert NS.""" - if not deferred_expert_work: - return - with record_function("muon::batched_expert_ns"): - ns_steps = group["ns_steps"] - for local_data, local_g in deferred_expert_work: - u = zeropower_via_newtonschulz5_batched( - local_g.to(COMM_DTYPE), steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:]) - local_data.mul_(1 - lr * weight_decay) - local_data.add_(u, alpha=-adjusted_lr) - - # Expand expert params by splitting on dim 0. - logger.debug("[_step_muon] before expand: %d params, expert_keys=%s", - len(params), self.expert_keys) - if self.expert_keys: - cache_key = tuple(id(p) for p in params) - cache = self._expert_expand_cache.get(cache_key) - - if cache is None: - # Cold path: full expansion + build cache metadata. - exp_names, exp_params = _expand_expert_params( - names, params, self.expert_keys) - - # Build per-expert-group info for hot-path grad updates. - grad_info = [] - exp_idx = 0 - for orig_idx, (n, p) in enumerate(zip(names, params)): - if not is_expert_param(n, self.expert_keys): - exp_idx += 1 - continue - - is_dt = isinstance(p.data, DTensor) - num_experts = (p.to_local() if is_dt else p.data).shape[0] - - # Detect TP mesh from the first expanded expert param. - tp_mesh = None - tp_pls = None - sample = exp_params[exp_idx] - if isinstance(sample.data, DTensor): - tp_mesh = sample.data.device_mesh - tp_pls = list(sample.data.placements) - - grad_info.append((orig_idx, num_experts, exp_idx, is_dt, - tp_mesh, tp_pls)) - exp_idx += num_experts - - self._expert_expand_cache[cache_key] = { - 'names': exp_names, - 'params': exp_params, - 'grad_info': grad_info, - } - names, params = exp_names, exp_params - else: - # Hot path: reuse cached params, only update expert grads. - for (orig_idx, num_experts, exp_start, is_dt, tp_mesh, - tp_pls) in cache['grad_info']: - p = params[orig_idx] - g = p.grad - local_grad = (g.to_local() - if is_dt and isinstance(g, DTensor) else g) - for i in range(num_experts): - expert_p = cache['params'][exp_start + i] - sg = local_grad[i] - if tp_mesh is not None: - expert_p.grad = DTensor.from_local( - sg, device_mesh=tp_mesh, placements=tp_pls) - else: - expert_p.grad = sg - p.grad = None - - names = cache['names'] - params = cache['params'] - else: - names, params = _expand_expert_params(names, params, - self.expert_keys) - logger.debug("[_step_muon] after expand: %d params", len(params)) - - param_dtensors = [] - name_dtensors = [] - - param_tensors = [] - name_tensors = [] - - # distributed_muon is a reference implementation for testing only. - # The parallel pipeline (all2all) path below is the production path. - if self.use_distributed_muon: - _run_deferred_expert_ns() - self.distributed_muon(names=names, - params=params, - group=group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits) - return - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - logger.debug( - "[route] %s → base (DTensor all-Replicate), " - "shape=%s, placements=%s", n, p.shape, p.placements) - param_tensors.append(p) - name_tensors.append(n) - else: - logger.debug( - "[route] %s → parallel (DTensor), shape=%s, " - "placements=%s, mesh=%s", n, p.shape, p.placements, - p.device_mesh.mesh_dim_names) - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - logger.debug("[route] %s → base (plain tensor), shape=%s", n, - p.data.shape) - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError(f"Unsupported parameter type: {type(p.data)}") - - logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, " - f"{len(param_tensors)} Tensors → base") - - def group_dtensors(dtensors, names): - # To support different placements, we group parameters by placements - # and run parallel Muon on each group. - - placement_to_params = defaultdict(lambda: ([], [])) - - assert len(dtensors) == len(names) - for p, n in zip(dtensors, names): - placement_to_params[tuple([p.placements, - p.device_mesh])][0].append(n) - placement_to_params[tuple([p.placements, - p.device_mesh])][1].append(p) - return placement_to_params - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - dtensor_group = group_dtensors(param_dtensors, name_dtensors) - - # Pre-launch the first chunk's A2A gather so that the NCCL - # communication overlaps with the (deferred) batched expert NS - # compute on the default CUDA stream. - prelaunch = None - if deferred_expert_work: - first_names, first_params = next(iter(dtensor_group.values())) - ordered, pts, rnk, csz = self._setup_parallel( - first_names, first_params, group, qk_logits) - first_chunk = ordered[:csz] - if first_chunk: - prelaunch = prelaunch_first_gather(first_chunk, pts, rnk, - group["none_grad"]) - - _run_deferred_expert_ns() - - first_group = True - for _, (names, params) in dtensor_group.items(): - pg = prelaunch if first_group else None - first_group = False - self.parallel( - names, - params, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - prelaunch_gather=pg, - ) - else: - _run_deferred_expert_ns() - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - ) - - def _register_states_for_offload(self): - """Register all optimizer state tensors with the CPU offload pool. - - Called once after the first step when states have been lazily created. - Offloads all param states (momentum buffers for Muon, moment1/moment2 - for AdamW) to free GPU memory between steps. - """ - pool = self._cpu_offload_pool - tracked = 0 - for group in self.param_groups: - for p in group["params"]: - if p not in self.state: - continue - state = self.state[p] - if group.get("use_muon", False): - if "momentum_buffer" in state: - pool.track(state["momentum_buffer"]) - tracked += 1 - else: - if "moment1" in state: - pool.track(state["moment1"]) - if "moment2" in state: - pool.track(state["moment2"]) - tracked += 1 - logger.info("[CPUOffload] Registered %d param states for offload", - tracked) - - @torch.no_grad - def step(self, closure=None, qk_logits=None): - """Perform a single optimization step. - - Args: - closure (Callable, optional): A closure that reevaluates the model - and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). - """ - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # H2D: reload optimizer states from CPU before computation. - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload() - - logger.debug("[Muon.step] expert_keys=%s, %d param groups", - self.expert_keys, len(self.param_groups)) - - for i, group in enumerate(self.param_groups): - if group["use_muon"]: - logger.debug("[Muon.step] group %d: use_muon=True, %d params", - i, len(group["params"])) - self._step_muon(group, qk_logits=qk_logits) - else: - logger.debug( - "[Muon.step] group %d: use_muon=False (AdamW), %d params", - i, len(group["params"])) - step_adamw(self.state, group) - - # D2H: offload optimizer states to CPU after computation. - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - return loss - - # ------------------------------------------------------------------ - # CPU offload public helpers - # ------------------------------------------------------------------ - - def turn_on_cpu_offload(self): - """Enable CPU offload for optimizer states.""" - if self.cpu_offload: - return - logger.info("[Muon] turn_on_cpu_offload") - self.cpu_offload = True - if not self.state: - return - self._cpu_offload_pool = CPUOffloadPool() - self._offload_initialized = False - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_off_cpu_offload(self): - """Disable CPU offload and keep optimizer states resident on GPU.""" - if not self.cpu_offload: - return - logger.info("[Muon] turn_off_cpu_offload") - if self._offload_initialized: - self._cpu_offload_pool.reload() - torch.cuda.current_stream().synchronize() - self._cpu_offload_pool = None - self._offload_initialized = False - self.cpu_offload = False - - # ------------------------------------------------------------------ - # Checkpoint support for cpu_offload - # ------------------------------------------------------------------ - - def state_dict(self) -> dict: - if self.cpu_offload: - raise RuntimeError( - "Muon.state_dict() requires turn_off_cpu_offload() before checkpoint save." - ) - return super().state_dict() - - def load_state_dict(self, state_dict: dict) -> None: - if self.cpu_offload: - raise RuntimeError( - "Muon.load_state_dict() requires turn_off_cpu_offload() before checkpoint load." - ) - super().load_state_dict(state_dict) - - # Invalidate adamw.py's module-level tensor caches so that - # the next step rebuilds them with the newly loaded state tensors. - _placement_cache.clear() - _tensor_cache.clear() diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/newton_schulz.py b/build/torch29-cxx11-rocm63-x86_64-linux/newton_schulz.py deleted file mode 100644 index d939264b69a34e7a3fa78859f34dc265a1159d59..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/newton_schulz.py +++ /dev/null @@ -1,240 +0,0 @@ -from itertools import repeat -from math import inf, sqrt - -import numpy as np -import torch - -from .matmul_transpose_triton import matmul_transpose_assign - -COMM_DTYPE = torch.bfloat16 -DEFAULT_CHUNK_SIZE_RATIO = 4 - - -def _optimal_quintic(l, u, max_iter=1000): - """ - Use the simplified Remez algorithm to find the optimal odd quintic approximant - to the constant function x -> 1 over the interval [l, u]. - - Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum - approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the - two interior equioscillation nodes q, r until convergence. Returns the - closed-form equioscillating solution when l ≈ u. - - Raises ValueError if any intermediate value (a, b, c, E, q, r) is non-finite - (NaN or inf). Raises RuntimeError if convergence is not reached within - max_iter iterations. - """ - assert 0 <= l <= u - if 1 - 5e-6 <= l / u: - return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5) - q = (3 * l + u) / 4 - r = (l + 3 * u) / 4 - E = inf - for _ in range(max_iter): - old_E = E - LHS = np.array( - [ - [l, l**3, l**5, 1], - [q, q**3, q**5, -1], - [r, r**3, r**5, 1], - [u, u**3, u**5, -1], - ] - ) - a, b, c, E = np.linalg.solve(LHS, np.ones(4)) - if not np.all(np.isfinite([a, b, c, E])): - raise ValueError( - f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}" - ) - q, r = np.sqrt( - (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / (10 * c) - ) - if not np.all(np.isfinite([q, r])): - raise ValueError(f"_optimal_quintic: non-finite node update q={q}, r={r}") - if abs(old_E - E) <= 1e-15: - break - else: - raise RuntimeError( - f"_optimal_quintic: did not converge after {max_iter} iterations" - ) - return float(a), float(b), float(c) - - -def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0): - """ - Compute the Polar Express coefficient series for `num_iters` quintic iterations. - - Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that - compose to map singular values from [l, 1] toward 1. At each step: - 1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion` - prevents near-zero singular values from stalling by raising the effective - lower bound; if it is active (cushion*u > l), the coefficients are - rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u]. - 2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the - last iteration, providing numerical headroom at the cost of a slightly slower - final convergence step. - 3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1). - - Returns a list of (a, b, c) tuples, one per iteration. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - u = 1 - assert 0 <= l <= u - safety_factor = 1 + safety_factor_eps - coefficients = [] - for iter in range(num_iters): - a, b, c = _optimal_quintic(max(l, cushion * u), u) - if cushion * u > l: - pl = a * l + b * l**3 + c * l**5 - pu = a * u + b * u**3 + c * u**5 - rescaler = 2 / (pl + pu) - a *= rescaler - b *= rescaler - c *= rescaler - if iter < num_iters - 1: - a /= safety_factor - b /= safety_factor**3 - c /= safety_factor**5 - coefficients.append((a, b, c)) - l = a * l + b * l**3 + c * l**5 - u = 2 - l - return coefficients - - -# Precomputed Polar Express coefficients (a, b, c) for 10 quintic Newton-Schulz -# iterations. Each tuple is the minimax-optimal (Remez/equioscillation) odd quintic -# approximant to x->1 over the current singular-value interval, computed once at -# import time and reused across all optimizer steps. -# -# Contrast with the former hardcoded NS coefficients (5 fixed tuples): -# - Former: empirically tuned to maximize slope at zero; did not converge -# singular values to 1, yielding US'V^T with S' ~ Uniform(0.5, 1.5) instead -# of the true polar factor UV^T. -# - Polar Express: analytically optimal per step, adapting to the shrinking -# singular-value interval [l, u] as iterations progress; converges all -# singular values to 1, producing the exact polar factor UV^T. -_coeffs_list = _optimal_composition( - l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02 -) - - -# This code is adapted from: -# KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py) -# NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress) -# matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon) -@torch.no_grad() -def _zeropower_via_newtonschulz5(G, steps): - """ - Compute the polar factor of G via the Polar Express method. - - Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c) - are the Polar Express coefficients from `_coeffs_list`. Each step is the - optimal odd quintic approximant to x -> 1 over the current singular-value - interval, minimizing the maximum approximation error (Remez / minimax criterion). - The composition maps singular values from [l, 1] to near 1, producing the - polar factor (orthogonal factor in the polar decomposition G = UP). - - `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2, - cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE - X = G # no manual typecast - - if G.size(0) > G.size(1): - X = X.T - - X = X / (X.norm() + 1e-7) - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - # Perform the NS iterations - for a, b, c in hs: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(0) > G.size(1): - X = X.T - - return X - - -@torch.no_grad() -def _zeropower_via_newtonschulz5_batched(G, steps): - """Batched polar factor computation for 3D (E, out, in) tensors. - - Same algorithm as ``_zeropower_via_newtonschulz5`` but uses - ``torch.bmm`` / ``torch.baddbmm`` instead of the 2D Triton kernel, - processing all E expert matrices in a single batched call. - """ - assert len(G.shape) == 3 - assert G.dtype == COMM_DTYPE - X = G - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - # Per-expert Frobenius norm. - X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) - - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - for a, b, c in hs: - buf1 = torch.bmm(X, X.transpose(-2, -1)) - buf2 = torch.bmm(buf1, buf1.transpose(-2, -1)) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.baddbmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - return X - - -_ns_per_shape: dict[tuple[int, ...], callable] = {} -_use_compile = True - - -def set_ns_compile(enabled: bool): - """Toggle torch.compile for Newton-Schulz iteration.""" - global _use_compile - _use_compile = enabled - - -def zeropower_via_newtonschulz5(G, steps=5): - if not _use_compile: - return _zeropower_via_newtonschulz5(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() - - -def zeropower_via_newtonschulz5_batched(G, steps=5): - """Compile-cached batched Newton-Schulz for 3D expert tensors.""" - if not _use_compile: - return _zeropower_via_newtonschulz5_batched(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile( - _zeropower_via_newtonschulz5_batched, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/optimizer/__init__.py b/build/torch29-cxx11-rocm63-x86_64-linux/optimizer/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/optimizer/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/pipeline.py b/build/torch29-cxx11-rocm63-x86_64-linux/pipeline.py deleted file mode 100644 index c0c2d515856182d8d15ad27dd4e4e093b29397d6..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/pipeline.py +++ /dev/null @@ -1,468 +0,0 @@ -import logging -from typing import Generator - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -from .core import _muon_state, adjust_lr_for_muon -from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5 -from .qk_clip import compute_scales - -logger = logging.getLogger(__name__) - -# ====================================================================== -# Stage helpers -# ====================================================================== - - -def _launch_gather( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Allocate gather buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_gather``). - gathered_grads: ``{id(p): empty_tensor}`` for owned params, - ``None`` for non-owned. - recv_counts: Per-source-rank element counts. - """ - # Allocate gathered-grad buffers - gathered_grads: dict[int, torch.Tensor | None] = {} - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - gathered_grads[id(p)] = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - else: - gathered_grads[id(p)] = None - - # Build send buffer – batch grad copies via torch.cat - # (1-2 fused kernels vs N individual narrow().copy_() calls). - send_counts = [0] * num_ranks - for p in params: - state = param_to_state[id(p)] - send_counts[state.worker_rank] += state.rank_numels[rank] - - total_send = sum(send_counts) - if total_send > 0: - # Group grad slices by destination rank in a single pass. - dst_to_grads = [[] for _ in range(num_ranks)] - for p in params: - state = param_to_state[id(p)] - n = state.rank_numels[rank] - if n > 0: - g = p.grad.to_local() - dst_to_grads[state.worker_rank].append(g.reshape(-1)) - - # Flatten in dst order and cat once. - all_slices = [] - for dst in range(num_ranks): - all_slices.extend(dst_to_grads[dst]) - send_buf = torch.cat(all_slices) - if send_buf.dtype != COMM_DTYPE: - send_buf = send_buf.to(COMM_DTYPE) - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += state.rank_numels[src] - recv_counts[src] = total - - recv_buf = torch.empty(sum(recv_counts), dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - logger.debug(f"send_buf size: {send_buf.numel()}, " - f"recv_buf size: {recv_buf.numel()}, " - f"recv_counts: {recv_counts}, " - f"send_counts: {send_counts}, " - f"process_group: {str(process_group)}") - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, gathered_grads, recv_counts - - -def _complete_gather( - recv_buf: torch.Tensor, - recv_counts: list[int], - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - param_to_state: dict[int, _muon_state], - rank: int, -) -> None: - """Reconstruct gathered grads from the recv buffer (in-place).""" - off = 0 - for src in range(len(recv_counts)): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - - indices = state.rank_indices[src] - - shard_view = gathered_grads[id(p)][indices] - n = shard_view.numel() - if n == 0: - continue - - sg = recv_buf.narrow(0, off + inner_off, n) - sg = sg.reshape(shard_view.shape) - gathered_grads[id(p)][indices] = sg - - inner_off += n - assert inner_off == block - off += block - - -def _compute_ns( - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - ns_steps: int, -) -> dict[int, torch.Tensor | None]: - """Run Newton-Schulz orthogonalization on owned parameters. - - Returns: - computed_us: ``{id(p): orthogonalized_update}`` for owned params. - """ - computed_us: dict[int, torch.Tensor | None] = {} - for p in owned_params: - u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps) - gathered_grads[id(p)] = None # free gathered grad - computed_us[id(p)] = u - return computed_us - - -def _launch_scatter( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, - computed_us: dict[int, torch.Tensor | None], -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor], list[int]]: - """Allocate scatter buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_scatter``). - scattered_us: Empty dict, populated by ``_complete_scatter`` with - zero-copy views into ``recv_buf``. - recv_counts: Per-source-rank element counts. - """ - # scattered_us is populated by _complete_scatter with zero-copy views - # into recv_buf, avoiding N empty_like allocations + N copy_ calls. - # Pre-seed entries for params whose local shard is empty (rank_numels == 0) - # so _update_params can iterate all params without KeyError. - scattered_us: dict[int, torch.Tensor] = {} - for p in params: - if param_to_state[id(p)].rank_numels[rank] == 0: - scattered_us[id(p)] = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - # Build send buffer – batch via torch.cat - # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls). - send_counts = [0] * num_ranks - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - for dst_rank in range(num_ranks): - send_counts[dst_rank] += state.rank_numels[dst_rank] - - total_send = sum(send_counts) - if total_send > 0: - # Cache u_full conversions to avoid redundant .to() per dst_rank. - u_fulls = {} - for p in owned_params: - u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous() - - # Collect slices in dst order (matches all-to-all send layout). - all_slices = [] - for dst_rank in range(num_ranks): - for p in owned_params: - state = param_to_state[id(p)] - su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten() - if su.numel() > 0: - all_slices.append(su) - - send_buf = torch.cat(all_slices) if all_slices else torch.empty( - 0, dtype=COMM_DTYPE, device="cuda") - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += state.rank_numels[rank] - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, scattered_us, recv_counts - - -def _complete_scatter( - recv_buf: torch.Tensor, - recv_counts: list[int], - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], -) -> None: - """Populate scattered_us with zero-copy views into recv_buf. - - Instead of pre-allocating tensors and copying, we assign views directly - from ``recv_buf``. This eliminates N ``empty_like`` + N ``copy_`` calls. - The underlying storage of ``recv_buf`` is kept alive through the views - until ``scattered_us`` is cleared after ``_update_params``. - """ - off = 0 - for src in range(len(recv_counts)): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = state.rank_numels[rank] - if n == 0: - continue - - scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - - inner_off += n - - assert inner_off == block - off += block - - -def _update_params( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], - lr: float, - weight_decay: float, -) -> None: - """Apply weight decay, Muon update, and optional QK clipping. - - Uses batched ``_foreach_mul_`` for weight decay and batched - ``_foreach_add_`` for the Muon update, grouping parameters by - adjusted_lr to minimize kernel launches while preserving float32 - precision for the alpha scaling. - """ - if not params: - return - - # Batched weight decay: p *= (1 - lr * wd) — single fused kernel. - p_locals = [p._local_tensor for p in params] - torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay) - - # Group params by adjusted_lr so _foreach_add_ can use a single - # alpha per group (preserves float32 precision for alpha scaling). - lr_groups: dict[float, tuple[list, list]] = {} - for p in params: - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - if adjusted_lr not in lr_groups: - lr_groups[adjusted_lr] = ([], []) - lr_groups[adjusted_lr][0].append(p._local_tensor) - lr_groups[adjusted_lr][1].append(scattered_us[id(p)]) - - for adjusted_lr, (p_group, u_group) in lr_groups.items(): - torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr) - - # QK clipping – applied directly on the local tensor to - # avoid DTensor sharding-propagation issues with _StridedShard. - for p in params: - state = param_to_state[id(p)] - if state.qk_clip_state is None: - continue - scales_full = compute_scales(p, state.qk_clip_state) - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = state.rank_indices[rank][0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - -# ====================================================================== -# Pre-launch helper for overlapping first chunk's gather with other work. -# ====================================================================== - - -@torch.no_grad() -def prelaunch_first_gather( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - none_grad: bool, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Launch the first chunk's A2A gather early for overlap with other compute. - - Call this *before* expensive GPU work (e.g. batched expert NS) so that - the NCCL all-to-all runs concurrently on the NCCL stream while the - default stream executes compute. - - Returns the same 4-tuple that ``_launch_gather`` produces, which should - be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - with record_function("muon::prelaunch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - return work, recv_buf, gathered_grads, recv_counts - - -# ====================================================================== -# Main generator – thin orchestrator that wires stages together. -# ====================================================================== - - -@torch.no_grad() -def muon_chunk_pipeline( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - ns_steps: int, - lr: float, - weight_decay: float, - none_grad: bool, - prelaunch_gather: tuple | None = None, -) -> Generator[None, None, None]: - """Process one chunk of parameters through the full Muon pipeline. - - Stages: gather -> compute (Newton-Schulz) -> scatter -> update. - - Each ``yield`` lets :func:`run_pipeline` interleave other chunks so - that communication and computation overlap across chunks. Async - communication is launched via ``async_op=True`` and completed after - the yield with ``work.wait()``. - - Overlap happens because :func:`run_pipeline` admits one new chunk - per iteration (staggered admission). While chunk *N* does NS - compute on the default CUDA stream, chunk *N+1*'s async all-to-all - runs concurrently on the NCCL stream — no separate ``comm_stream`` - is required. - - If ``prelaunch_gather`` is provided, the gather was already launched - by :func:`prelaunch_first_gather` and we skip launching it again. - - Yields exactly **2** times: - - 1. After launching async all-to-all gather (or immediately if pre-launched). - 2. After launching async all-to-all scatter. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - if prelaunch_gather is not None: - # Gather was pre-launched; none_grad already handled by caller. - work, recv_buf, gathered_grads, recv_counts = prelaunch_gather - else: - # Normal path: launch async gather. - with record_function("muon::launch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - yield # --- YIELD 1: other chunks can launch their gather --- - - with record_function("muon::wait_gather"): - work.wait() - _complete_gather(recv_buf, recv_counts, owned_params, gathered_grads, - param_to_state, rank) - del recv_buf - - # Stage 3: Newton-Schulz orthogonalization. - with record_function("muon::newton_schulz"): - computed_us = _compute_ns(owned_params, gathered_grads, ns_steps) - gathered_grads.clear() - - # Stages 4-5: launch async scatter. - with record_function("muon::launch_scatter"): - work, recv_buf, scattered_us, recv_counts = _launch_scatter( - params, owned_params, param_to_state, rank, num_ranks, - process_group, computed_us) - computed_us.clear() - - yield # --- YIELD 2: other chunks can launch their scatter --- - - with record_function("muon::wait_scatter"): - work.wait() - _complete_scatter(recv_buf, recv_counts, params, param_to_state, rank, - scattered_us) - del recv_buf - - # Stage 6: apply parameter updates. - with record_function("muon::update_params"): - _update_params(params, param_to_state, rank, scattered_us, lr, - weight_decay) - scattered_us.clear() diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/qk_clip.py b/build/torch29-cxx11-rocm63-x86_64-linux/qk_clip.py deleted file mode 100644 index 2aba711b3004b7f09e7141da7ef834bd61cc2430..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/qk_clip.py +++ /dev/null @@ -1,198 +0,0 @@ -import logging -import math -from dataclasses import dataclass - -import torch -from torch.distributed.tensor import DTensor - -from .core import normalize_fqn - -logger = logging.getLogger(__name__) - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - and return (kind, layer_index). - - Supported kinds: - MHA/GQA: 'wq', 'wk', 'q_proj', 'k_proj' - MLA: 'wq_b' (Q up-proj), 'wkv_b' (KV up-proj) - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.1.attn.wq_b.weight' -> ('wq_b', 1) - 'model.0.attn.wkv_b.weight' -> ('wkv_b', 0) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = normalize_fqn(name).split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: str | None # 'wq'/'q_proj'/'wq_b' or 'wk'/'k_proj'/'wkv_b' or None - indices: list[int] # which heads to consider for clipping - head_dim: int # from config (qk_head_dim for MLA wq_b) - threshold: float # from config - logit: torch.Tensor | None - - # MLA-specific fields - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 - - -def get_qk_clip_info(clip_config, n, qk_logits): - """Extract QK clipping info for a named parameter. - - Args: - clip_config: QK clipping configuration dict (or None). - MHA/GQA keys: head_dim, threshold, q_indices, k_indices - MLA extra keys: is_mla=True, qk_nope_head_dim, qk_rope_head_dim, v_head_dim - n: Parameter name string. - qk_logits: Dict mapping layer indices to logit tensors (or None). - - Returns: - QKClipInfo instance with clipping configuration for this parameter. - """ - if clip_config is None: - return None - - head_dim = clip_config.get('head_dim') - threshold = clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - # In TP settings, qk_logits may be DTensor - # We convert it to full tensor here for simplicity - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - -def compute_scales(p, qk_clip_state): - """Compute per-head scaling factors for QK clipping. - - Returns scales tensor (√γ per head) if any head exceeds threshold, else None. - For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim. - """ - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - # Check if any head exceeds threshold before allocating. - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - - if not head_scales: - return None - - # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows - if qk_clip_state.is_mla and kind == 'wkv_b': - effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim - else: - effective_head_dim = head_dim - - H_global = p.shape[0] // effective_head_dim - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale - return scales_full - - -def qk_clip(p, scales, info): - """Apply per-head scaling to a Q/K projection weight matrix. - - Args: - p: Parameter (nn.Parameter or raw tensor). - scales: [n_heads] tensor, each element = √γ_h. - info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions. - - MLA sub-region scaling per Algorithm 1 (MuonClip): - wq_b: q_nope rows → √γ, q_pe rows → γ - wkv_b: k_nope rows → √γ, v rows → unchanged - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not info.is_mla: - # MHA/GQA: uniform √γ applied to all rows in each head - W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: vectorized sub-region scaling within each head - if info.kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_head_dim = qk_nope + info.qk_rope_head_dim - W_3d = W.view(-1, qk_head_dim, W.shape[1]) # [H, qk_head_dim, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # q_nope → √γ - W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1, - 1)) # q_pe → γ - - elif info.kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - W_3d = W.view(-1, kv_stride, W.shape[1]) # [H, kv_stride, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # k_nope → √γ - # v rows: not touched (k_R shared rotary unchanged) diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/__init__.py b/build/torch29-cxx11-rocm64-x86_64-linux/__init__.py deleted file mode 100644 index 239c7a65f8293e7d0df28f05fce645af56d628c0..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .muon import Muon - -__all__ = [ - "Muon", -] diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/_ops.py b/build/torch29-cxx11-rocm64-x86_64-linux/_ops.py deleted file mode 100644 index 034bff088659b5df6f6d401feb18c89dc5f33b29..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _optimizer_8d53b78_dirty -ops = torch.ops._optimizer_8d53b78_dirty - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_optimizer_8d53b78_dirty::{op_name}" \ No newline at end of file diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so b/build/torch29-cxx11-rocm64-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so deleted file mode 100755 index b11d926ee80d66a0e597271dd7ecf3e693493358..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/_optimizer_8d53b78_dirty.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:f3186e65c7ef03c5272229d1b155fa3b4309ae2c119f149e20bbae41c64cd754 -size 1865232 diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/adamw.py b/build/torch29-cxx11-rocm64-x86_64-linux/adamw.py deleted file mode 100644 index b5a95816a9f5b9e1889eaadae65373bfbced809a..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/adamw.py +++ /dev/null @@ -1,271 +0,0 @@ -import logging -from collections import defaultdict -from typing import cast - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -def fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: float | torch.Tensor, - weight_decay: float, - eps: float, - maximize: bool, -) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: dict | None = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) - - -def _to_local(t): - """Unwrap DTensor to local tensor for fused ops.""" - return t._local_tensor if isinstance(t, DTensor) else t - - -# --------------------------------------------------------------------------- -# Caches for eliminating per-step Python overhead. -# -# Placement grouping and tensor list assembly are identical every step -# (params don't change placement, moment/step tensors are the same objects -# after initialisation). We cache them keyed by id() of the param list -# stored in param_groups (stable across steps). -# -# Only gradients change each step and must be collected fresh. -# --------------------------------------------------------------------------- - -# id(group["params"]) → dict[placement_key, list[param]] -_placement_cache: dict[int, dict[tuple, list]] = {} - -# id(placement_group_list) → (params_local, moment1, moment2, state_steps) -_tensor_cache: dict[int, tuple[list, list, list, list]] = {} - - -def _step_adamw_params_slow(optimizer_state, params, group): - """Uncached fallback for the rare case where some params lack grads.""" - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - g = p.grad - if g is None: - continue - state = optimizer_state[p] - params_with_grads.append(_to_local(p)) - grads.append(_to_local(g)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - if not params_with_grads: - return - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw_params(optimizer_state, params, group): - """Run fused AdamW on a list of parameters sharing the same placement. - - After the first call, cached tensor lists (params_local, moment1, - moment2, state_steps) are reused — only gradients are collected fresh. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - params: List of parameters to update. - group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay. - """ - # Collect grads — the only thing that changes each step. - with record_function("adamw::collect_grads"): - grads = [] - for p in params: - g = p.grad - if g is None: - # Rare: fall back to slow path that filters per-param. - _step_adamw_params_slow(optimizer_state, params, group) - return - grads.append(_to_local(g)) - - tensor_key = id(params) - if tensor_key not in _tensor_cache: - with record_function("adamw::init_tensor_cache"): - params_local = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - state = optimizer_state[p] - params_local.append(_to_local(p)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(p.grad) - state["moment2"] = torch.zeros_like(p.grad) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - _tensor_cache[tensor_key] = (params_local, moment1, moment2, - state_steps) - - params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key] - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - with record_function("adamw::fused_adamw"): - fused_adamw( - params_local, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw(optimizer_state, group): - """Dispatch AdamW step, grouping parameters by type and placement. - - Placement grouping is cached after the first call since params never - change their placement between steps. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - group: Parameter group dict. - """ - params = group["params"] - placement_key = id(params) - - if placement_key not in _placement_cache: - with record_function("adamw::group_by_placement"): - placement_to_params: dict[tuple, - list[torch.Tensor]] = defaultdict(list) - for p in params: - match p: - case DTensor(): - logger.debug( - "[AdamW] DTensor param: shape=%s, placements=%s, " - "mesh=%s, grad=%s", p.shape, p.placements, - p.device_mesh.mesh_dim_names, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple( - [p.placements, p.device_mesh])].append(p) - case torch.Tensor(): - logger.debug( - "[AdamW] plain param: shape=%s, grad=%s", p.shape, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple([torch.Tensor, - None])].append(p) - - logger.debug("[AdamW] %d placement groups, %d total params", - len(placement_to_params), len(params)) - - _placement_cache[placement_key] = dict(placement_to_params) - - for group_params in _placement_cache[placement_key].values(): - step_adamw_params(optimizer_state, group_params, group) diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/async_utils.py b/build/torch29-cxx11-rocm64-x86_64-linux/async_utils.py deleted file mode 100644 index a45c530ac9cad88e3555ec1047a6aa59f225347e..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/async_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import logging -from typing import Generator - -logger = logging.getLogger(__name__) - - -class _Task: - """Internal: wraps a generator, advances one yield at a time.""" - - def __init__(self, generator: Generator[None, None, None], index: int): - self._generator = generator - self._index = index - self._steps_completed = 0 - self.step() # run to first yield - - def step(self) -> bool: - try: - next(self._generator) - self._steps_completed += 1 - logger.debug("pipeline[%d] completed stage %d", self._index, - self._steps_completed) - return True - except StopIteration: - logger.debug("pipeline[%d] finished after %d stages", self._index, - self._steps_completed) - return False - - def close(self): - self._generator.close() - - -def run_pipeline( - pipelines: Generator[Generator[None, None, None], None, None], - max_concurrent: int, -) -> None: - """Run generator-based pipelines with bounded concurrency. - - Each pipeline is a generator that yields at stage boundaries. - The runtime interleaves pipelines so communication and computation - overlap across chunks. - """ - if max_concurrent <= 0: - raise ValueError(f"max_concurrent must be > 0, got {max_concurrent}") - - have_new = True - task_index = 0 - previous_tasks: list[_Task] = [] - - try: - while have_new or previous_tasks: - running_tasks: list[_Task] = [] - - # Admit one new pipeline per iteration (staggered admission). - # Admitting one at a time ensures that while chunk N does NS - # compute on the default stream, chunk N+1's NCCL all-to-all - # runs concurrently on the NCCL stream — creating real - # communication/computation overlap on the GPU. - if have_new and len(previous_tasks) < max_concurrent: - try: - gen = next(pipelines) - task = _Task(gen, task_index) - task_index += 1 - running_tasks.append(task) - except StopIteration: - have_new = False - - # Advance every previously-yielded task by one step. - for task in previous_tasks: - if task.step(): - running_tasks.append(task) - - previous_tasks = running_tasks - except BaseException: - # Clean up all in-flight generators to release GPU resources. - for task in previous_tasks: - task.close() - raise diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/core.py b/build/torch29-cxx11-rocm64-x86_64-linux/core.py deleted file mode 100644 index c69d515afef305ad0ed66374095fa2d2468d99cc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/core.py +++ /dev/null @@ -1,219 +0,0 @@ -import logging -import math -from dataclasses import dataclass -from typing import List - -import torch -from torch.distributed import ProcessGroup -from torch.distributed.tensor import DTensor - -# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into -# parameter FQNs. Activation checkpointing similarly inserts -# "_checkpoint_wrapped_module". Strip these so name-based matching (skip_keys, -# expert_keys, QK layer parsing) works regardless of wrapper nesting. -_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) - -logger = logging.getLogger(__name__) - - -def normalize_fqn(name: str) -> str: - """Strip torch.compile / checkpoint wrapper components from a parameter FQN.""" - return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS) - - -@dataclass -class _muon_state: - worker_rank: int - process_group: ProcessGroup - rank_indices: dict[int, tuple] # local_rank -> per-dim indices - rank_numels: dict[int, int] # local_rank -> numel - name: str - qk_clip_state: torch.Tensor | None = None - - -def _batch_momentum( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update (no nesterov).""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - - -def _batch_momentum_nesterov( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update with nesterov correction.""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - nesterov_terms = torch._foreach_mul(momentum_bufs, momentum) - torch._foreach_add_(grads, nesterov_terms) - - -_compiled_momentum: dict[bool, callable] = {} -_use_momentum_compile = True - - -def set_momentum_compile(enabled: bool): - """Toggle torch.compile for batched momentum.""" - global _use_momentum_compile - _use_momentum_compile = enabled - - -def batch_pre_ortho( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, - nesterov: bool, -) -> None: - """Batched momentum update on lists of plain tensors. - - Mirrors dion's ``muon_update_pre_orthogonalize``. - Inputs must be plain CUDA tensors (not DTensor). - Modifies ``momentum_bufs`` and (for nesterov) ``grads`` in-place. - - When compile is enabled, uses separately compiled functions for - nesterov=True/False to avoid graph breaks from the branch. - """ - fn = _batch_momentum_nesterov if nesterov else _batch_momentum - if _use_momentum_compile: - if nesterov not in _compiled_momentum: - _compiled_momentum[nesterov] = torch.compile(fn) - fn = _compiled_momentum[nesterov] - fn(grads, momentum_bufs, momentum) - - -def _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay): - """Weight-decay + update on plain tensors. - - Not compiled: per-param @torch.compile caused ~0.25ms TorchDynamo cache - lookup per call × 256+ params = massive overhead. The pipeline path uses - batched _foreach_* ops instead; this function remains for base() and - distributed_muon(). - """ - p_data.mul_(1 - lr * weight_decay) - p_data.add_(u_data, alpha=-adjusted_lr) - - -def update_p(p, u, lr, adjusted_lr, weight_decay): - """Apply weight decay and orthogonalized update to parameter. - - Args: - p: Parameter (torch.nn.Parameter or DTensor). - u: Orthogonalized update tensor. - lr: Base learning rate. - adjusted_lr: Size-adjusted learning rate. - weight_decay: Weight decay coefficient. - """ - # Unwrap Parameter -> underlying data tensor. - p_data = p.data if isinstance(p, torch.nn.Parameter) else p - # Unwrap DTensor -> local CUDA tensor for compiled kernel. - if isinstance(p_data, DTensor): - p_data = p_data._local_tensor - u_data = u._local_tensor if isinstance(u, DTensor) else u - _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay) - - -def adjust_lr_for_muon(lr, param_shape): - """Scale learning rate based on parameter matrix dimensions. - - Args: - lr: Base learning rate. - param_shape: Shape of the parameter tensor. - - Returns: - Adjusted learning rate. - """ - A, B = param_shape[:2] - # We adjust the learning rate and weight decay based on the size of the parameter matrix - # as described in the paper - adjusted_ratio = 0.2 * math.sqrt(max(A, B)) - adjusted_lr = lr * adjusted_ratio - return adjusted_lr - - -def _match_key(parts, key): - """Check if key matches as contiguous components in parts. - - Single-component keys (e.g. "experts") match any single component. - Multi-component keys (e.g. "experts.w1") match as a contiguous subsequence. - """ - key_parts = key.split(".") - key_len = len(key_parts) - if key_len == 1: - return key in parts - return any(parts[i:i + key_len] == key_parts - for i in range(len(parts) - key_len + 1)) - - -def is_expert_param(name, expert_keys): - """Check if a parameter name matches any expert key (component-level).""" - if not expert_keys: - return False - parts = normalize_fqn(name).split(".") - return any(_match_key(parts, key) for key in expert_keys) - - -def default_is_muon(name, x, expert_keys=None): - normalized = normalize_fqn(name) - parts = normalized.split(".") - skip_keys = [ - "embed_tokens", - "lm_head", - "tok_embeddings", - "output", - "mhc_attn", - "mhc_ffn", - "lambda_proj", - ] - if any(key in parts for key in skip_keys): - logger.info( - "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d", - normalized, name, x.ndim) - return False - effective_ndim = x.ndim - is_expert = is_expert_param(name, expert_keys) - if is_expert: - effective_ndim -= 1 - result = effective_ndim >= 2 - logger.info( - "[is_muon] %s (orig: %s): ndim=%d, expert=%s, effective_ndim=%d → %s", - normalized, name, x.ndim, is_expert, effective_ndim, - "Muon" if result else "AdamW") - return result - - -def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None): - if is_muon_func is None: - is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys) - - muon_params, muon_names = [], [] - non_muon_params, non_muon_names = [], [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - non_muon_names.append(n) - - logger.info("[param_groups] expert_keys=%s, Muon=%d, AdamW=%d", - expert_keys, len(muon_names), len(non_muon_names)) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/cpu_offload.py b/build/torch29-cxx11-rocm64-x86_64-linux/cpu_offload.py deleted file mode 100644 index fb5e69154a1d4a6c884491413a37a9acf0f66c80..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/cpu_offload.py +++ /dev/null @@ -1,206 +0,0 @@ -"""CPU offloading for optimizer states. - -Manages a pinned CPU memory pool and async CUDA streams to offload -optimizer state tensors (momentum buffers, Adam moments) to CPU between -optimizer steps, freeing GPU memory. - -All tracked tensors are packed into a single flat pinned CPU buffer -(per dtype). D2H and H2D copies are performed per-tensor directly -between individual GPU tensors and their slice of the CPU flat buffer -— no GPU staging buffer is allocated, so there is **no temporary GPU -memory spike** during offload or reload. - -Individual tensor storages are freed after offload via -``untyped_storage().resize_(0)``, preserving tensor identity so -downstream caches remain valid. -""" - -import logging -from collections import defaultdict - -import torch -from torch.distributed.tensor import DTensor - -logger = logging.getLogger(__name__) - - -class CPUOffloadPool: - """Pinned CPU memory pool for async optimizer state offloading. - - Tracked tensors are grouped by dtype. Each group gets a single flat - pinned CPU buffer. D2H / H2D copies are per-tensor (into slices of - the flat buffer) to avoid allocating a GPU staging buffer. - """ - - def __init__(self): - self._managed: list[torch.Tensor] = [] - self._storage_nbytes: dict[int, int] = {} # id(t) → bytes - - # Per-dtype group: populated on first offload. - # dtype → dict with keys: - # "indices" : list[int] managed-list indices - # "offsets" : list[tuple[int,int]] (start, numel) in flat buf - # "total" : int total numel - # "cpu_flat" : Tensor pinned CPU buffer - self._groups: dict[torch.dtype, dict] = {} - - self._offload_stream: torch.cuda.Stream | None = None - self._device: torch.device | None = None - self._initialized: bool = False - self._logged: bool = False - - # ------------------------------------------------------------------ - @staticmethod - def _local(t: torch.Tensor) -> torch.Tensor: - """Unwrap DTensor to its local CUDA tensor.""" - return t._local_tensor if isinstance(t, DTensor) else t - - def _ensure_stream(self): - if self._offload_stream is None: - self._offload_stream = torch.cuda.Stream(device=self._device) - - # ------------------------------------------------------------------ - def track(self, tensor: torch.Tensor): - """Register a GPU tensor for CPU offloading. Idempotent.""" - tid = id(tensor) - if tid in self._storage_nbytes: - return - local = self._local(tensor) - if self._device is None: - self._device = local.device - storage = local.untyped_storage() - # Skip tensors with empty storage (e.g. empty FSDP shards) - if storage.size() == 0: - return - self._storage_nbytes[tid] = storage.size() - self._managed.append(tensor) - - # ------------------------------------------------------------------ - def _init_buffers(self): - """Build per-dtype flat buffers on first offload.""" - # Group managed tensors by dtype. - dtype_map: dict[torch.dtype, list[tuple[int, int]]] = defaultdict(list) - for idx, t in enumerate(self._managed): - local = self._local(t) - dtype_map[local.dtype].append((idx, local.numel())) - - total_cpu_bytes = 0 - for dtype, entries in dtype_map.items(): - offsets: list[tuple[int, int]] = [] - indices: list[int] = [] - off = 0 - for idx, n in entries: - indices.append(idx) - offsets.append((off, n)) - off += n - cpu_flat = torch.empty(off, dtype=dtype, device="cpu", pin_memory=True) - self._groups[dtype] = { - "indices": indices, - "offsets": offsets, - "total": off, - "cpu_flat": cpu_flat, - } - total_cpu_bytes += off * cpu_flat.element_size() - - self._initialized = True - logger.info( - "[CPUOffload] Pool initialized: %d tensors, %d dtype group(s), " - "%.2f MB pinned CPU memory", - len(self._managed), - len(self._groups), - total_cpu_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def offload(self): - """Per-tensor async D2H into CPU flat buffer, then free GPU storage.""" - if not self._managed: - return - if not self._initialized: - self._init_buffers() - self._ensure_stream() - - # Offload stream waits for compute to finish. - compute_event = torch.cuda.current_stream(self._device).record_event() - self._offload_stream.wait_event(compute_event) - - offloaded_bytes = 0 - - # Per-tensor D2H copies directly into CPU flat buffer slices. - # No GPU staging buffer → no temporary GPU memory spike. - with torch.cuda.stream(self._offload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - cpu_flat[off : off + n].copy_(local.reshape(-1), non_blocking=True) - - offloaded_bytes += grp["total"] * cpu_flat.element_size() - - # Wait for all D2H copies to land, then free GPU storage. - self._offload_stream.synchronize() - for t in self._managed: - storage = self._local(t).untyped_storage() - if storage.size() != 0: - storage.resize_(0) - else: - raise RuntimeError( - f"Tensor storage is already freed (size=0) before offload. " - f"This indicates a double-free or external interference. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - - if not self._logged: - logger.info( - "[CPUOffload] Offloaded %.2f MB (GPU → CPU)", - offloaded_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def reload(self): - """Per-tensor H2D from CPU flat buffer on the default stream. - - Runs on the current (default) CUDA stream to avoid stream - interaction issues with the parallel Muon pipeline. Since - pinned CPU memory is the source, the copies overlap with - GPU idle time between steps. - """ - if not self._managed or not self._initialized: - return - - reloaded_bytes = 0 - - # Re-allocate all GPU storages first. - for t in self._managed: - local = self._local(t) - storage = local.untyped_storage() - if storage.size() != 0: - raise RuntimeError( - f"Storage should have been freed (size=0) before reload, " - f"but got size={storage.size()}. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}" - ) - storage.resize_(self._storage_nbytes[id(t)]) - - # Per-tensor H2D copies from CPU flat buffer slices. - # non_blocking=True with pinned source allows DMA overlap. - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off : off + n], non_blocking=True) - - reloaded_bytes += grp["total"] * cpu_flat.element_size() - - if not self._logged: - logger.info( - "[CPUOffload] Reloaded %.2f MB (CPU → GPU)", reloaded_bytes / (1024**2) - ) diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/distributed/utils.py b/build/torch29-cxx11-rocm64-x86_64-linux/distributed/utils.py deleted file mode 100644 index 683096bcd50489e7a09959753fdbc29ba0f3b6a6..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/distributed/utils.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import (Placement, Shard, - _StridedShard) - - -def _is_shard(placement: Placement) -> bool: - """Check if a placement is a shard type (Shard or _StridedShard). - - In PyTorch 2.10+, _StridedShard no longer inherits from Shard, so - ``placement.is_shard()`` returns False for _StridedShard. This helper - handles both old and new hierarchies. - """ - return isinstance(placement, (Shard, _StridedShard)) - - -def get_slices_of_dtensor( - target: DTensor | torch.Tensor, - local_rank: int, - shard_mesh: DeviceMesh, - shard_placements: tuple[Placement], -) -> tuple[slice | torch.Tensor, ...]: - """ - Get per-dimension indices for a given rank's shard of the target tensor. - - Uses ``Shard.local_shard_size_and_offset`` and - ``_StridedShard.local_shard_size_and_offset`` for correct handling of - both contiguous and strided (non-contiguous) sharding. - - Args: - target (DTensor | torch.Tensor): The target tensor (for its shape). - local_rank (int): The local rank within the shard group. - shard_mesh (DeviceMesh): The shard mesh (only shard dimensions). - shard_placements (tuple[Placement]): The shard placements. - - Returns: - A tuple of indices (one per tensor dim). Each element is either: - - A ``slice`` (for contiguous or unsharded dims) - - A 1-D ``torch.LongTensor`` of indices (for strided sharding) - """ - - # find the global rank of the local rank in the shard mesh - rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank] - - rank_coords = (shard_mesh.mesh == rank).nonzero() - - assert len(rank_coords) == 1 - rank_coords = tuple(rank_coords[0].tolist()) - - assert len(rank_coords) == len(shard_placements) - - # Track per-shard-dim indices. - # None means "not yet sharded on this dim". - dim_indices: dict[int, torch.Tensor] = {} - - # Caution: Assuming replicate-to-shard of the shard mesh goes with - # left-to-right sharding. This is ensured by the sorting logic of - # construct_shard_mesh function. - for mesh_dim_idx, (rank_coord, placement) in enumerate( - zip(rank_coords, shard_placements)): - assert _is_shard(placement) - - num_chunks = shard_mesh.mesh.shape[mesh_dim_idx] - shard_dim = placement.dim - - # Current effective size on this dim (may already be sub-sharded) - if shard_dim in dim_indices: - curr_size = len(dim_indices[shard_dim]) - else: - curr_size = target.size()[shard_dim] - - # Compute indices for this level of sharding - if isinstance(placement, _StridedShard): - _shard_size, offsets = _StridedShard.local_shard_size_and_offset( - placement, - curr_size, - num_chunks, - rank_coord, - return_first_offset=False) - new_indices = torch.tensor(offsets, dtype=torch.long) - else: - shard_size, offset = Shard.local_shard_size_and_offset( - curr_size, num_chunks, rank_coord) - new_indices = torch.arange(offset, - offset + shard_size, - dtype=torch.long) - - # Compose with previous indices on this dim - if shard_dim in dim_indices: - dim_indices[shard_dim] = dim_indices[shard_dim][new_indices] - else: - dim_indices[shard_dim] = new_indices - - # Build result tuple - result: list[slice | torch.Tensor] = [] - for d in range(len(target.size())): - if d not in dim_indices: - result.append(slice(None)) - else: - indices = dim_indices[d] - # Convert contiguous indices to slice for efficiency - if len(indices) > 0: - start = indices[0].item() - expected = torch.arange(start, - start + len(indices), - dtype=torch.long) - if torch.equal(indices, expected): - result.append(slice(start, start + len(indices))) - else: - result.append(indices) - else: - result.append(slice(0, 0)) - - return tuple(result) - - -_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, - ProcessGroup]] = dict() - - -def construct_shard_mesh( - placements: tuple[Placement], - mesh: DeviceMesh, -) -> tuple[DeviceMesh, ProcessGroup, tuple[Placement, ...]]: - """Construct shard sub-mesh and ProcessGroup for all-to-all communication. - - Given a DTensor's placements and device mesh, extracts the "shard group" - — the set of ranks that together hold all shards of the same replica — - and creates a ProcessGroup for all-to-all among them. - - Steps: - 1. Sort placements: Replicate first, then Shard by (dim, granularity). - 2. Permute the mesh tensor to match the sorted order. - 3. Collapse Replicate dims → list of shard sub-meshes (one per replica). - 4. Create/retrieve a cached ProcessGroup for the current rank's sub-mesh. - - Example — 8 GPUs, mesh shape (2, 2, 2), - placements ``[Shard(0), Replicate, _StridedShard(0)]``:: - - Step 1 — Sort: [Replicate, _StridedShard(0), Shard(0)] - Permutation: [1, 2, 0] - - Step 2 — Permute mesh dims by [1, 2, 0]: - Original: Permuted: - [[[0,1],[2,3]], [[[0,2],[1,3]], - [[4,5],[6,7]]] [[4,6],[5,7]]] - - Step 3 — Unbind replicate dim (dim 0), giving 2 shard sub-meshes: - sub-mesh 0 = [[0,2],[1,3]] (replica group 0) - sub-mesh 1 = [[4,6],[5,7]] (replica group 1) - shard_placements = (_StridedShard(0), Shard(0)) - - Step 4 — Rank 0 → ProcessGroup([0,1,4,5]) - Rank 2 → ProcessGroup([2,3,6,7]) - - Returns: - ``(shard_mesh, process_group, shard_placements)`` - """ - my_rank = dist.get_rank() - assert mesh.mesh.device.type == 'cpu' - - # -- Fast path: 1D all-shard mesh → reuse existing PG. ---------------- - # Reuses the mesh's existing ProcessGroup directly, avoiding the - # overhead of dist.new_group(). The standard path below also handles - # subset calls safely via use_local_synchronization=True, but this - # fast path is still beneficial for the common 1D shard case. - if mesh.ndim == 1 and len(placements) == 1 and _is_shard(placements[0]): - key = (*mesh.mesh.shape, *mesh.mesh.flatten().tolist()) - if key not in _ranks_to_dist_cache: - _ranks_to_dist_cache[key] = (mesh, mesh.get_group()) - return (*_ranks_to_dist_cache[key], tuple(placements)) - - mesh_tensor = mesh.mesh.clone() - - # -- Step 1: Sort placements (Replicate first, then Shard by dim). ------ - # _StridedShard comes BEFORE regular Shard on the same dim so that - # get_slices_of_dtensor applies the outer sharding first, matching - # DTensor's left-to-right (outer-to-inner) composition order. - def _sort_key(item): - index, placement = item - assert not placement.is_partial(), "Partial placement not supported" - if placement.is_replicate(): - return (-1, 0, index) - assert _is_shard(placement), f"Unsupported: {type(placement)}" - split = (-1 / placement.split_factor if isinstance( - placement, _StridedShard) else 0) - return (placement.dim, split, index) - - indexed = sorted(enumerate(placements), key=_sort_key) - perm, sorted_placements = zip(*indexed) - - # -- Step 2: Permute mesh to match sorted placement order. -------------- - sorted_mesh = mesh_tensor.permute(perm) - - # -- Step 3: Collapse replicate dims → list of shard sub-meshes. -------- - # E.g. mesh (2, 3, 4, 4) with [R, R, S(0), S(1)] → 6 sub-meshes of (4, 4) - num_rep = sum(1 for p in sorted_placements if p.is_replicate()) - if num_rep > 0: - if num_rep > 1: - sorted_mesh = sorted_mesh.flatten(0, num_rep - 1) - shard_meshes = list(torch.unbind(sorted_mesh, dim=0)) - else: - shard_meshes = [sorted_mesh] - shard_placements = sorted_placements[num_rep:] - assert len(shard_placements) == len(set(shard_placements)) - - # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. -- - # Each rank only creates the group it belongs to, using - # use_local_synchronization=True so that only group members need to - # coordinate. This avoids deadlocks when different PP stages call - # construct_shard_mesh for different parameters. - def _cache_key(t: torch.Tensor) -> tuple: - return (*t.shape, *t.flatten().tolist()) - - my_key = None - for sm in shard_meshes: - if (my_rank == sm).any().item(): - key = _cache_key(sm) - assert my_key is None, "Rank appears in multiple shard groups" - my_key = key - if key not in _ranks_to_dist_cache: - pg = dist.new_group(sm.flatten().tolist(), - use_local_synchronization=True) - _ranks_to_dist_cache[key] = ( - DeviceMesh(device_type="cuda", mesh=sm), - pg, - ) - - return (*_ranks_to_dist_cache[my_key], shard_placements) diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/matmul_transpose_triton.py b/build/torch29-cxx11-rocm64-x86_64-linux/matmul_transpose_triton.py deleted file mode 100644 index 792de23d82c3fb45fe33d397ab9b76a0787259d0..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/matmul_transpose_triton.py +++ /dev/null @@ -1,122 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], - restore_value=['y'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -@torch.library.custom_op("muon::matmul_transpose_assign", - mutates_args=("d_out", )) -def matmul_transpose_assign(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """Compute d_out = d_in @ d_in.T using an optimized Triton kernel.""" - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -@matmul_transpose_assign.register_fake -def _(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """FakeTensor impl: d_out is already allocated, mutation is declared.""" - pass diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/metadata.json b/build/torch29-cxx11-rocm64-x86_64-linux/metadata.json deleted file mode 100644 index c55a35717622f1dd5c8ba376ea3a814cbcc10d78..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/metadata.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/muon.py b/build/torch29-cxx11-rocm64-x86_64-linux/muon.py deleted file mode 100644 index 14c0e22471fa6d47a51ed95e0e0c341dc18d5194..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/muon.py +++ /dev/null @@ -1,1068 +0,0 @@ -import logging -import types -from collections import defaultdict -from typing import Any - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Replicate, Shard -from torch.profiler import record_function - -from .adamw import _placement_cache, _tensor_cache, step_adamw -from .async_utils import run_pipeline -from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho, - get_default_muon_param_groups, is_expert_param, update_p) -from .cpu_offload import CPUOffloadPool -from .distributed.utils import (_is_shard, construct_shard_mesh, - get_slices_of_dtensor) -from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO, - _zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5_batched) -from .pipeline import muon_chunk_pipeline, prelaunch_first_gather -from .qk_clip import compute_scales, get_qk_clip_info, qk_clip - -logger = logging.getLogger(__name__) - - -def _expand_expert_params(names, params, expert_keys): - """Expand expert params by splitting on dim 0 (expert dimension). - - Params whose name matches any key in ``expert_keys`` are treated as - expert-parallel tensors. Their outermost dimension is the expert - dimension: an ``(E, out, in)`` tensor becomes ``E`` separate 2D - ``nn.Parameter`` views so that in-place updates propagate back to - the original storage. - - Non-expert params with ``ndim > 2`` trigger an ``AssertionError`` — - if they are expert params, their key must be added to ``expert_keys``. - - The grad must already be set on each expert param (e.g. after momentum). - - For DTensor expert params, placements that shard on dim 0 (expert dim) - are consumed by the split. Non-dim-0 shard placements (e.g. TP) are - preserved: each 2D slice is wrapped as a DTensor on the corresponding - submesh so the parallel pipeline handles the TP communication. - """ - expanded_names = [] - expanded_params = [] - - for n, p in zip(names, params): - is_expert = is_expert_param(n, expert_keys) - is_dtensor = isinstance(p.data, DTensor) - - if is_expert: - if is_dtensor: - logger.debug( - "[expand_expert] %s: expert DTensor, shape=%s, " - "placements=%s, mesh=%s, local_shape=%s", n, p.shape, - p.placements, p.device_mesh.mesh_dim_names, - p.to_local().shape) - else: - logger.debug( - "[expand_expert] %s: expert plain tensor, shape=%s", n, - p.data.shape) - - if not is_expert: - assert p.data.ndim <= 2, ( - f"Param {n} has ndim={p.data.ndim} but does not match " - f"expert_keys={expert_keys}. If this is an expert param, " - f"add its key to expert_keys.") - expanded_names.append(n) - expanded_params.append(p) - continue - - g = p.grad - assert g is not None, ( - f"Expert param {n} must have grad set before expansion") - - tp_mesh = None - tp_placements_2d = None - - if is_dtensor: - local_data = p.to_local() - local_grad = g.to_local() if isinstance(g, DTensor) else g - - # Find non-dim-0 shard placements (e.g. TP sharding). - # After splitting on dim 0, Shard(k) becomes Shard(k-1). - tp_dim_indices = [] - tp_placements_2d = [] - for i, pl in enumerate(p.placements): - if _is_shard(pl) and pl.dim != 0: - tp_dim_indices.append(i) - tp_placements_2d.append(Shard(pl.dim - 1)) - - if tp_dim_indices: - tp_dim_names = tuple(p.device_mesh.mesh_dim_names[i] - for i in tp_dim_indices) - if len(tp_dim_names) == 1: - tp_mesh = p.device_mesh[tp_dim_names[0]] - else: - tp_mesh = p.device_mesh[tp_dim_names] - else: - local_data = p.data - local_grad = g - - # Expand: split dim 0, reshape each slice to 2D. - num_local_experts = local_data.shape[0] - for i in range(num_local_experts): - slice_data = local_data[i] - slice_grad = local_grad[i] - - if tp_mesh is not None: - # Wrap as DTensor on TP submesh so the pipeline handles - # TP communication (gather/scatter across TP ranks). - dt_data = DTensor.from_local(slice_data, - device_mesh=tp_mesh, - placements=tp_placements_2d) - dt_grad = DTensor.from_local(slice_grad, - device_mesh=tp_mesh, - placements=tp_placements_2d) - expert_param = torch.nn.Parameter(dt_data, requires_grad=False) - expert_param.grad = dt_grad - else: - expert_param = torch.nn.Parameter(slice_data, - requires_grad=False) - expert_param.grad = slice_grad - - expanded_names.append(f"{n}[{i}]") - expanded_params.append(expert_param) - - p.grad = None # allow expert grad storage to be freed after pipeline - - return expanded_names, expanded_params - - -class Muon(torch.optim.Optimizer): - """ - Muon - MomentUm Orthogonalized by Newton-schulz - - Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- - processing step, in which each 2D parameter's update is replaced with the nearest orthogonal - matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has - the advantage that it can be stably run in bfloat16 on the GPU. - - Some warnings: - - We believe this optimizer is unlikely to work well for training with small batch size. - - We believe it may not work well for finetuning pretrained models, but we haven't tested this. - - Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. - lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) - momentum: The momentum used by the internal SGD. (0.95 is a good default) - nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) - ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. - Parameters that are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW instead. - adamw_lr: The learning rate for the internal AdamW. - adamw_betas: The betas for the internal AdamW. - adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - warmup_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher warmup_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. - chunk_size : Batch size of parameters to process in each - all2all gather/compute/scatter step. - Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified. - use_distributed_muon: Use distributed muon by Liu et al. (2024). - For testing purpose only. - expert_keys: List of strings to identify expert-parallel parameters. - If any key appears in a parameter's name, its outermost - dimension is treated as the expert dimension and expanded - into per-expert 2D params for Muon. For example, - ``expert_keys=["experts"]`` matches any param whose name - contains "experts". 3D+ params not matched by any key - will raise an error. - """ - - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config=None, - warmup_step=5, - chunk_size=-1, - use_distributed_muon=False, - expert_keys=None): - defaults = dict( - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - nesterov=nesterov, - ns_steps=ns_steps, - adamw_betas=adamw_betas, - adamw_eps=adamw_eps, - none_grad=none_grad, - use_muon=True, - ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) - super().__init__(params, defaults) - - self.debug = debug - self.clip_config = clip_config if clip_config is not None else { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100, - } - self.warmup_step = warmup_step - self.chunk_size = chunk_size - self.use_distributed_muon = use_distributed_muon - self.expert_keys = expert_keys - self.cpu_offload = False - self._cpu_offload_pool: CPUOffloadPool | None = None - self._offload_initialized = False - self._parallel_cache: dict[tuple[str, ...], dict] = {} - self._expert_expand_cache: dict[tuple[int, ...], dict] = {} - - def _calc_flops(self, G, steps): - assert len(G.shape) == 2 - M, N = G.shape - if M > N: - M, N = N, M - - return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) - - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - p.placements, p.device_mesh) - - return shard_mesh, shard_pg, shard_placements - - def init_state_and_assign_params(self, names, params, group, qk_logits): - param_to_state = {} - param_to_flops = {} - - total_flops = 0 - for p in params: - g = p.grad - if g is None: - continue - assert g.ndim == 2, "Muon only supports 2D parameters." - - flops = self._calc_flops(g, group["ns_steps"]) - param_to_flops[id(p)] = flops - total_flops += flops - - if self.debug: - logger.debug("Total TFLOPs for Muon: %.2f TFLOPs", - total_flops / 1e12) - - paired = list(zip(names, params)) - - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) - - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) - - round_robin = 0 - mesh = ordered_params[0].device_mesh - placements = ordered_params[0].placements - - shard_mesh, shard_pg, shard_placements = self.get_shard_mesh( - ordered_params[0]) - shard_mesh_flattened = shard_mesh.mesh.flatten() - num_ranks = dist.get_world_size(group=shard_pg) - - for n, p in zip(ordered_names, ordered_params): - if mesh != p.device_mesh: - raise ValueError("All parameters must be on the same mesh.") - if placements != p.placements: - raise ValueError("All parameters must have same placements.") - - worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks - round_robin = (round_robin + 1) % len(shard_mesh_flattened) - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - # Precompute per-rank indices and numels for all-to-all. - rank_indices: dict[int, tuple] = {} - rank_numels: dict[int, int] = {} - for r in range(num_ranks): - indices = get_slices_of_dtensor(p, r, shard_mesh, - shard_placements) - rank_indices[r] = indices - numel = 1 - for idx, dim_size in zip(indices, p.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) - rank_numels[r] = numel - - param_to_state[id(p)] = _muon_state( - worker_rank=worker_rank, - process_group=shard_pg, - rank_indices=rank_indices, - rank_numels=rank_numels, - name=n, - qk_clip_state=qk_clip_state, - ) - - return param_to_state, ordered_params - - def base(self, names, params, group, lr, weight_decay, qk_logits): - # Momentum is already applied by _step_muon before this method. - for n, p in zip(names, params): - g = p.grad - if g is None: - continue - - u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) - - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - def distributed_muon( - self, - names: list[str], - params: list[torch.nn.Parameter], - group: dict[str, Any], - lr: float, - weight_decay: float, - qk_logits: list[torch.Tensor | DTensor] | None, - ): - """Batched Distributed Muon — for testing/correctness verification only. - - Uses all-gather to reconstruct full tensors, computes Newton-Schulz on - the full grad, then slices back to local shards. This is simpler but - slower than the parallel pipeline (all2all) path, so it serves as a - reference implementation for verifying correctness. - """ - with record_function("distributed_muon"): - # Momentum is already applied by _step_muon before this method. - ns_steps = group["ns_steps"] - - # Separate plain tensors (no communication) from DTensors. - plain_names, plain_params = [], [] - dtensor_names, dtensor_params = [], [] - for n, p in zip(names, params): - if p.grad is None: - continue - if isinstance(p.data, DTensor): - dtensor_names.append(n) - dtensor_params.append(p) - else: - plain_names.append(n) - plain_params.append(p) - - # Process plain tensors per-param (no communication). - for n, p in zip(plain_names, plain_params): - u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE), - steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, - qk_logits) - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - if not dtensor_params: - return - - # Group DTensors by (placements, mesh) for batched all-gather. - placement_groups: dict[tuple, - tuple[list, - list]] = defaultdict(lambda: ([], [])) - for n, p in zip(dtensor_names, dtensor_params): - key = (p.placements, p.device_mesh) - placement_groups[key][0].append(n) - placement_groups[key][1].append(p) - - logger.info( - "distributed_muon: %d placement groups, %d total dtensors", - len(placement_groups), len(dtensor_params)) - - for (placements, mesh), (grp_names, - grp_params) in placement_groups.items(): - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - placements, mesh) - rank = dist.get_rank(shard_pg) - world_size = dist.get_world_size(shard_pg) - - logger.info(" group: %d params, placements=%s, world_size=%d", - len(grp_params), placements, world_size) - - # Separate params that can be batched (all shard dims evenly - # divisible) from those needing per-param full_tensor - # (e.g. MoE gate weights with fewer rows than shard ranks). - # all_gather_into_tensor requires equal buffer sizes across - # ranks, so uneven splits must use DTensor full_tensor(). - batch_names, batch_params = [], [] - single_names, single_params = [], [] - for n, p in zip(grp_names, grp_params): - even = all(p.shape[pl.dim] % - shard_mesh.mesh.shape[dim_idx] == 0 - for dim_idx, pl in enumerate(shard_placements)) - if even: - batch_names.append(n) - batch_params.append(p) - else: - single_names.append(n) - single_params.append(p) - - # Process uneven-split params per-param via full_tensor(). - for n, p in zip(single_names, single_params): - with record_function("distributed_muon::newton_schulz"): - g_full = p.grad.full_tensor().to(COMM_DTYPE) - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - if not batch_params: - continue - - logger.info(" batched=%d, single=%d", len(batch_params), - len(single_params)) - - # Concat all local grad shards into a single flat buffer. - with record_function("distributed_muon::gather"): - grad_locals = [ - p.grad.to_local().to(COMM_DTYPE).flatten() - for p in batch_params - ] - numels = [g.numel() for g in grad_locals] - grad_concat = torch.cat(grad_locals) - del grad_locals - - # Single all-gather (replaces N separate full_tensor). - grad_gathered = torch.empty( - grad_concat.numel() * world_size, - dtype=COMM_DTYPE, - device="cuda", - ) - dist.all_gather_into_tensor(grad_gathered, - grad_concat, - group=shard_pg) - - total_numel = grad_concat.numel() - del grad_concat - - # Precompute per-param offsets within the concat buffer. - offsets = [] - off = 0 - for ne in numels: - offsets.append(off) - off += ne - - # Per-param: reconstruct full grad → NS → local update. - for i, (n, p) in enumerate(zip(batch_names, batch_params)): - with record_function("distributed_muon::newton_schulz"): - g_full = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - for r in range(world_size): - r_start = r * total_numel + offsets[i] - shard = grad_gathered[r_start:r_start + numels[i]] - indices = get_slices_of_dtensor( - p, r, shard_mesh, shard_placements) - g_full[indices] = shard.reshape( - g_full[indices].shape) - - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - def _setup_parallel(self, names, params, group, qk_logits): - """Compute (or retrieve cached) parallel pipeline metadata. - - Returns: - (ordered_params, param_to_state, rank, chunk_size) - """ - cache_key = tuple(names) - - if cache_key not in self._parallel_cache: - # First call: compute metadata and populate cache. - param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) - - shard_pg = param_to_state[id(ordered_params[0])].process_group - rank = dist.get_rank(group=shard_pg) - - if self.chunk_size == -1: - shard_ranks = dist.get_world_size(shard_pg) - chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO - elif self.chunk_size > 0: - chunk_size = self.chunk_size - else: - raise ValueError( - "chunk_size must be -1 or a positive integer.") - - ordered_names = [ - param_to_state[id(p)].name for p in ordered_params - ] - name_to_state = { - param_to_state[id(p)].name: param_to_state[id(p)] - for p in ordered_params - } - self._parallel_cache[cache_key] = { - 'ordered_names': ordered_names, - 'name_to_state': name_to_state, - 'rank': rank, - 'chunk_size': chunk_size, - } - else: - # Cached path: rebuild param_to_state with current id(p) keys. - cache = self._parallel_cache[cache_key] - rank = cache['rank'] - chunk_size = cache['chunk_size'] - - name_to_param = dict(zip(names, params)) - ordered_params = [name_to_param[n] for n in cache['ordered_names']] - - param_to_state = {} - for p, n in zip(ordered_params, cache['ordered_names']): - cached_state = cache['name_to_state'][n] - param_to_state[id(p)] = _muon_state( - worker_rank=cached_state.worker_rank, - process_group=cached_state.process_group, - rank_indices=cached_state.rank_indices, - rank_numels=cached_state.rank_numels, - name=n, - qk_clip_state=get_qk_clip_info(self.clip_config, n, - qk_logits), - ) - - return ordered_params, param_to_state, rank, chunk_size - - def parallel(self, - names, - params, - group, - lr, - weight_decay, - qk_logits, - prelaunch_gather=None): - """ - Perform a parallel optimization step using Muon. - - Parameters are chunked and each chunk is processed by a - :func:`muon_chunk_pipeline` generator. :func:`run_pipeline` - interleaves multiple chunks so that communication and computation - overlap across chunks (the same overlap previously achieved by the - warmup + main-loop index scheduling). - - If ``prelaunch_gather`` is provided, it is passed to the first - chunk's generator to skip re-launching the already in-flight - A2A gather. - """ - - # Momentum is already applied by _step_muon before this method. - - ordered_params, param_to_state, rank, chunk_size = ( - self._setup_parallel(names, params, group, qk_logits)) - - def pipelines(): - first = True - for start in range(0, len(ordered_params), chunk_size): - chunk = ordered_params[start:start + chunk_size] - if chunk: - kwargs = dict( - params=chunk, - param_to_state=param_to_state, - rank=rank, - ns_steps=group["ns_steps"], - lr=lr, - weight_decay=weight_decay, - none_grad=group["none_grad"], - ) - if first and prelaunch_gather is not None: - kwargs['prelaunch_gather'] = prelaunch_gather - first = False - yield muon_chunk_pipeline(**kwargs) - - with record_function("muon::pipeline"): - run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1) - - def _step_muon(self, group, qk_logits=None): - params = group["params"] - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - # Apply momentum to all params before routing/expansion. - # Batched using _foreach_* ops (compiled, fullgraph=True). - with record_function("muon::momentum"): - active_params = [p for p in params if p.grad is not None] - if active_params: - # Ensure momentum buffers exist (avoid zeros_like when already present). - for p in active_params: - if "momentum_buffer" not in self.state[p]: - self.state[p]["momentum_buffer"] = torch.zeros_like( - p.grad) - - # Extract local tensors for compiled batch function. - local_grads = [ - p.grad._local_tensor - if isinstance(p.grad, DTensor) else p.grad - for p in active_params - ] - local_bufs = [ - self.state[p]["momentum_buffer"]._local_tensor - if isinstance(self.state[p]["momentum_buffer"], DTensor) - else self.state[p]["momentum_buffer"] - for p in active_params - ] - - # Wrap momentum as tensor for torch.compile. - batch_pre_ortho(local_grads, local_bufs, - torch.tensor(momentum), group["nesterov"]) - - # For non-nesterov, the result is the momentum buffer. - if not group["nesterov"]: - for p in active_params: - p.grad = self.state[p]["momentum_buffer"] - - # Identify batched experts for deferred NS. - # Detection is cheap (condition checks only); actual NS compute is - # deferred so it can overlap with the first chunk's A2A gather. - deferred_expert_work = [] - if self.expert_keys: - batched_expert_indices = [] - for i, (n, p) in enumerate(zip(names, params)): - if not (is_expert_param(n, self.expert_keys) - and p.grad is not None): - continue - # Eligible: plain tensor, or DTensor with no non-dim-0 shards. - if isinstance(p.data, DTensor): - has_tp = any( - _is_shard(pl) and pl.dim != 0 for pl in p.placements) - if has_tp: - continue - batched_expert_indices.append(i) - - if batched_expert_indices: - # Save refs for deferred NS; free grads from param list. - for i in batched_expert_indices: - p = params[i] - g = p.grad - local_g = (g._local_tensor - if isinstance(g, DTensor) else g) - local_data = (p.data._local_tensor if isinstance( - p.data, DTensor) else p.data) - deferred_expert_work.append((local_data, local_g)) - p.grad = None - - # Remove batched experts from lists before expansion. - keep = sorted( - set(range(len(params))) - set(batched_expert_indices)) - names = [names[i] for i in keep] - params = [params[i] for i in keep] - - def _run_deferred_expert_ns(): - """Execute deferred batched expert NS.""" - if not deferred_expert_work: - return - with record_function("muon::batched_expert_ns"): - ns_steps = group["ns_steps"] - for local_data, local_g in deferred_expert_work: - u = zeropower_via_newtonschulz5_batched( - local_g.to(COMM_DTYPE), steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:]) - local_data.mul_(1 - lr * weight_decay) - local_data.add_(u, alpha=-adjusted_lr) - - # Expand expert params by splitting on dim 0. - logger.debug("[_step_muon] before expand: %d params, expert_keys=%s", - len(params), self.expert_keys) - if self.expert_keys: - cache_key = tuple(id(p) for p in params) - cache = self._expert_expand_cache.get(cache_key) - - if cache is None: - # Cold path: full expansion + build cache metadata. - exp_names, exp_params = _expand_expert_params( - names, params, self.expert_keys) - - # Build per-expert-group info for hot-path grad updates. - grad_info = [] - exp_idx = 0 - for orig_idx, (n, p) in enumerate(zip(names, params)): - if not is_expert_param(n, self.expert_keys): - exp_idx += 1 - continue - - is_dt = isinstance(p.data, DTensor) - num_experts = (p.to_local() if is_dt else p.data).shape[0] - - # Detect TP mesh from the first expanded expert param. - tp_mesh = None - tp_pls = None - sample = exp_params[exp_idx] - if isinstance(sample.data, DTensor): - tp_mesh = sample.data.device_mesh - tp_pls = list(sample.data.placements) - - grad_info.append((orig_idx, num_experts, exp_idx, is_dt, - tp_mesh, tp_pls)) - exp_idx += num_experts - - self._expert_expand_cache[cache_key] = { - 'names': exp_names, - 'params': exp_params, - 'grad_info': grad_info, - } - names, params = exp_names, exp_params - else: - # Hot path: reuse cached params, only update expert grads. - for (orig_idx, num_experts, exp_start, is_dt, tp_mesh, - tp_pls) in cache['grad_info']: - p = params[orig_idx] - g = p.grad - local_grad = (g.to_local() - if is_dt and isinstance(g, DTensor) else g) - for i in range(num_experts): - expert_p = cache['params'][exp_start + i] - sg = local_grad[i] - if tp_mesh is not None: - expert_p.grad = DTensor.from_local( - sg, device_mesh=tp_mesh, placements=tp_pls) - else: - expert_p.grad = sg - p.grad = None - - names = cache['names'] - params = cache['params'] - else: - names, params = _expand_expert_params(names, params, - self.expert_keys) - logger.debug("[_step_muon] after expand: %d params", len(params)) - - param_dtensors = [] - name_dtensors = [] - - param_tensors = [] - name_tensors = [] - - # distributed_muon is a reference implementation for testing only. - # The parallel pipeline (all2all) path below is the production path. - if self.use_distributed_muon: - _run_deferred_expert_ns() - self.distributed_muon(names=names, - params=params, - group=group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits) - return - - for n, p in zip(names, params): - if p is None or p.grad is None: - continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - logger.debug( - "[route] %s → base (DTensor all-Replicate), " - "shape=%s, placements=%s", n, p.shape, p.placements) - param_tensors.append(p) - name_tensors.append(n) - else: - logger.debug( - "[route] %s → parallel (DTensor), shape=%s, " - "placements=%s, mesh=%s", n, p.shape, p.placements, - p.device_mesh.mesh_dim_names) - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - logger.debug("[route] %s → base (plain tensor), shape=%s", n, - p.data.shape) - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError(f"Unsupported parameter type: {type(p.data)}") - - logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, " - f"{len(param_tensors)} Tensors → base") - - def group_dtensors(dtensors, names): - # To support different placements, we group parameters by placements - # and run parallel Muon on each group. - - placement_to_params = defaultdict(lambda: ([], [])) - - assert len(dtensors) == len(names) - for p, n in zip(dtensors, names): - placement_to_params[tuple([p.placements, - p.device_mesh])][0].append(n) - placement_to_params[tuple([p.placements, - p.device_mesh])][1].append(p) - return placement_to_params - - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) - - dtensor_group = group_dtensors(param_dtensors, name_dtensors) - - # Pre-launch the first chunk's A2A gather so that the NCCL - # communication overlaps with the (deferred) batched expert NS - # compute on the default CUDA stream. - prelaunch = None - if deferred_expert_work: - first_names, first_params = next(iter(dtensor_group.values())) - ordered, pts, rnk, csz = self._setup_parallel( - first_names, first_params, group, qk_logits) - first_chunk = ordered[:csz] - if first_chunk: - prelaunch = prelaunch_first_gather(first_chunk, pts, rnk, - group["none_grad"]) - - _run_deferred_expert_ns() - - first_group = True - for _, (names, params) in dtensor_group.items(): - pg = prelaunch if first_group else None - first_group = False - self.parallel( - names, - params, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - prelaunch_gather=pg, - ) - else: - _run_deferred_expert_ns() - - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - ) - - def _register_states_for_offload(self): - """Register all optimizer state tensors with the CPU offload pool. - - Called once after the first step when states have been lazily created. - Offloads all param states (momentum buffers for Muon, moment1/moment2 - for AdamW) to free GPU memory between steps. - """ - pool = self._cpu_offload_pool - tracked = 0 - for group in self.param_groups: - for p in group["params"]: - if p not in self.state: - continue - state = self.state[p] - if group.get("use_muon", False): - if "momentum_buffer" in state: - pool.track(state["momentum_buffer"]) - tracked += 1 - else: - if "moment1" in state: - pool.track(state["moment1"]) - if "moment2" in state: - pool.track(state["moment2"]) - tracked += 1 - logger.info("[CPUOffload] Registered %d param states for offload", - tracked) - - @torch.no_grad - def step(self, closure=None, qk_logits=None): - """Perform a single optimization step. - - Args: - closure (Callable, optional): A closure that reevaluates the model - and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). - """ - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # H2D: reload optimizer states from CPU before computation. - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload() - - logger.debug("[Muon.step] expert_keys=%s, %d param groups", - self.expert_keys, len(self.param_groups)) - - for i, group in enumerate(self.param_groups): - if group["use_muon"]: - logger.debug("[Muon.step] group %d: use_muon=True, %d params", - i, len(group["params"])) - self._step_muon(group, qk_logits=qk_logits) - else: - logger.debug( - "[Muon.step] group %d: use_muon=False (AdamW), %d params", - i, len(group["params"])) - step_adamw(self.state, group) - - # D2H: offload optimizer states to CPU after computation. - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - return loss - - # ------------------------------------------------------------------ - # CPU offload public helpers - # ------------------------------------------------------------------ - - def turn_on_cpu_offload(self): - """Enable CPU offload for optimizer states.""" - if self.cpu_offload: - return - logger.info("[Muon] turn_on_cpu_offload") - self.cpu_offload = True - if not self.state: - return - self._cpu_offload_pool = CPUOffloadPool() - self._offload_initialized = False - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_off_cpu_offload(self): - """Disable CPU offload and keep optimizer states resident on GPU.""" - if not self.cpu_offload: - return - logger.info("[Muon] turn_off_cpu_offload") - if self._offload_initialized: - self._cpu_offload_pool.reload() - torch.cuda.current_stream().synchronize() - self._cpu_offload_pool = None - self._offload_initialized = False - self.cpu_offload = False - - # ------------------------------------------------------------------ - # Checkpoint support for cpu_offload - # ------------------------------------------------------------------ - - def state_dict(self) -> dict: - if self.cpu_offload: - raise RuntimeError( - "Muon.state_dict() requires turn_off_cpu_offload() before checkpoint save." - ) - return super().state_dict() - - def load_state_dict(self, state_dict: dict) -> None: - if self.cpu_offload: - raise RuntimeError( - "Muon.load_state_dict() requires turn_off_cpu_offload() before checkpoint load." - ) - super().load_state_dict(state_dict) - - # Invalidate adamw.py's module-level tensor caches so that - # the next step rebuilds them with the newly loaded state tensors. - _placement_cache.clear() - _tensor_cache.clear() diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/newton_schulz.py b/build/torch29-cxx11-rocm64-x86_64-linux/newton_schulz.py deleted file mode 100644 index d939264b69a34e7a3fa78859f34dc265a1159d59..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/newton_schulz.py +++ /dev/null @@ -1,240 +0,0 @@ -from itertools import repeat -from math import inf, sqrt - -import numpy as np -import torch - -from .matmul_transpose_triton import matmul_transpose_assign - -COMM_DTYPE = torch.bfloat16 -DEFAULT_CHUNK_SIZE_RATIO = 4 - - -def _optimal_quintic(l, u, max_iter=1000): - """ - Use the simplified Remez algorithm to find the optimal odd quintic approximant - to the constant function x -> 1 over the interval [l, u]. - - Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum - approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the - two interior equioscillation nodes q, r until convergence. Returns the - closed-form equioscillating solution when l ≈ u. - - Raises ValueError if any intermediate value (a, b, c, E, q, r) is non-finite - (NaN or inf). Raises RuntimeError if convergence is not reached within - max_iter iterations. - """ - assert 0 <= l <= u - if 1 - 5e-6 <= l / u: - return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5) - q = (3 * l + u) / 4 - r = (l + 3 * u) / 4 - E = inf - for _ in range(max_iter): - old_E = E - LHS = np.array( - [ - [l, l**3, l**5, 1], - [q, q**3, q**5, -1], - [r, r**3, r**5, 1], - [u, u**3, u**5, -1], - ] - ) - a, b, c, E = np.linalg.solve(LHS, np.ones(4)) - if not np.all(np.isfinite([a, b, c, E])): - raise ValueError( - f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}" - ) - q, r = np.sqrt( - (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / (10 * c) - ) - if not np.all(np.isfinite([q, r])): - raise ValueError(f"_optimal_quintic: non-finite node update q={q}, r={r}") - if abs(old_E - E) <= 1e-15: - break - else: - raise RuntimeError( - f"_optimal_quintic: did not converge after {max_iter} iterations" - ) - return float(a), float(b), float(c) - - -def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0): - """ - Compute the Polar Express coefficient series for `num_iters` quintic iterations. - - Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that - compose to map singular values from [l, 1] toward 1. At each step: - 1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion` - prevents near-zero singular values from stalling by raising the effective - lower bound; if it is active (cushion*u > l), the coefficients are - rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u]. - 2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the - last iteration, providing numerical headroom at the cost of a slightly slower - final convergence step. - 3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1). - - Returns a list of (a, b, c) tuples, one per iteration. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - u = 1 - assert 0 <= l <= u - safety_factor = 1 + safety_factor_eps - coefficients = [] - for iter in range(num_iters): - a, b, c = _optimal_quintic(max(l, cushion * u), u) - if cushion * u > l: - pl = a * l + b * l**3 + c * l**5 - pu = a * u + b * u**3 + c * u**5 - rescaler = 2 / (pl + pu) - a *= rescaler - b *= rescaler - c *= rescaler - if iter < num_iters - 1: - a /= safety_factor - b /= safety_factor**3 - c /= safety_factor**5 - coefficients.append((a, b, c)) - l = a * l + b * l**3 + c * l**5 - u = 2 - l - return coefficients - - -# Precomputed Polar Express coefficients (a, b, c) for 10 quintic Newton-Schulz -# iterations. Each tuple is the minimax-optimal (Remez/equioscillation) odd quintic -# approximant to x->1 over the current singular-value interval, computed once at -# import time and reused across all optimizer steps. -# -# Contrast with the former hardcoded NS coefficients (5 fixed tuples): -# - Former: empirically tuned to maximize slope at zero; did not converge -# singular values to 1, yielding US'V^T with S' ~ Uniform(0.5, 1.5) instead -# of the true polar factor UV^T. -# - Polar Express: analytically optimal per step, adapting to the shrinking -# singular-value interval [l, u] as iterations progress; converges all -# singular values to 1, producing the exact polar factor UV^T. -_coeffs_list = _optimal_composition( - l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02 -) - - -# This code is adapted from: -# KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py) -# NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress) -# matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon) -@torch.no_grad() -def _zeropower_via_newtonschulz5(G, steps): - """ - Compute the polar factor of G via the Polar Express method. - - Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c) - are the Polar Express coefficients from `_coeffs_list`. Each step is the - optimal odd quintic approximant to x -> 1 over the current singular-value - interval, minimizing the maximum approximation error (Remez / minimax criterion). - The composition maps singular values from [l, 1] to near 1, producing the - polar factor (orthogonal factor in the polar decomposition G = UP). - - `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2, - cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE - X = G # no manual typecast - - if G.size(0) > G.size(1): - X = X.T - - X = X / (X.norm() + 1e-7) - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - # Perform the NS iterations - for a, b, c in hs: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(0) > G.size(1): - X = X.T - - return X - - -@torch.no_grad() -def _zeropower_via_newtonschulz5_batched(G, steps): - """Batched polar factor computation for 3D (E, out, in) tensors. - - Same algorithm as ``_zeropower_via_newtonschulz5`` but uses - ``torch.bmm`` / ``torch.baddbmm`` instead of the 2D Triton kernel, - processing all E expert matrices in a single batched call. - """ - assert len(G.shape) == 3 - assert G.dtype == COMM_DTYPE - X = G - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - # Per-expert Frobenius norm. - X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) - - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list)) - ) - for a, b, c in hs: - buf1 = torch.bmm(X, X.transpose(-2, -1)) - buf2 = torch.bmm(buf1, buf1.transpose(-2, -1)) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.baddbmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - return X - - -_ns_per_shape: dict[tuple[int, ...], callable] = {} -_use_compile = True - - -def set_ns_compile(enabled: bool): - """Toggle torch.compile for Newton-Schulz iteration.""" - global _use_compile - _use_compile = enabled - - -def zeropower_via_newtonschulz5(G, steps=5): - if not _use_compile: - return _zeropower_via_newtonschulz5(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() - - -def zeropower_via_newtonschulz5_batched(G, steps=5): - """Compile-cached batched Newton-Schulz for 3D expert tensors.""" - if not _use_compile: - return _zeropower_via_newtonschulz5_batched(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile( - _zeropower_via_newtonschulz5_batched, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/optimizer/__init__.py b/build/torch29-cxx11-rocm64-x86_64-linux/optimizer/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/optimizer/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/pipeline.py b/build/torch29-cxx11-rocm64-x86_64-linux/pipeline.py deleted file mode 100644 index c0c2d515856182d8d15ad27dd4e4e093b29397d6..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/pipeline.py +++ /dev/null @@ -1,468 +0,0 @@ -import logging -from typing import Generator - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -from .core import _muon_state, adjust_lr_for_muon -from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5 -from .qk_clip import compute_scales - -logger = logging.getLogger(__name__) - -# ====================================================================== -# Stage helpers -# ====================================================================== - - -def _launch_gather( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Allocate gather buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_gather``). - gathered_grads: ``{id(p): empty_tensor}`` for owned params, - ``None`` for non-owned. - recv_counts: Per-source-rank element counts. - """ - # Allocate gathered-grad buffers - gathered_grads: dict[int, torch.Tensor | None] = {} - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - gathered_grads[id(p)] = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - else: - gathered_grads[id(p)] = None - - # Build send buffer – batch grad copies via torch.cat - # (1-2 fused kernels vs N individual narrow().copy_() calls). - send_counts = [0] * num_ranks - for p in params: - state = param_to_state[id(p)] - send_counts[state.worker_rank] += state.rank_numels[rank] - - total_send = sum(send_counts) - if total_send > 0: - # Group grad slices by destination rank in a single pass. - dst_to_grads = [[] for _ in range(num_ranks)] - for p in params: - state = param_to_state[id(p)] - n = state.rank_numels[rank] - if n > 0: - g = p.grad.to_local() - dst_to_grads[state.worker_rank].append(g.reshape(-1)) - - # Flatten in dst order and cat once. - all_slices = [] - for dst in range(num_ranks): - all_slices.extend(dst_to_grads[dst]) - send_buf = torch.cat(all_slices) - if send_buf.dtype != COMM_DTYPE: - send_buf = send_buf.to(COMM_DTYPE) - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += state.rank_numels[src] - recv_counts[src] = total - - recv_buf = torch.empty(sum(recv_counts), dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - logger.debug(f"send_buf size: {send_buf.numel()}, " - f"recv_buf size: {recv_buf.numel()}, " - f"recv_counts: {recv_counts}, " - f"send_counts: {send_counts}, " - f"process_group: {str(process_group)}") - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, gathered_grads, recv_counts - - -def _complete_gather( - recv_buf: torch.Tensor, - recv_counts: list[int], - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - param_to_state: dict[int, _muon_state], - rank: int, -) -> None: - """Reconstruct gathered grads from the recv buffer (in-place).""" - off = 0 - for src in range(len(recv_counts)): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - - indices = state.rank_indices[src] - - shard_view = gathered_grads[id(p)][indices] - n = shard_view.numel() - if n == 0: - continue - - sg = recv_buf.narrow(0, off + inner_off, n) - sg = sg.reshape(shard_view.shape) - gathered_grads[id(p)][indices] = sg - - inner_off += n - assert inner_off == block - off += block - - -def _compute_ns( - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - ns_steps: int, -) -> dict[int, torch.Tensor | None]: - """Run Newton-Schulz orthogonalization on owned parameters. - - Returns: - computed_us: ``{id(p): orthogonalized_update}`` for owned params. - """ - computed_us: dict[int, torch.Tensor | None] = {} - for p in owned_params: - u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps) - gathered_grads[id(p)] = None # free gathered grad - computed_us[id(p)] = u - return computed_us - - -def _launch_scatter( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, - computed_us: dict[int, torch.Tensor | None], -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor], list[int]]: - """Allocate scatter buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_scatter``). - scattered_us: Empty dict, populated by ``_complete_scatter`` with - zero-copy views into ``recv_buf``. - recv_counts: Per-source-rank element counts. - """ - # scattered_us is populated by _complete_scatter with zero-copy views - # into recv_buf, avoiding N empty_like allocations + N copy_ calls. - # Pre-seed entries for params whose local shard is empty (rank_numels == 0) - # so _update_params can iterate all params without KeyError. - scattered_us: dict[int, torch.Tensor] = {} - for p in params: - if param_to_state[id(p)].rank_numels[rank] == 0: - scattered_us[id(p)] = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - # Build send buffer – batch via torch.cat - # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls). - send_counts = [0] * num_ranks - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - for dst_rank in range(num_ranks): - send_counts[dst_rank] += state.rank_numels[dst_rank] - - total_send = sum(send_counts) - if total_send > 0: - # Cache u_full conversions to avoid redundant .to() per dst_rank. - u_fulls = {} - for p in owned_params: - u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous() - - # Collect slices in dst order (matches all-to-all send layout). - all_slices = [] - for dst_rank in range(num_ranks): - for p in owned_params: - state = param_to_state[id(p)] - su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten() - if su.numel() > 0: - all_slices.append(su) - - send_buf = torch.cat(all_slices) if all_slices else torch.empty( - 0, dtype=COMM_DTYPE, device="cuda") - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += state.rank_numels[rank] - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, scattered_us, recv_counts - - -def _complete_scatter( - recv_buf: torch.Tensor, - recv_counts: list[int], - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], -) -> None: - """Populate scattered_us with zero-copy views into recv_buf. - - Instead of pre-allocating tensors and copying, we assign views directly - from ``recv_buf``. This eliminates N ``empty_like`` + N ``copy_`` calls. - The underlying storage of ``recv_buf`` is kept alive through the views - until ``scattered_us`` is cleared after ``_update_params``. - """ - off = 0 - for src in range(len(recv_counts)): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = state.rank_numels[rank] - if n == 0: - continue - - scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - - inner_off += n - - assert inner_off == block - off += block - - -def _update_params( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], - lr: float, - weight_decay: float, -) -> None: - """Apply weight decay, Muon update, and optional QK clipping. - - Uses batched ``_foreach_mul_`` for weight decay and batched - ``_foreach_add_`` for the Muon update, grouping parameters by - adjusted_lr to minimize kernel launches while preserving float32 - precision for the alpha scaling. - """ - if not params: - return - - # Batched weight decay: p *= (1 - lr * wd) — single fused kernel. - p_locals = [p._local_tensor for p in params] - torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay) - - # Group params by adjusted_lr so _foreach_add_ can use a single - # alpha per group (preserves float32 precision for alpha scaling). - lr_groups: dict[float, tuple[list, list]] = {} - for p in params: - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - if adjusted_lr not in lr_groups: - lr_groups[adjusted_lr] = ([], []) - lr_groups[adjusted_lr][0].append(p._local_tensor) - lr_groups[adjusted_lr][1].append(scattered_us[id(p)]) - - for adjusted_lr, (p_group, u_group) in lr_groups.items(): - torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr) - - # QK clipping – applied directly on the local tensor to - # avoid DTensor sharding-propagation issues with _StridedShard. - for p in params: - state = param_to_state[id(p)] - if state.qk_clip_state is None: - continue - scales_full = compute_scales(p, state.qk_clip_state) - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = state.rank_indices[rank][0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - -# ====================================================================== -# Pre-launch helper for overlapping first chunk's gather with other work. -# ====================================================================== - - -@torch.no_grad() -def prelaunch_first_gather( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - none_grad: bool, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Launch the first chunk's A2A gather early for overlap with other compute. - - Call this *before* expensive GPU work (e.g. batched expert NS) so that - the NCCL all-to-all runs concurrently on the NCCL stream while the - default stream executes compute. - - Returns the same 4-tuple that ``_launch_gather`` produces, which should - be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - with record_function("muon::prelaunch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - return work, recv_buf, gathered_grads, recv_counts - - -# ====================================================================== -# Main generator – thin orchestrator that wires stages together. -# ====================================================================== - - -@torch.no_grad() -def muon_chunk_pipeline( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - ns_steps: int, - lr: float, - weight_decay: float, - none_grad: bool, - prelaunch_gather: tuple | None = None, -) -> Generator[None, None, None]: - """Process one chunk of parameters through the full Muon pipeline. - - Stages: gather -> compute (Newton-Schulz) -> scatter -> update. - - Each ``yield`` lets :func:`run_pipeline` interleave other chunks so - that communication and computation overlap across chunks. Async - communication is launched via ``async_op=True`` and completed after - the yield with ``work.wait()``. - - Overlap happens because :func:`run_pipeline` admits one new chunk - per iteration (staggered admission). While chunk *N* does NS - compute on the default CUDA stream, chunk *N+1*'s async all-to-all - runs concurrently on the NCCL stream — no separate ``comm_stream`` - is required. - - If ``prelaunch_gather`` is provided, the gather was already launched - by :func:`prelaunch_first_gather` and we skip launching it again. - - Yields exactly **2** times: - - 1. After launching async all-to-all gather (or immediately if pre-launched). - 2. After launching async all-to-all scatter. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - if prelaunch_gather is not None: - # Gather was pre-launched; none_grad already handled by caller. - work, recv_buf, gathered_grads, recv_counts = prelaunch_gather - else: - # Normal path: launch async gather. - with record_function("muon::launch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - yield # --- YIELD 1: other chunks can launch their gather --- - - with record_function("muon::wait_gather"): - work.wait() - _complete_gather(recv_buf, recv_counts, owned_params, gathered_grads, - param_to_state, rank) - del recv_buf - - # Stage 3: Newton-Schulz orthogonalization. - with record_function("muon::newton_schulz"): - computed_us = _compute_ns(owned_params, gathered_grads, ns_steps) - gathered_grads.clear() - - # Stages 4-5: launch async scatter. - with record_function("muon::launch_scatter"): - work, recv_buf, scattered_us, recv_counts = _launch_scatter( - params, owned_params, param_to_state, rank, num_ranks, - process_group, computed_us) - computed_us.clear() - - yield # --- YIELD 2: other chunks can launch their scatter --- - - with record_function("muon::wait_scatter"): - work.wait() - _complete_scatter(recv_buf, recv_counts, params, param_to_state, rank, - scattered_us) - del recv_buf - - # Stage 6: apply parameter updates. - with record_function("muon::update_params"): - _update_params(params, param_to_state, rank, scattered_us, lr, - weight_decay) - scattered_us.clear() diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/qk_clip.py b/build/torch29-cxx11-rocm64-x86_64-linux/qk_clip.py deleted file mode 100644 index 2aba711b3004b7f09e7141da7ef834bd61cc2430..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/qk_clip.py +++ /dev/null @@ -1,198 +0,0 @@ -import logging -import math -from dataclasses import dataclass - -import torch -from torch.distributed.tensor import DTensor - -from .core import normalize_fqn - -logger = logging.getLogger(__name__) - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - and return (kind, layer_index). - - Supported kinds: - MHA/GQA: 'wq', 'wk', 'q_proj', 'k_proj' - MLA: 'wq_b' (Q up-proj), 'wkv_b' (KV up-proj) - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.1.attn.wq_b.weight' -> ('wq_b', 1) - 'model.0.attn.wkv_b.weight' -> ('wkv_b', 0) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = normalize_fqn(name).split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: str | None # 'wq'/'q_proj'/'wq_b' or 'wk'/'k_proj'/'wkv_b' or None - indices: list[int] # which heads to consider for clipping - head_dim: int # from config (qk_head_dim for MLA wq_b) - threshold: float # from config - logit: torch.Tensor | None - - # MLA-specific fields - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 - - -def get_qk_clip_info(clip_config, n, qk_logits): - """Extract QK clipping info for a named parameter. - - Args: - clip_config: QK clipping configuration dict (or None). - MHA/GQA keys: head_dim, threshold, q_indices, k_indices - MLA extra keys: is_mla=True, qk_nope_head_dim, qk_rope_head_dim, v_head_dim - n: Parameter name string. - qk_logits: Dict mapping layer indices to logit tensors (or None). - - Returns: - QKClipInfo instance with clipping configuration for this parameter. - """ - if clip_config is None: - return None - - head_dim = clip_config.get('head_dim') - threshold = clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - # In TP settings, qk_logits may be DTensor - # We convert it to full tensor here for simplicity - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - -def compute_scales(p, qk_clip_state): - """Compute per-head scaling factors for QK clipping. - - Returns scales tensor (√γ per head) if any head exceeds threshold, else None. - For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim. - """ - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - # Check if any head exceeds threshold before allocating. - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - - if not head_scales: - return None - - # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows - if qk_clip_state.is_mla and kind == 'wkv_b': - effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim - else: - effective_head_dim = head_dim - - H_global = p.shape[0] // effective_head_dim - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale - return scales_full - - -def qk_clip(p, scales, info): - """Apply per-head scaling to a Q/K projection weight matrix. - - Args: - p: Parameter (nn.Parameter or raw tensor). - scales: [n_heads] tensor, each element = √γ_h. - info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions. - - MLA sub-region scaling per Algorithm 1 (MuonClip): - wq_b: q_nope rows → √γ, q_pe rows → γ - wkv_b: k_nope rows → √γ, v rows → unchanged - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not info.is_mla: - # MHA/GQA: uniform √γ applied to all rows in each head - W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: vectorized sub-region scaling within each head - if info.kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_head_dim = qk_nope + info.qk_rope_head_dim - W_3d = W.view(-1, qk_head_dim, W.shape[1]) # [H, qk_head_dim, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # q_nope → √γ - W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1, - 1)) # q_pe → γ - - elif info.kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - W_3d = W.view(-1, kv_stride, W.shape[1]) # [H, kv_stride, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # k_nope → √γ - # v rows: not touched (k_R shared rotary unchanged) diff --git a/docs/expert_parallel.md b/docs/expert_parallel.md deleted file mode 100644 index c037a998bc44ffea02503dacf41f7124508a1612..0000000000000000000000000000000000000000 --- a/docs/expert_parallel.md +++ /dev/null @@ -1,264 +0,0 @@ -# Expert Parallelism in torchtitan - -torchtitan (0.2.0)의 expert parallelism 구현을 정리한 문서. -Muon optimizer의 MoE 지원에 필요한 배경 지식. - -Reference: `torchtitan/distributed/expert_parallel.py`, `torchtitan/distributed/parallel_dims.py` - -## Overview - -torchtitan은 MoE expert weights에 대해 4가지 parallelism 전략을 제공: - -| Config | TP | EP | ETP | Expert Weight Placements | Token Dispatch | -|--------|----|----|-----|--------------------------|----------------| -| TP Only | >1 | 1 | - | `[Shard(1/2)]` on TP mesh | None | -| EP Only | 1 | >1 | - | `[Shard(0)]` on EP mesh | All-to-all | -| EP+ETP (etp=tp) | >1 | >1 | =tp | `[Shard(0), Shard(1/2)]` on [EP, TP] mesh | All-to-all on EP | -| EP+ETP (etp=1) | >1 | >1 | 1 | `[Shard(0)]` on EP mesh | Sequence parallel on TP | - -Expert weights shape: `(num_experts, out_dim, in_dim)` (w1, w3) / `(num_experts, in_dim, out_dim)` (w2). - -## EP가 dp_shard를 빌리는 구조 - -EP는 새로운 물리적 차원이 아니라 `dp_shard`를 분해해서 사용: - -``` -dp_shard = dp_shard_mod_ep * dp_shard_in_ep - -ETP=TP일 때: ep = dp_shard_in_ep * cp -ETP=1일 때: ep = dp_shard_in_ep * cp * tp -``` - -기존 mesh `[pp, dp_replicate, dp_shard, cp, tp]`가 EP 활성화 시: - -``` -[pp, dp_replicate, dp_shard_mod_ep, dp_shard_in_ep, cp, tp] -``` - -로 확장됨. `dp_shard_mod_ep`는 값이 1이어도 mesh에 유지 (FSDP wrapping 일관성). - -### 예시: 8 GPUs, ep=4, dp_shard=8, tp=1, cp=1 - -``` -dp_shard_in_ep = ep / cp = 4 -dp_shard_mod_ep = dp_shard * cp / ep = 2 - -mesh: [dp_shard_mod_ep=2, dp_shard_in_ep=4] -EP mesh: [dp_shard_in_ep=4] → expert들을 4-way로 분배 -FSDP mesh: [dp_shard_mod_ep=2] → expert FSDP는 2-way로 shard -``` - -## Submesh 매핑 - -```python -# Data loading (no communication) -dp = [dp_replicate, dp_shard_mod_ep, dp_shard_in_ep] - -# Non-expert parameter sharding (FSDP) -dp_shard_cp = [dp_shard_mod_ep, dp_shard_in_ep, cp] - -# Expert parameter sharding (EFSDP) — dp_shard_in_ep 제외 -dp_mod_ep = [dp_replicate?, dp_shard_mod_ep] - -# Expert parallelism mesh -ep = [dp_shard_in_ep, cp, (tp if etp==1)] - -# Loss all-reduce -dp_cp = [dp_replicate, dp_shard_mod_ep, dp_shard_in_ep, cp] -``` - -## 4가지 전략 상세 - -### 1. TensorParallel (TP Only, EP=1) - -EP 없이 TP만 사용. Expert weights를 TP mesh에서 column/row-wise sharding: - -```python -# expert_parallel.py: TensorParallel -w1: [Shard(1)] on TP mesh # column-wise (out_dim) -w2: [Shard(2)] on TP mesh # row-wise (out_dim, 3D에서 dim 2) -w3: [Shard(1)] on TP mesh # column-wise (out_dim) -``` - -Token dispatch 없음. 일반 TP와 동일하게 동작. - -### 2. ExpertParallel (EP Only, TP=1) - -Expert dim (dim 0)으로 sharding. Token all-to-all dispatch: - -```python -# expert_parallel.py: ExpertParallel -w1, w2, w3: [Shard(0)] on EP mesh # expert dim으로 분배 -``` - -Forward pass: -1. Router가 각 token을 expert에 할당 -2. `all_to_all_single`으로 token을 해당 expert의 rank로 dispatch -3. 각 rank가 local expert에서 compute -4. `all_to_all_single`으로 결과를 원래 rank로 combine - -### 3. ExpertTensorParallel (EP+TP, ETP=TP) - -EP와 TP를 동시에 2D로 적용: - -```python -# expert_parallel.py: ExpertTensorParallel (extends ExpertParallel) -w1: [Shard(0), Shard(1)] on [EP, TP] mesh # expert + column -w2: [Shard(0), Shard(2)] on [EP, TP] mesh # expert + row -w3: [Shard(0), Shard(1)] on [EP, TP] mesh # expert + column -``` - -Token dispatch: -1. TP mesh에서 input을 Replicate (gradient는 Partial) -2. EP mesh에서 all-to-all dispatch (ExpertParallel과 동일) -3. All-to-all은 EP mesh에서만 발생, TP 통신은 weight sharding으로 처리 - -### 4. ReordererSequenceParallel (EP+TP, ETP=1) - -TP hardware를 EP에 빌려줌. TP mesh가 sequence parallel로 동작: - -```python -# expert_parallel.py: ReordererSequenceParallel -# Expert weights: [Shard(0)] on EP mesh (TP 안 씀) -# Token split: batch*seq_len을 TP rank 수로 나눠서 분배 - -# EP mesh = [dp_shard_in_ep, cp, tp] ← tp가 EP에 포함됨 -``` - -TP rank들이 token을 나눠 처리 (sequence parallel). Expert weight에는 TP sharding 없음. - -## EFSDP (Expert FSDP) - -Expert parameter에 대한 FSDP는 non-expert parameter와 **다른 mesh**를 사용: - -```python -# parallelize.py: apply_fsdp -# Non-expert: dp_shard_cp mesh 전체로 shard -fully_shard(transformer_block, mesh=dp_shard_cp_mesh) - -# Expert (EP 활성화 시): dp_mod_ep mesh로만 shard -# dp_shard_in_ep는 이미 EP에서 사용 중이므로 제외 -fully_shard(transformer_block.moe.experts, mesh=dp_mod_ep_mesh) -``` - -### Dynamic shard placement - -Expert 수보다 `dp_mod_ep * ep`가 클 때 (expert dim으로 더 쪼갤 수 없을 때), -dim 0 대신 dim 1로 shard. - -**torchtitan 코드** (`torchtitan/models/llama4/infra/parallelize.py:339-359`): - -```python -# NOTE: EP alreadys shards the routed experts on dim 0 (num_experts). -# When dp_mod_ep * ep > num_experts, FSDP default dim-0 sharding -# causes inefficiency, so we choose to do FSDP sharding on dim-1. -_experts_shard_placement_fn = None -if ( - dp_mod_ep_mesh.size() * ep_degree - > transformer_block.moe.experts.num_experts -): - _experts_shard_placement_fn = lambda param: Shard(1) - -fully_shard( - transformer_block.moe.experts, - **fsdp_mod_ep_config, # mesh=dp_mod_ep_mesh - reshard_after_forward=reshard_after_forward, - shard_placement_fn=_experts_shard_placement_fn, -) -``` - -`dp_mod_ep_mesh` 구성 (`parallelize.py:140-159`): - -```python -dp_mod_ep_mesh_dim_names = [] -if parallel_dims.ep_enabled: - if parallel_dims.dp_replicate_enabled: - dp_mod_ep_mesh_dim_names.append("dp_replicate") - dp_mod_ep_mesh_dim_names.append("dp_shard_mod_ep") -# → dp_mod_ep_mesh = world_mesh[tuple(dp_mod_ep_mesh_dim_names)] -``` - -### 실제 placement 검증 결과 - -8 GPUs, `num_experts=2`, `etp=1` 기준: - -#### num_experts=8 (기본) - -모든 config에서 expert weights는 **dim 0 (expert dim)으로만 shard**: - -| Config | Expert Placements | Mesh | -|--------|-------------------|------| -| ep=8 | `[Shard(0)]` | `[ep=8]` | -| ep=4, fsdp=2 | `[_StridedShard(0), Shard(0)]` | `[dp_shard_mod_ep=2, ep=4]` | -| ep=2, fsdp=4 | `[_StridedShard(0), Shard(0)]` | `[dp_shard_mod_ep=4, ep=2]` | -| ep=2, hsdp=2+2 | `[Replicate(), _StridedShard(0), Shard(0)]` | `[dp_rep=2, dp_shard_mod_ep=2, ep=2]` | - -EFSDP는 `_StridedShard(dim=0)`, EP는 `Shard(dim=0)`. 비-dim-0 shard 없음. - -#### num_experts=2 (expert 수 < EFSDP shard count) - -`dp_mod_ep * ep > num_experts` 조건 충족 시 **EFSDP가 Shard(1)로 전환**: - -| Config | 조건 | Expert Placements | Mesh | -|--------|------|-------------------|------| -| ep=2, fsdp=4 | 4*2=8 > 2 | `[Shard(1), Shard(0)]` | `[dp_shard_mod_ep=4, ep=2]` | -| ep=2, hsdp=2+2 | 2*2=4 > 2 | `[Replicate(), Shard(1), Shard(0)]` | `[dp_rep=2, dp_shard_mod_ep=2, ep=2]` | - -- EFSDP: `Shard(1)` on `dp_shard_mod_ep` → out_dim을 shard (w1: 2816/4=704) -- EP: `Shard(0)` on `ep` → expert dim을 shard (2/2=1) -- `_StridedShard`가 아닌 일반 `Shard` 사용 - -## Gradient Clipping with EP - -EP parameter와 non-EP parameter의 gradient norm을 별도로 계산 후 합산: - -```python -# distributed/utils.py: _clip_grad_norm_with_ep -ep_norm = get_total_norm(ep_grads, ...) -non_ep_norm = get_total_norm(non_ep_grads, ...) -total_norm = (ep_norm**p + non_ep_norm**p) ** (1/p) -``` - -EP parameter 판별: `device_mesh.mesh_dim_names`에 "ep" 포함 여부. - -## Muon optimizer에서의 처리 - -현재 Muon optimizer의 MoE 지원: - -1. **`_expand_expert_params`**: 3D expert weight를 expert dim (dim 0)으로 split하여 2D param으로 확장 -2. **TP가 있을 때**: non-dim-0 shard (TP)를 TP submesh에 DTensor로 wrap - - 3D `(Shard(0), Shard(1))` → 2D `(Shard(0),)` on TP submesh -3. **`construct_shard_mesh` fast path**: 1D submesh에서 `dist.new_group()` deadlock 방지 - -### Muon이 지원하는 config - -| Config | 지원 | 비고 | -|--------|------|------| -| TP Only (EP=1) | O | expert를 TP submesh DTensor로 처리 | -| EP Only (TP=1) | O | expert를 plain tensor로 처리 (base mode) | -| FSDP + TP | O | FSDP는 expert dim, TP는 out/in dim | -| HSDP + TP | O | Replicate + FSDP + TP | -| EP Only (많은 experts) | O | EFSDP `Shard(0)` → plain tensor | -| EP + FSDP (적은 experts) | 미테스트 | EFSDP `Shard(1)` → 아래 참조 | -| EP + TP (ETP=TP) | 미테스트 | 2D expert DTensor `[Shard(0), Shard(1/2)]` | -| EP + TP (ETP=1) | 미테스트 | EP mesh에 TP가 포함된 경우 | - -### EFSDP Shard(1)과 Muon의 호환성 - -Muon은 placement-agnostic. `_expand_expert_params`의 non-dim-0 shard 처리 로직이 -TP뿐 아니라 EFSDP `Shard(1)`에도 동일하게 적용됨 (변수명만 `tp_*`일 뿐 로직은 generic): - -``` -3D: (Shard(1), Shard(0)) on [dp_shard_mod_ep=4, ep=2] - local shape: (1, 704, 2048) - -_expand_expert_params: - 1. non-dim-0 shard 탐색 → Shard(1) on dp_shard_mod_ep - 2. submesh 추출 → dp_shard_mod_ep (1D, size 4) - 3. dim 0 split → (704, 2048) - 4. DTensor wrap → Shard(0) on dp_shard_mod_ep - = 일반 FSDP sharded 2D 텐서와 동일 - -→ parallel()/distributed_muon()이 all-gather → Newton-Schulz → scatter 처리. - construct_shard_mesh fast path 적용 (1D submesh, deadlock 없음). -``` diff --git a/docs/implementation.md b/docs/implementation.md deleted file mode 100644 index 0b5c88ea4d316257f9ff303f7d36194a7288e99c..0000000000000000000000000000000000000000 --- a/docs/implementation.md +++ /dev/null @@ -1,321 +0,0 @@ -# Muon Optimizer: Implementation Guide - -This document explains the internal architecture of the Muon optimizer for reviewers and new contributors. It covers the execution paths, the parallel pipeline design, and the distributed sharding utilities. - -## Table of Contents - -1. [Overview](#overview) -2. [Entry Point and Parameter Routing](#entry-point-and-parameter-routing) -3. [Execution Paths](#execution-paths) -4. [Parallel Pipeline (the core feature)](#parallel-pipeline) -5. [MoE Expert Weight Support](#moe-expert-weight-support-expert_keys) -6. [Distributed Utilities](#distributed-utilities) -7. [Newton-Schulz Orthogonalization](#newton-schulz-orthogonalization) -8. [QK Clipping](#qk-clipping) -9. [AdamW for Non-Muon Parameters](#adamw-for-non-muon-parameters) -10. [Source File Map](#source-file-map) - ---- - -## Overview - -Muon (MomentUm Orthogonalized by Newton-schulz) applies standard SGD-momentum and then replaces each 2D parameter's update with the nearest orthogonal matrix via a Newton-Schulz iteration. The iteration runs stably in bfloat16 on GPU. - -The optimizer supports arbitrary N-D sharding configurations: FSDP2, TP, or hybrid setups like `2 TP x 2 DP-Replicate x 2 DP-Shard`. This generality is what drives most of the code complexity. - -## Entry Point and Parameter Routing - -**File:** `muon.py` — `Muon.step()` / `Muon._step_muon()` - -Users must provide parameter groups with `use_muon=True/False` flags (via `get_default_muon_param_groups()`). At each step: - -1. **Non-Muon groups** → `step_adamw()` (fused AdamW). -2. **Muon groups** → `_step_muon()`, which further classifies each parameter: - -``` -_step_muon(group) - | - +-- momentum update (batched _foreach_* ops) - +-- _expand_expert_params() -- 3D expert params → per-expert 2D views (cached) - | - +-- DTensor, all Replicate placements --> base() (no sharding) - +-- DTensor, sharded --> parallel() (pipelined all-to-all) - +-- plain Tensor --> base() (single device) -``` - -Parameters are classified by their DTensor placements: -- **Fully replicated** DTensors and plain tensors use `base()` — standard single-device Muon. -- **Sharded** DTensors use `parallel()` — the pipelined all-to-all approach described below. -- `distributed_muon()` exists as a **test-only reference implementation** for correctness verification. - -## Execution Paths - -### base() — Single Device - -Straightforward per-parameter loop: momentum update → Newton-Schulz orthogonalization → parameter update → optional QK clipping. - -### distributed_muon() — Full Gather (test-only) - -Reference implementation for correctness verification. Uses batched all-gather to reconstruct full tensors, computes Newton-Schulz on the full grad, then slices back to local shards. Simple but communication-heavy — not used in production. - -### parallel() — Pipelined All-to-All - -This is the main advanced feature. Instead of all-gathering the full parameter, it uses **all-to-all** to distribute work: each rank "owns" a subset of parameters and is responsible for their Newton-Schulz computation. - -## Parallel Pipeline - -### Design Motivation - -Newton-Schulz is compute-intensive. The key insight is that each rank only needs to orthogonalize the parameters it "owns" — not all parameters. So the flow is: - -1. **Gather**: Each rank sends its local gradient shard to the owning rank via all-to-all. -2. **Compute**: The owning rank runs Newton-Schulz on the full (gathered) gradient. -3. **Scatter**: The owning rank sends the orthogonalized update back to all ranks via all-to-all. -4. **Update**: Each rank applies weight decay and the update to its local shard. - -To overlap communication and computation, parameters are split into **chunks**, and multiple chunks are processed concurrently. - -### Architecture - -``` -muon.py: parallel() - | - +-- init_state_and_assign_params() -- assigns ownership, precomputes indices - | - +-- pipelines() generator -- yields muon_chunk_pipeline() per chunk - | - +-- run_pipeline(pipelines, max_concurrent=warmup_step+1) - | - +-- interleaves chunks at yield boundaries -``` - -### The Chunk Pipeline Generator - -**File:** `pipeline.py` — `muon_chunk_pipeline()` - -Each chunk is a generator that yields **2 times**, creating stages separated by async communication: - -``` - YIELD 1 YIELD 2 - | | -[Build bufs + async gather a2a] --> [wait + NS compute + async scatter a2a] --> [wait + Update params] -``` - -- **Async communication**: `dist.all_to_all_single(..., async_op=True)` launches non-blocking communication. The generator yields immediately after, allowing other chunks to run. `work.wait()` completes the operation after the yield. -- **Chunk-level overlap**: `run_pipeline()` interleaves multiple chunks at yield boundaries, so while chunk N waits for its communication, chunk N+1 can launch its own. - -### The Pipeline Scheduler - -**File:** `async_utils.py` — `run_pipeline()` - -A simple round-robin scheduler: - -```python -while have_new or previous_tasks: - # Admit one new pipeline if below concurrency limit - if have_new and len(previous_tasks) < max_concurrent: - task = next(pipelines) # runs to first yield - # Advance all existing tasks by one yield - for task in previous_tasks: - task.step() # runs to next yield -``` - -`max_concurrent = warmup_step + 1` controls how many chunks can be in-flight simultaneously. Higher values increase memory usage but improve communication/computation overlap. - -### Ownership Assignment - -**File:** `muon.py` — `init_state_and_assign_params()` - -Parameters are sorted by FLOP cost (descending) and assigned to ranks in round-robin order across the shard mesh. This balances compute load across ranks. - -### Precomputed Shard Indices - -Instead of computing per-rank shard indices on every step, they are precomputed once during `init_state_and_assign_params()` and stored in `_muon_state`: - -```python -@dataclass -class _muon_state: - worker_rank: int # which rank owns this param's computation - process_group: ProcessGroup # the all-to-all communication group - rank_indices: dict[int, tuple] # rank -> per-dim indices into full tensor - rank_numels: dict[int, int] # rank -> number of elements in shard - name: str - qk_clip_state: QKClipInfo | None -``` - -`rank_indices[r]` is a tuple of `slice` or `torch.Tensor` per dimension, describing which elements of the full tensor rank `r` owns. `rank_numels[r]` is the total number of elements in that shard. These are used directly in the pipeline's gather and scatter stages. - -### Pipeline Stages in Detail - -#### Stages 1-2: Gather - -1. **Allocate** receive buffers for gathered gradients (only on owning ranks). -2. **Build send buffer**: Each rank flattens its local gradient shard for each destination rank. -3. **Async all-to-all**: `dist.all_to_all_single(..., async_op=True)` launches gather. -4. **Yield 1**: Other chunks can launch their gather while this one waits. -5. **`work.wait()`**: Complete the gather. -6. **Reconstruct**: The owning rank places received shards into the full gradient using `rank_indices`. - -#### Stage 3: Compute - -The owning rank runs `_zeropower_via_newtonschulz5()` on the full gathered gradient. This is the most compute-intensive stage. Runs inline (no yield) since it is synchronous GPU work. - -#### Stages 4-5: Scatter - -Inverse of gather: -1. **Allocate** receive buffers for the orthogonalized update `U`. -2. **Build send buffer**: The owning rank slices `U` using `rank_indices` for each destination rank. -3. **Async all-to-all**: `dist.all_to_all_single(..., async_op=True)` launches scatter. -4. **Yield 2**: Other chunks can launch their scatter while this one waits. -5. **`work.wait()`**: Complete the scatter. -6. **Copy** received shards into local update buffers. - -#### Stage 6: Update - -Each rank applies weight decay and the Muon update to its local parameter shard. Also applies QK clipping if configured. - -## MoE Expert Weight Support (`expert_keys`) - -**File:** `muon.py` — `_expand_expert_params()` - -MoE models have 3D expert weights with shape `(num_experts, out_dim, in_dim)`. Since Muon operates on 2D matrices, expert params need special handling. - -### Configuration - -Pass `expert_keys` to both `get_default_muon_param_groups()` and `Muon()`: - -```python -params = get_default_muon_param_groups(model, expert_keys=["experts"]) -optim = Muon(params, expert_keys=["experts"], ...) -``` - -Any parameter whose name contains a string in `expert_keys` is treated as an expert-parallel parameter. Non-matching 3D+ parameters raise `AssertionError` to catch misconfiguration. - -### How It Works - -`_expand_expert_params()` runs after momentum and before routing to `base()`/`parallel()`/`distributed_muon()`: - -1. **Split on dim 0**: A 3D `(E, out, in)` tensor becomes `E` separate 2D `(out, in)` `nn.Parameter` views. Views share storage with the original, so in-place updates propagate back. -2. **Placement remapping**: When the original is a DTensor, `Shard(k)` on dim `k > 0` becomes `Shard(k-1)` on the 2D slice (since dim 0 is consumed by the split). -3. **Submesh wrapping**: Non-dim-0 shard placements are preserved by wrapping each 2D slice as a DTensor on the corresponding submesh. This is **placement-agnostic** — the same logic handles TP `Shard(1/2)`, EFSDP `Shard(1)`, or any other non-dim-0 sharding. - -### Placement-Agnostic Design - -The expansion logic does not care *why* a dimension is sharded — only whether it's on dim 0 (consumed by split) or not (preserved on submesh): - -| Original Placement | After Expansion | -|-------------------|-----------------| -| `Shard(0)` (EP) | Consumed by split → plain tensor | -| `Shard(1)` (TP or EFSDP) | `Shard(0)` on submesh → 2D DTensor | -| `Shard(2)` (TP row-wise) | `Shard(1)` on submesh → 2D DTensor | -| `Replicate` | Ignored (not a shard) | -| `_StridedShard(0)` (EFSDP) | Consumed by split → plain tensor | - -After expansion, the 2D params flow through the standard routing: DTensors with shard placements go to `parallel()`, plain tensors go to `base()`. - -For EP/EFSDP background and torchtitan integration details, see [`docs/expert_parallel.md`](expert_parallel.md). - -## Distributed Utilities - -**File:** `distributed/utils.py` - -These utilities solve the problem of mapping from a DTensor's arbitrary sharding configuration to the concrete indices each rank owns. - -### `construct_shard_mesh(placements, mesh)` - -Given a DTensor's placements and device mesh, this function: - -1. **Sorts** placements: Replicate dims first, then Shard dims by dimension (with `_StridedShard` before regular `Shard` on the same dim, so the outer sharding is applied first). -2. **Permutes** the mesh accordingly. -3. **Separates** replicate dims from shard dims — each replicate group gets its own shard sub-mesh. -4. **Creates** a ProcessGroup for the current rank's shard mesh. - -Returns `(shard_mesh, process_group, shard_placements)` — used for all-to-all communication. - -**Why this is needed:** A model might use `[Replicate, Shard(0), _StridedShard(0)]` across a 3D mesh. The optimizer needs to identify which ranks participate in the same shard group (share the same data) and create a ProcessGroup for them. - -### `get_slices_of_dtensor(target, local_rank, shard_mesh, shard_placements)` - -Computes the exact indices that a given rank owns in the full tensor. Handles both contiguous (`Shard`) and strided (`_StridedShard`) sharding, including composed multi-level sharding on the same dimension. - -Returns a tuple of `slice` (contiguous) or `torch.LongTensor` (strided) per dimension. - -**Example:** With `[Shard(0), _StridedShard(0)]` on a (16, 2048) tensor across 4 ranks: -- Rank 0 might own rows `[0, 4, 8, 12]` (strided) -- Rank 1 might own rows `[1, 5, 9, 13]` -- etc. - -### PyTorch 2.10 Compatibility - -In PyTorch 2.10, `_StridedShard` no longer inherits from `Shard`. The helper `_is_shard()` handles both old and new hierarchies: - -```python -def _is_shard(placement): - return isinstance(placement, (Shard, _StridedShard)) -``` - -## Newton-Schulz Orthogonalization - -**File:** `newton_schulz.py` - -`_zeropower_via_newtonschulz5()` computes the polar factor of a matrix using the Polar Express method — quintic Newton-Schulz iterations with analytically optimal (minimax/Remez) coefficients precomputed by `_optimal_composition()`. The default configuration uses 10 iterations with `l=1e-3`, converging all singular values to 1 to produce the exact polar factor `UV^T`. Wrapped by `zeropower_via_newtonschulz5()` which adds per-shape `torch.compile` caching with CUDA graph support. - -Each iteration uses `matmul_transpose_assign()` (a Triton kernel for `X @ X^T`) for efficiency. - -**File:** `matmul_transpose_triton.py` - -The `matmul_transpose_assign(d_in, d_out)` kernel computes `d_out = d_in @ d_in^T` in-place. It exploits symmetry by computing only upper-triangle blocks and mirroring. - -## QK Clipping - -**File:** `qk_clip.py` - -Optional dynamic clipping for attention head projections (Q and K weight matrices). When the maximum QK logit for a head exceeds a threshold, the corresponding rows of the weight matrix are scaled down by `sqrt(threshold / logit)`. - -**In the parallel pipeline:** QK clipping is applied per-row using each row's global head index. This correctly handles strided sharding where local rows may be interleaved across multiple heads: - -```python -# pipeline.py: _update_params() -ratio = p.shape[0] // scales_full.shape[0] # rows per head -idx0 = state.rank_indices[rank][0] # which global rows this rank owns -row_scales = scales_full[idx0 // ratio] # map each row to its head's scale -p._local_tensor.mul_(row_scales.view(-1, 1)) -``` - -## AdamW for Non-Muon Parameters - -**File:** `adamw.py` - -Parameters not eligible for Muon (1D parameters, embeddings, LM head) are optimized with fused AdamW via `torch._fused_adamw_`. Parameters are grouped by device/dtype and DTensor placement before the fused call. - -## Source File Map - -| File | Lines | Purpose | -|------|-------|---------| -| `muon.py` | ~815 | Optimizer class, parameter routing, 3 execution paths, MoE expert expansion + caching | -| `pipeline.py` | ~400 | Generator-based parallel pipeline (gather/compute/scatter/update) | -| `async_utils.py` | ~75 | Pipeline scheduler with bounded concurrency | -| `core.py` | ~175 | `_muon_state` dataclass, batched momentum/update helpers, param grouping | -| `distributed/utils.py` | ~230 | Shard mesh construction, DTensor index computation | -| `newton_schulz.py` | ~190 | Polar Express coefficients, Newton-Schulz iteration + compile/CUDA graph | -| `matmul_transpose_triton.py` | ~130 | Triton kernel for symmetric matmul | -| `qk_clip.py` | ~135 | QK logit clipping | -| `adamw.py` | ~170 | Fused AdamW for non-Muon params | - -### Dependency Graph - -``` -matmul_transpose_triton.py (leaf) - | - newton_schulz.py (leaf + triton) - | - core.py ---- qk_clip.py (leaf, distributed/utils) - | | | - | pipeline.py --- async_utils.py - | | - | adamw.py - | | - muon.py (all above) - | - __init__.py -``` diff --git a/docs/muon-clip.md b/docs/muon-clip.md deleted file mode 100644 index bcb77228e11beb4a2407085f34d19ebc12e54d56..0000000000000000000000000000000000000000 --- a/docs/muon-clip.md +++ /dev/null @@ -1,317 +0,0 @@ -# QK-Clip for MuonClip Optimizer (MLA) - -> Reference: [Kimi K2 Technical Report](https://arxiv.org/pdf/2507.20534), Section 2.1, Algorithm 1 - -## 개요 - -QK-Clip은 Muon optimizer에서 발생하는 attention logit explosion을 방지하기 위한 **weight rescaling** 기법이다. -forward/backward에는 개입하지 않고, optimizer step **이후**에 weight를 rescale하여 logit 성장을 원천 차단한다. - -## Algorithm 1: MuonClip - -``` -for each training step t: - // 1. Muon optimizer step - for each weight W: - Mt = µ·Mt-1 + Gt - Ot = Newton-Schulz(Mt) · √max(n,m) · 0.2 - Wt = Wt-1 - η·(Ot + λ·Wt-1) - - // 2. QK-Clip - for each attention head h: - S^h_max ← forward에서 기록한 head h의 max pre-softmax logit - if S^h_max > τ: - γ ← τ / S^h_max - W^h_qc ← W^h_qc · √γ (query compressed, q_nope) - W^h_kc ← W^h_kc · √γ (key compressed, k_nope) - W^h_qr ← W^h_qr · γ (query rotary, q_pe) - // k_R (shared rotary, k_pe): 안 건드림 -``` - -## 기존 코드 → MLA 수도코드 - -### 현재 코드 구조 (MHA/GQA) - -``` -parse_qk_layer(name) → wq/wk 여부 판별, layer index 추출 -get_qk_clip_info(config, n) → QKClipInfo (kind, indices, head_dim, threshold, logit) -compute_scales(p, info) → per-head √γ scales 텐서 반환 -qk_clip(p, scales, head_dim) → W.view(-1, head_dim, in_dim).mul_(scales) -``` - -현재 코드는 head_dim이 균일하고, Q/K weight 전체에 동일한 √γ를 적용한다. - -### MLA에서 달라지는 점 - -| 항목 | MHA/GQA (현재) | MLA | -|---|---|---| -| Q weight | `wq` / `q_proj` | `wq_b` (up-proj from LoRA) | -| K weight | `wk` / `k_proj` | `wkv_b` (k_nope + v 합쳐져 있음) | -| Q head stride | `qk_head_dim` (균일) | `qk_head_dim` = `qk_nope_head_dim + qk_rope_head_dim` | -| K head stride | `qk_head_dim` (균일) | `kv_stride` = `qk_nope_head_dim + v_head_dim` | -| Q scaling | 전체 √γ | nope → √γ, rope → γ (서로 다름) | -| K scaling | 전체 √γ | k_nope → √γ, v → 1.0 (부분만) | -| shared k_pe | 없음 | `wkv_a` 뒷부분, 안 건드림 | - -### 수도코드: parse_qk_layer (MLA 확장) - -```python -def parse_qk_layer(name: str) -> tuple[str | None, int]: - parts = normalize_fqn(name).split('.') - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - # MHA/GQA: wq, wk, q_proj, k_proj - # MLA: wq_b (Q up-proj), wkv_b (KV up-proj) - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 -``` - -### 수도코드: QKClipInfo (MLA 확장) - -```python -@dataclass -class QKClipInfo: - kind: str | None # 'wq_b' or 'wkv_b' (MLA) / 'wq','wk' (MHA) - indices: list[int] # clipping 대상 head indices - head_dim: int # 기존 MHA용 (uniform stride) - threshold: float - logit: torch.Tensor | None - - # MLA 전용 필드 - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 -``` - -### 수도코드: get_qk_clip_info (MLA 확장) - -```python -def get_qk_clip_info(clip_config, n, qk_logits): - if clip_config is None: - return None - - threshold = clip_config['threshold'] - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=clip_config['head_dim'], # qk_head_dim (for wq_b) - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - # 기존 MHA/GQA 경로 - return QKClipInfo( - kind=kind, indices=indices, - head_dim=clip_config['head_dim'], - threshold=threshold, logit=logit, - ) -``` - -### 수도코드: compute_scales (MLA 확장) - -기존과 동일하게 per-head γ를 계산한다. (γ 결정은 MHA와 동일) -달라지는 건 `qk_clip` 적용 시 head 내부를 sub-region별로 나눠서 다른 변환을 쓰는 것이다. - -```python -def compute_scales(p, qk_clip_state): - """기존 코드와 동일. per-head √γ 반환.""" - kind = qk_clip_state.kind - indices = qk_clip_state.indices - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) # √γ - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - - if not head_scales: - return None - - H_global = p.shape[0] // qk_clip_state.head_dim # MLA: head_dim = qk_head_dim or kv_stride - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale # √γ_h - - return scales_full -``` - -### 수도코드: qk_clip (MLA 확장) - -per-head scales(√γ)는 동일하게 받되, head 내부 sub-region에 다른 함수를 적용한다. - -```python -def qk_clip(p, scales, head_dim, is_mla=False, kind=None, info=None): - """ - scales: [n_heads] 텐서, 각 원소 = √γ_h - - is_mla=False: 기존 MHA/GQA (head 내 uniform √γ) - is_mla=True: MLA (head 내 sub-region별 다른 변환) - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not is_mla: - # 기존: 모든 행에 √γ 균일 적용 - W.view(-1, head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: head별로 sub-region 분리 적용 - if kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_rope = info.qk_rope_head_dim - qk_head_dim = qk_nope + qk_rope - - for h in range(len(scales)): - sqrt_gamma = scales[h].item() - if sqrt_gamma >= 1.0: - continue - gamma = sqrt_gamma * sqrt_gamma # √γ → γ - s = h * qk_head_dim - - W[s : s + qk_nope] *= sqrt_gamma # q_nope → √γ - W[s + qk_nope : s + qk_head_dim] *= gamma # q_pe → γ - - elif kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - - for h in range(len(scales)): - sqrt_gamma = scales[h].item() - if sqrt_gamma >= 1.0: - continue - s = h * kv_stride - - W[s : s + qk_nope] *= sqrt_gamma # k_nope → √γ - # v 행: 안 건드림 -``` - -### 수도코드: GQA에서 wkv_b indices 처리 - -Q head → KV head 매핑이 필요하다. -여러 Q head가 같은 KV head를 공유하므로, **group 내 최소 gamma** 기준으로 한 번만 적용해야 한다. - -```python -def build_k_indices_for_mla(clip_config, n_heads, n_kv_heads): - """ - Q head 기준 logit으로부터 KV head indices를 생성한다. - q_indices가 Q head index 기준이라면, - k_indices는 대응되는 KV head index로 변환해야 한다. - - 주의: 같은 KV head에 매핑되는 여러 Q head 중 - 가장 큰 logit (= 가장 작은 gamma)을 사용해야 한다. - """ - heads_per_kv = n_heads // n_kv_heads - q_indices = clip_config.get('q_indices', list(range(n_heads))) - - # Q head → KV head 매핑 - # logit 텐서에서 같은 kv_head에 대응되는 Q head들 중 max를 취하는 것은 - # compute_scales_mla 내부에서 min(gamma) 로 처리됨 - - k_indices = [] - seen = set() - for q_idx in q_indices: - kv_idx = q_idx // heads_per_kv - if kv_idx not in seen: - k_indices.append(kv_idx) - seen.add(kv_idx) - - return k_indices -``` - -### 수도코드: 호출 흐름 (통합) - -```python -# optimizer step 이후 호출되는 부분 (기존 코드 구조 유지) - -for name, param in model.named_parameters(): - info = get_qk_clip_info(clip_config, name, qk_logits) - if info is None or info.kind is None: - continue - - scales = compute_scales(param, info) # per-head √γ (MHA/MLA 공통) - if scales is not None: - qk_clip(param, scales, info.head_dim, - is_mla=info.is_mla, kind=info.kind, info=info) -``` - -### 수도코드: clip_config 예시 - -```python -# MHA/GQA (기존) -clip_config = { - 'head_dim': 128, - 'threshold': 100.0, - 'q_indices': list(range(n_heads)), - 'k_indices': list(range(n_kv_heads)), -} - -# MLA (확장) -clip_config = { - 'is_mla': True, - 'head_dim': 192, # qk_head_dim (= qk_nope + qk_rope) - 'qk_nope_head_dim': 128, - 'qk_rope_head_dim': 64, - 'v_head_dim': 128, - 'threshold': 100.0, - 'q_indices': list(range(n_heads)), - 'k_indices': list(range(n_kv_heads)), # build_k_indices_for_mla로 생성 -} -``` - -## 행 인덱스 매핑 테이블 - -| 알고리즘 기호 | 텐서 | 행 범위 | scale | -|---|---|---|---| -| W^h_qc | `wq_b.weight` | `[h*qk_head_dim : h*qk_head_dim + qk_nope_head_dim]` | √γ | -| W^h_qr | `wq_b.weight` | `[h*qk_head_dim + qk_nope_head_dim : (h+1)*qk_head_dim]` | γ | -| W^h_kc | `wkv_b.weight` | `[kv_h*kv_stride : kv_h*kv_stride + qk_nope_head_dim]` | √γ | -| k_R | `wkv_a` output 뒷부분 | - | 안 건드림 | - -- `kv_stride = qk_nope_head_dim + v_head_dim` -- `kv_h = h // (n_heads // n_kv_heads)` (GQA head 매핑) - -## 하이퍼파라미터 - -| 파라미터 | 값 | 비고 | -|---|---|---| -| τ (threshold) | 100 | K2 full-scale 학습 | -| τ (aggressive) | 30 | 소규모 ablation, 성능 저하 없음 확인 | - -## 참고사항 - -- **Self-deactivation**: K2에서 초기 70k step 동안 12.7%의 head만 trigger됨. 이후 모든 head의 S_max가 τ 아래로 내려가면서 자연스럽게 비활성화. -- **DP/TP 환경**: S^h_max를 all-reduce로 모든 rank에서 max 수집 필요. -- **GQA 중복 적용 방지**: 같은 KV head를 공유하는 Q head group에서 가장 작은 gamma(= 가장 큰 logit)를 기준으로 KV weight를 한 번만 scaling. `compute_scales_mla`에서 `min(gamma)` 로직으로 처리. -- **wq_b_gate**: attention logit이 아닌 output gate에만 관여하므로 QK-Clip 대상 아님. -- **기존 logit soft-cap**: forward-level safety net으로 남겨두되, optimizer-level QK-Clip을 추가하는 것이 논문의 접근법. diff --git a/docs/muon/balanced.png b/docs/muon/balanced.png new file mode 100644 index 0000000000000000000000000000000000000000..2076978a5a0149d598b419bfc45c508405dca0df --- /dev/null +++ b/docs/muon/balanced.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9933e2cd5490513593dd6cf1c5c4f18b7f33fd6e6b11c696784269c2bb78055b +size 98003 diff --git a/docs/muon/distributed_muon.png b/docs/muon/distributed_muon.png new file mode 100644 index 0000000000000000000000000000000000000000..26544c9e035afae48d1b32cd6ae729c600a47f33 --- /dev/null +++ b/docs/muon/distributed_muon.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:31caea472991fd24a7934bf211b5adcbf154b5295bfe364bba5b603851c2cfae +size 407912 diff --git a/docs/muon/distributed_muon_execution.png b/docs/muon/distributed_muon_execution.png new file mode 100644 index 0000000000000000000000000000000000000000..824c728b78c73ca0d5b70a169ed2e5e50a59946c --- /dev/null +++ b/docs/muon/distributed_muon_execution.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:72ab4d8076f1e182900d71636dd22c32b20bface38890cef72a0c94c496d5f02 +size 57140 diff --git a/docs/muon/imbalance.png b/docs/muon/imbalance.png new file mode 100644 index 0000000000000000000000000000000000000000..d63f0a034912195910cfac8a49f0533ac99968b1 --- /dev/null +++ b/docs/muon/imbalance.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c71d5faed05d46b2269fefa3b6bea6791d7bf51744f47aa4bb8c311eda1b27ff +size 56528 diff --git a/docs/muon/main.tex b/docs/muon/main.tex new file mode 100644 index 0000000000000000000000000000000000000000..41ad35f7b47a0bcfd9cc5bcc879b5fa1bf56c6f4 --- /dev/null +++ b/docs/muon/main.tex @@ -0,0 +1,142 @@ +\documentclass{article} +\usepackage{graphicx} +\usepackage{hyperref} +\usepackage{amsmath} +\usepackage{caption} +\usepackage{tgtermes} +\usepackage{float} +\usepackage[a4paper, margin=1in]{geometry} +\usepackage{booktabs} +\usepackage{algorithm} +\usepackage{algorithmicx} +\usepackage{algpseudocode} +\date{} + +\begin{document} + +{\LARGE \bfseries Parallelize Muon with FSDP2 \par} +\vspace{1em} % 제목 아래 간격 조정 + +\section*{Motivation} + +\begin{figure}[H] + \centering + \includegraphics[width=0.8\textwidth]{distributed_muon.png} + \caption*{Distributed Muon by Moonlight} +\end{figure} + +While a distributed version of Muon is available, it has the drawback of redundant computations across GPUs. + +\begin{figure}[H] + \centering + \includegraphics[width=1.0\textwidth]{distributed_muon_execution.png} + \caption*{Execution timeline of Distributed Muon} +\end{figure} + +\begin{itemize} + \item \texttt{C[i]} : Compute Newton-Schulz(G) for i-th gradient + \item \texttt{AG[i]} : AllGather i-th gradient + \item \texttt{G[i]} : Gather i-th gradient + \item \texttt{SC[i]} : Scatter i-th gradient +\end{itemize} +\clearpage +\section*{Algorithm} + +\subsection*{Parallel Muon} + +\begin{algorithm} +\caption{Parallel Muon} +\textbf{Require:} DP partitioned gradient $\mathbf{g}$, DP partitioned Momentum $\mathbf{m}$, DP partitioned parameter $\mathbf{p}$, momentum $\mu$, local rank $\mathbf{r}$ +\begin{algorithmic}[1] +\State \texttt{// Apply momentum to $\mathbf{g}$ using local partitioned momentum $\mathbf{m}$} +\State $\mathbf{g'} \gets \text{update\_with\_momentum}(\mathbf{g}, \mathbf{m}, \mu)$ +\State \texttt{// Schedule $\mathbf{g'}$ to rank $\mathbf{R}$} +\State $\mathbf{R} \gets \text{schedule}(\mathbf{g'}, \text{dp\_group})$ +\State \texttt{// Gather $\mathbf{g'}$ across DP into a full matrix $\mathbf{G}$ to rank $\mathbf{R}$} +\State $\mathbf{G} \gets \text{gather}(\mathbf{g'}, \text{dp\_group}, \text{dst=}\mathbf{R})$ +\State \texttt{// Calculate Newton-Schulz only in $\mathbf{R}$} +\If{$\mathbf{r}$ == $\mathbf{R}$} + \State $\mathbf{u} \gets \text{Newton-Schulz}(\mathbf{G})$ +\Else + \State $\mathbf{u} \gets None$ +\EndIf + +\State \texttt{// Scatter a full matrix $\mathbf{u}$ across DP} +\State $\mathbf{u'} \gets \text{scatter}(\mathbf{u},\text{dp\_group},\text{src=}\mathbf{R})$ +\State \texttt{// Apply DP partitioned $\mathbf{u'}$ to $\mathbf{p}$} +\State $\mathbf{p'} \gets \text{apply\_update}(\mathbf{p}, \mathbf{u'})$ +\State \textbf{return $\mathbf{p'}$} +\end{algorithmic} +\end{algorithm} + +We eliminate redundant computation by assigning each parameter to a specific GPU. + +However, without proper scheduling, this optimization can lead to poor GPU utilization. In particular, although redundant computation is avoided by assigning each parameter to a specific rank, it causes idle time—since all other ranks must wait for the scatter communication to complete before proceeding. + +\begin{figure}[H] + \centering + \includegraphics[width=1.0\textwidth]{naive_execution.png} + \caption*{Execution timeline of Parallel Muon} +\end{figure} + +\subsection*{Scheduling Sub-Operations} + +We can schedule the whole sub-operations as follows, due to the following reasons: +\begin{itemize} + \item There are no dependencies between parameters. + \item GPUs can execute computation and communication concurrently. +\end{itemize} + +\begin{figure}[H] + \centering + \includegraphics[width=1.0\textwidth]{pipelined.png} + \caption*{Execution timeline of re-scheduled Parallel Muon} +\end{figure} + +We define the chunk size $C$ as the number of GPUs and schedule each sub-operation in batches of size $C$. This scheduling allows each GPU to continue computation even while waiting for collective communication to complete. + +\textbf{[Algorithm]} (To be written) +\clearpage +\subsection*{Load Balancing} + +If parameters in a chunk have imbalanced computation loads, idle bubbles may occur. \\ +To mitigate this, we apply load balancing based on per-parameter FLOPs. + +\vspace{1em} +\textbf{Imbalanced (Round Robin)} + +\begin{figure}[H] + \centering + \includegraphics[width=1.0\textwidth]{imbalance.png} +\end{figure} + +\textbf{After Load Balancing} + +\begin{figure}[H] + \centering + \includegraphics[width=1.0\textwidth]{balanced.png} +\end{figure} + +\section*{Implementation} + +The full implementation is available in \texttt{optimizer/torch-ext/optimizer/muon.py}. +To enable concurrent computation and communication, we use separate compute and communication streams (\texttt{torch.cuda.Stream}) and use \texttt{torch.cuda.Event} to synchronize between sub-operations. + +Thanks to the simplicity of \texttt{torch.DTensor} and \texttt{torch.distributed}, the implementation remains straightforward and low in complexity. + +\section*{Evaluation} +We evaluated the performance using 10B model currently in development, achieving 151 TFLOPS per GPU during the optimizer step. + +\begin{table}[H] + \centering + \begin{tabular}{@{}lllll@{}} + \toprule + Model Size & TFLOPs for Muon & GPUs & Elapsed time & TFLOPS/GPU \\ + \midrule + 10B & 847.45 & 4xMI250 (8 devices) & 1.4 s & 151 \\ + \bottomrule + \end{tabular} +\end{table} +Based on the breakdown, 7\% of the time is attributed to updating sharded gradients and parameters, 78\% to GEMM operations, and the remaining 15\% to non-overlapped communication overhead. + +\end{document} \ No newline at end of file diff --git a/docs/muon/naive_execution.png b/docs/muon/naive_execution.png new file mode 100644 index 0000000000000000000000000000000000000000..e8f3c4ce721cda02eb95f569c58739d36008b525 --- /dev/null +++ b/docs/muon/naive_execution.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:eaacd3625f33cee9735ed0d96b95f98c696dfc771976be970a38c991e2ce84ab +size 42729 diff --git a/docs/muon/parallel_muon.pdf b/docs/muon/parallel_muon.pdf new file mode 100644 index 0000000000000000000000000000000000000000..8321c572edfae32e963a013d69187d58971fc27e --- /dev/null +++ b/docs/muon/parallel_muon.pdf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c1a88537a50ecc3db52d6e148d3513b31e2c9810c09df0da8f6aff03fa652fe5 +size 654538 diff --git a/docs/muon/pipelined.png b/docs/muon/pipelined.png new file mode 100644 index 0000000000000000000000000000000000000000..7e3d51f98c8f2e501704298c6ec48dca08203884 --- /dev/null +++ b/docs/muon/pipelined.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1f8043cc58e7d8d9da5694ad7bccd1b9fe0210349b9aa9a62652a97f75cf097 +size 64316 diff --git a/docs/optimizations.md b/docs/optimizations.md deleted file mode 100644 index 5a492d6f6ee5e72765e9d91e0b5210a68af08e18..0000000000000000000000000000000000000000 --- a/docs/optimizations.md +++ /dev/null @@ -1,125 +0,0 @@ -# Performance Optimizations (vs. main) - -Summary of optimizations on branch `perf/pipelined-distributed-muon-clean` relative to `main`. - ---- - -## 1. Batched Momentum (`core.py`) - -**Before:** Per-param `update_g()` — one `torch.add` + optional `torch.add_` per parameter. - -**After:** `_batch_pre_ortho()` — `_foreach_mul_`, `_foreach_add_` on lists of local tensors (unwrapped from DTensor). Single fused kernel per batch instead of N individual kernels. - -**Impact:** Eliminates N per-param Python-loop overhead + N small kernel launches. Scales with parameter count. - ---- - -## 2. Pipeline Buffer Packing (`pipeline.py`) - -### Gather send buffer - -**Before:** Per-param `.to(COMM_DTYPE).contiguous()` followed by per-destination `append` to list, then `torch.cat` on the per-dst lists. - -**After:** Collect all grad slices in destination order in a single pass, then one `torch.cat` call. Avoids intermediate per-destination lists and redundant dtype conversions. - -### Scatter send buffer - -**Before:** Per-param, per-destination-rank: index `u_full[indices].flatten()`, append to per-dst list, then flatten+cat. - -**After:** Cache `u_full` conversions (avoid redundant `.to()` per dst_rank). Collect all slices in dst order in one pass, single `torch.cat`. - -**Impact:** Fewer kernel launches, less Python overhead, reduced intermediate allocations. - ---- - -## 3. Zero-Copy Scatter (`pipeline.py`) - -**Before:** `_launch_scatter` pre-allocates `torch.empty_like(p.to_local())` for every param. `_complete_scatter` copies from recv_buf into these pre-allocated tensors via `copy_()`. - -**After:** `_complete_scatter` assigns **views** into `recv_buf` directly (via `recv_buf.narrow(...).view_as(...)`). No pre-allocation, no copy. The recv_buf storage stays alive through the views until `_update_params` consumes them. - -**Impact:** Eliminates N `empty_like` allocations + N `copy_` kernel launches per scatter stage. - ---- - -## 4. Batched Parameter Update (`pipeline.py`) - -**Before:** Per-param loop calling `update_p()` (which unwraps DTensor, applies weight decay, applies update individually). - -**After:** Batched using `_foreach_mul_` (weight decay) and `_foreach_add_` (Muon update), grouped by `adjusted_lr` to preserve float32 alpha precision. Single kernel per group instead of per param. - -**Impact:** Reduces N per-param kernel launches to 1-2 batched kernel launches. - ---- - -## 5. Parallel Metadata Caching (`muon.py`) - -**Before:** `init_state_and_assign_params()` called every step — sorts params by FLOP cost, assigns ownership via round-robin, precomputes per-rank indices/numels for all-to-all. - -**After:** `_parallel_cache` keyed by `tuple(names)`. First call computes and caches `ordered_names`, `name_to_state`, `rank`, `chunk_size`. Subsequent calls reuse cached metadata, only rebuilding `param_to_state` with current `id(p)` keys (since param objects are stable but ids may change for QK clip updates). - -**Impact:** Eliminates repeated sorting, mesh construction, and index precomputation on every step. - ---- - -## 6. Expert Param Expansion Caching (`muon.py`) - -**Before:** `_expand_expert_params()` called every step — for each expert param `(E, out, in)`, creates E `nn.Parameter` wrappers (triggers `aten::detach`), indexes data and grad (`aten::select`), and wraps in DTensor for TP. - -**After:** `_expert_expand_cache` keyed by `tuple(id(p) for p in params)`. Cold path runs `_expand_expert_params` once and caches: - -- `expanded_names` / `expanded_params` — the nn.Parameter wrappers with stable data views -- `grad_info` — per-expert-group metadata (orig param index, num experts, expanded start index, DTensor flag, TP mesh/placements) - -Hot path reuses cached nn.Parameter objects (data views are stable since optimizer updates happen in-place on the same storage). Only updates `.grad` on each cached expert param by slicing the current step's gradient. - -**Eliminated on hot path:** - -- `nn.Parameter()` construction — removes `aten::detach` -- `local_data[i]` data slicing — removes half of `aten::select` + `aten::as_strided` -- `DTensor.from_local()` for data — only needed for grad now -- `is_expert_param()` name matching per step - -**Still required per step:** - -- `local_grad[i]` — grad tensor changes each step (nesterov) -- `DTensor.from_local(slice_grad, ...)` — for TP expert grads -- `p.grad = None` — freeing original 3D grad storage - -**Impact:** ~8ms CPU overhead reduction per step at production scale (64 GPUs, 48 local experts). - ---- - -## 7. Newton-Schulz Compile + CUDA Graph (`newton_schulz.py`) - -**Before:** `_zeropower_via_newtonschulz5()` called directly every time. - -**After:** `zeropower_via_newtonschulz5()` wrapper with per-shape `torch.compile` caching + CUDA graph (`triton.cudagraphs=True`). Each unique shape gets its own compiled function stored in `_ns_per_shape`. Toggled via `set_ns_compile(enabled)`. - -**Impact:** After warmup, NS iterations run as CUDA graphs — eliminates per-step compilation overhead and CPU-GPU synchronization. - ---- - -## 8. Removed `small_param_numel_threshold` (`muon.py`) - -**Before:** Small sharded DTensors (below threshold, default 65536) fell back to `distributed_muon()` which used per-param `full_tensor()` + redistribute. - -**After:** All sharded DTensors go to `parallel()`. `distributed_muon()` is retained as a test-only reference implementation. Uneven shard splits (e.g., MoE gate weights with fewer rows than shard ranks) are handled inline via `full_tensor()` fallback within the batched distributed_muon path. - -**Impact:** Simpler routing, no silent fallback to slower path. - ---- - -## Summary Table - -| Optimization | Location | Category | Kernel Launches Saved | -|---|---|---|---| -| Batched momentum | `core.py` | CPU + GPU | N per-param → 2-3 batched | -| Buffer packing (gather) | `pipeline.py` | CPU + GPU | N cat+cast → 1 cat+cast | -| Buffer packing (scatter) | `pipeline.py` | CPU + GPU | N cat → 1 cat | -| Zero-copy scatter | `pipeline.py` | GPU memory | N alloc+copy → 0 | -| Batched param update | `pipeline.py` | CPU + GPU | N update → 1-2 batched | -| Parallel metadata cache | `muon.py` | CPU | Sort+index per step → once | -| Expert expand cache | `muon.py` | CPU | N detach+select → grad-only | -| NS compile + CUDA graph | `newton_schulz.py` | GPU | JIT warmup → graph replay | -| Remove small_param_threshold | `muon.py` | Routing | Simpler, unified path | diff --git a/docs/pytorch-2.10-tp-fix.md b/docs/pytorch-2.10-tp-fix.md deleted file mode 100644 index 3533cad5e5c5cfcc63d99c7d94ad2577db182646..0000000000000000000000000000000000000000 --- a/docs/pytorch-2.10-tp-fix.md +++ /dev/null @@ -1,151 +0,0 @@ -# PyTorch 2.10 Tensor Parallelism Fix - -## Summary - -PyTorch 2.10 changed the class hierarchy for `_StridedShard`, breaking our -`distributed/utils.py` code that handles DTensor sharding. This document -records the root cause, every change made so far, and the one remaining issue. - ---- - -## 1. Root Cause: `_StridedShard` class hierarchy change - -| Version | MRO | -|---------|-----| -| PyTorch < 2.10 | `_StridedShard -> Shard -> Placement` | -| PyTorch 2.10 | `_StridedShard -> StridedShard -> Placement` | - -**Consequences:** - -- `isinstance(strided_shard, Shard)` returns `False` -- `strided_shard.is_shard()` returns `False` -- Our `construct_shard_mesh()` treated `_StridedShard` as an unsupported - placement and raised `AssertionError`. - -### When does `_StridedShard` appear? - -When two parallelism dimensions shard the same tensor dimension. -For example, `fsdp+tp` or `hsdp+tp` configurations where both TP and -FSDP shard dimension 0 of Q/K/V projection weights: - -``` -TP : Shard(0) → each TP rank gets 2048/4 = 512 rows -FSDP: Shard(0) on top → each FSDP rank further splits those rows -``` - -PyTorch represents the second sharding as `_StridedShard(dim=0, split_factor=N)` -to indicate non-contiguous (interleaved) row ownership. - ---- - -## 2. Completed Fixes - -### 2.1 `_is_shard()` helper (`distributed/utils.py`) - -Added a helper that correctly identifies both `Shard` and `_StridedShard`: - -```python -def _is_shard(placement: Placement) -> bool: - return isinstance(placement, (Shard, _StridedShard)) -``` - -Used in `construct_shard_mesh()` where the old code called `placement.is_shard()`. - -### 2.2 Rewritten `get_slices_of_dtensor()` (`distributed/utils.py`) - -Old code assumed contiguous slicing (`start = rank * shard_size`), which is -wrong for `_StridedShard`. - -New code uses PyTorch's own offset-computation methods: - -| Placement type | API used | -|----------------|----------| -| `Shard` | `Shard.local_shard_size_and_offset(size, chunks, rank)` → `(size, offset)` | -| `_StridedShard` | `_StridedShard.local_shard_size_and_offset(instance, size, chunks, rank, return_first_offset=False)` → `(size, offsets_list)` | - -Return type changed from `tuple[slice, ...]` to `tuple[slice | torch.Tensor, ...]`: -- `slice` for contiguous ranges (Shard or contiguous StridedShard result) -- `torch.LongTensor` of indices for non-contiguous ranges - -Composed sharding (multiple placements on the same dim) is handled by -indexing: `dim_indices[shard_dim] = dim_indices[shard_dim][new_indices]`. - -### 2.3 Updated `numel_for_rank()` (`core.py`) - -Now handles both `slice` and `torch.Tensor` index types: - -```python -for idx, dim_size in zip(indices, param.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) -``` - -### 2.4 Updated pipeline stages (`pipeline.py`) - -- **`_gather_grads()`**: Uses `gathered_grads[id(p)][indices] = sg` (index - assignment) instead of view-based copy, works for both slice and tensor - indices. -- **`_scatter_us()`**: `u_full[indices].flatten()` works for both index types. -- **`_update_params()` QK clipping**: Applies clipping on `p._local_tensor` - directly instead of through DTensor operations, avoiding sharding - propagation errors with `_StridedShard`. - ---- - -## 3. Test Results After Fixes - -| Test configuration | Status | -|--------------------|--------| -| base | PASS | -| fsdp | PASS | -| hsdp | PASS | -| tp | PASS | -| hsdp+tp (QK clip off) | PASS | -| hsdp+tp (QK clip on) | PASS | -| fsdp+tp (QK clip off) | PASS | -| fsdp+tp (QK clip on) | PASS | - -All 24 tests pass (126s, `--skip-verify`). - ---- - -## 4. Fixed: QK Clipping with Strided Sharding - -### Problem - -With strided (non-contiguous) sharding, local rows are **interleaved across -multiple heads**. For example with `fsdp+tp` (`dp_shard=2, tp=4`): - -- Q/K projection global shape: `(2048, 2048)`, `head_dim=128`, `16 heads` -- Each rank owns 256 rows, but they span 4 heads with 64 rows per head -- `view(-1, head_dim, cols)` assumes contiguous head blocks → wrong for - interleaved rows → shape mismatch error - -### Fix Applied - -For strided sharding, apply scales **per-row** based on each row's global -head index instead of using the head-block view: - -```python -if isinstance(weight_indices[0], slice): - # Contiguous case: view-based approach still works - ... -else: - # Strided case: per-row scaling - head_per_row = weight_indices[0] // ratio - row_scales = scales_full[head_per_row] - local_p.mul_(row_scales.view(-1, 1)) -``` - ---- - -## 5. Files Modified - -| File | Changes | -|------|---------| -| `torch-ext/optimizer/distributed/utils.py` | Added `_is_shard()`, rewrote `get_slices_of_dtensor()`, fixed `construct_shard_mesh()` | -| `torch-ext/optimizer/core.py` | Updated `numel_for_rank()` for `slice | Tensor` indices | -| `torch-ext/optimizer/pipeline.py` | Updated `_gather_grads()`, `_scatter_us()`, `_update_params()` QK clipping | diff --git a/flake.lock b/flake.lock new file mode 100644 index 0000000000000000000000000000000000000000..368754a84e467fe6ba68962628649fc9ab6121cc --- /dev/null +++ b/flake.lock @@ -0,0 +1,167 @@ +{ + "nodes": { + "flake-compat": { + "locked": { + "lastModified": 1747046372, + "narHash": "sha256-CIVLLkVgvHYbgI2UpXvIIBJ12HWgX+fjA8Xf8PUmqCY=", + "owner": "edolstra", + "repo": "flake-compat", + "rev": "9100a0f413b0c601e0533d1d94ffd501ce2e7885", + "type": "github" + }, + "original": { + "owner": "edolstra", + "repo": "flake-compat", + "type": "github" + } + }, + "flake-compat_2": { + "locked": { + "lastModified": 1733328505, + "narHash": "sha256-NeCCThCEP3eCl2l/+27kNNK7QrwZB1IJCrXfrbv5oqU=", + "owner": "edolstra", + "repo": "flake-compat", + "rev": "ff81ac966bb2cae68946d5ed5fc4994f96d0ffec", + "type": "github" + }, + "original": { + "owner": "edolstra", + "repo": "flake-compat", + "type": "github" + } + }, + "flake-utils": { + "inputs": { + "systems": "systems" + }, + "locked": { + "lastModified": 1731533236, + "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=", + "owner": "numtide", + "repo": "flake-utils", + "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b", + "type": "github" + }, + "original": { + "owner": "numtide", + "repo": "flake-utils", + "type": "github" + } + }, + "flake-utils_2": { + "inputs": { + "systems": "systems_2" + }, + "locked": { + "lastModified": 1731533236, + "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=", + "owner": "numtide", + "repo": "flake-utils", + "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b", + "type": "github" + }, + "original": { + "owner": "numtide", + "repo": "flake-utils", + "type": "github" + } + }, + "hf-nix": { + "inputs": { + "flake-compat": "flake-compat_2", + "flake-utils": "flake-utils_2", + "nixpkgs": "nixpkgs" + }, + "locked": { + "lastModified": 1748598786, + "owner": "huggingface", + "repo": "hf-nix", + "rev": "6ca679441494139fde1f2355691ddb5dc8170269", + "type": "github" + }, + "original": { + "owner": "huggingface", + "repo": "hf-nix", + "type": "github" + } + }, + "kernel-builder": { + "inputs": { + "flake-compat": "flake-compat", + "flake-utils": "flake-utils", + "hf-nix": "hf-nix", + "nixpkgs": [ + "kernel-builder", + "hf-nix", + "nixpkgs" + ] + }, + "locked": { + "lastModified": 1749822059, + "narHash": "sha256-zype8KSqESZUIQpsY6sbf4f9pPxM/Zwem+KuH5LeHFk=", + "owner": "huggingface", + "repo": "kernel-builder", + "rev": "96abd968baa5fa16217413050fa7372d5db3baa5", + "type": "github" + }, + "original": { + "owner": "huggingface", + "repo": "kernel-builder", + "type": "github" + } + }, + "nixpkgs": { + "locked": { + "lastModified": 1747820358, + "narHash": "sha256-fTqsZsUX6M3yeEvgyQvXcbGmT2CaRVyVwsi8eK29Oj4=", + "owner": "danieldk", + "repo": "nixpkgs", + "rev": "d3c1681180717528068082103bf323147de6ab0b", + "type": "github" + }, + "original": { + "owner": "danieldk", + "ref": "cudatoolkit-12.9-kernel-builder", + "repo": "nixpkgs", + "type": "github" + } + }, + "root": { + "inputs": { + "kernel-builder": "kernel-builder" + } + }, + "systems": { + "locked": { + "lastModified": 1681028828, + "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", + "owner": "nix-systems", + "repo": "default", + "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", + "type": "github" + }, + "original": { + "owner": "nix-systems", + "repo": "default", + "type": "github" + } + }, + "systems_2": { + "locked": { + "lastModified": 1681028828, + "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", + "owner": "nix-systems", + "repo": "default", + "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", + "type": "github" + }, + "original": { + "owner": "nix-systems", + "repo": "default", + "type": "github" + } + } + }, + "root": "root", + "version": 7 +} diff --git a/optimizer/dummy.cu b/optimizer/dummy.cu index 4a37c61c2feb239dbe063d08fd763ee7d0bb86ad..9a9780b635e46a946e7c836cd390c24e41da3385 100644 --- a/optimizer/dummy.cu +++ b/optimizer/dummy.cu @@ -3,4 +3,4 @@ namespace { __global__ void dummy() { // This kernel does nothing but serves as a placeholder } -} // namespace +} diff --git a/test/README.md b/test/README.md deleted file mode 100644 index 35ae009acbca01163cc3a3cddac000fa957c0a4f..0000000000000000000000000000000000000000 --- a/test/README.md +++ /dev/null @@ -1,42 +0,0 @@ -# Muon Optimizer Test - -This directory contains a test script for the **Muon optimizer**. - -## Prerequisites - -- **GPU Requirement** - - All tests require **8 GPUs** by default. - - If you have fewer GPUs available: - - Modify the parallelism configurations in `test_muon.py`. - -- **Model Access** - - The tests require access to the private model repository: - - `Motif-Technologies/Motif-2.6B-4layer-random` on Hugging Face. - - Set your Hugging Face token via the environment variable `HF_TOKEN`. - - If you don’t have access, please contact the maintainer. - -- **Using a Different Model (Optional)** - - You may modify the test to use a different model by: - - Updating the model name in `conftest.py::inputs`. - - Adjusting the tensor parallel rules in `utils.py::_apply_tp`. - -## Usage - -- To execute the test with 8 GPUs, simply run: - -```bash -./run_test.sh -``` - -- To check the other available options, you can use: - -```bash -pytest --help -... -Custom options: - --measure-perf Measure execution time and peak memory usage during optimizer step. - --do-profile Enable profiling during tests. - --skip-verify Skip verification of optimizer step correctness with sequential implementation. - This can be useful when GPU memory is limited. -... -``` diff --git a/test/__init__.py b/test/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/test/conftest.py b/test/conftest.py deleted file mode 100644 index 814753950f7bfb909e8c37bdcb0b96bc8c6ce266..0000000000000000000000000000000000000000 --- a/test/conftest.py +++ /dev/null @@ -1,183 +0,0 @@ -import logging - -import pytest -import torch -import torch.distributed as dist -from packaging import version -from transformers import AutoModelForCausalLM - -logger = logging.getLogger(__name__) -logging.basicConfig(level=logging.INFO) - -# Raise dynamo recompile limit so that compiled momentum (batch_pre_ortho) -# does not fall back to eager mode when the test suite runs 30+ model -# configurations with different tensor shapes in a single process. -torch._dynamo.config.recompile_limit = 64 - -SEED = 0xdeadbeef - - -def pytest_addoption(parser): - parser.addoption( - "--measure-perf", - action="store_true", - default=False, - help= - "Measure execution time and peak memory usage during optimizer step.", - ) - - parser.addoption( - "--do-profile", - action="store_true", - default=False, - help="Enable profiling during tests.", - ) - - parser.addoption( - "--skip-verify", - action="store_true", - default=False, - help= - "Skip verification of optimizer step correctness with sequential implementation.\n" - "This can be useful when GPU memory is limited.", - ) - - -def pytest_configure(config): - if config.getoption( - "--do-profile") and not config.getoption("--measure-perf"): - raise pytest.UsageError( - "--do-profile requires --measure-perf. Please enable both flags.") - - -@pytest.fixture(scope="session") -def measure_perf(request): - return request.config.getoption("--measure-perf") - - -@pytest.fixture(scope="session") -def do_profile(request): - return request.config.getoption("--do-profile") - - -@pytest.fixture(scope="session") -def skip_verify(request): - return request.config.getoption("--skip-verify") - - -@pytest.fixture(scope="session", autouse=True) -def init_dist(request): - if version.parse(torch.__version__) < version.parse("2.8"): - pytest.skip("torch>=2.8.0 is required for parallel muon") - return - - try: - dist.init_process_group(backend="nccl") - torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count()) - except Exception as e: - print(f"Failed to initialize torch.distributed: {e}") - pytest.skip("Failed to initialize torch.distributed") - - if dist.get_world_size() != 8: - pytest.skip("Need 8 processes in dist group. " - "You can run with `torchrun --nproc-per-node=8 " - "--local-ranks-filter 0 -m pytest " - "test_rms_norm_sequence_parallel.py`." - "To run with less than 8 gpus, modify " - "the test cases accordingly.") - - yield - dist.destroy_process_group() - - -@pytest.fixture(scope="session") -def inputs(): - """Load Motif-2.6B model and generate random gradients for testing. - Returns: - tuple[torch.nn.Module, list[torch.Tensor], dict[int, torch.Tensor]]: - - torch.nn.Module: The Motif-2.6B model. - - list[torch.Tensor]: A list of random gradients for each model parameter. - - dict[int, torch.Tensor]: A dictionary mapping layer indices to random QK logits. - """ - model_name = "Motif-Technologies/Motif-2.6B-4layer-random" - - torch.manual_seed(SEED) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(SEED) - - model = AutoModelForCausalLM.from_pretrained( - model_name, - trust_remote_code=True, - ) - logger.info( - f"Loaded model {model_name}. ({len(list(model.parameters()))} parameters)" - ) - - grads: list[torch.Tensor] = [] - for param in model.parameters(): - grad = torch.randn_like(param, device=param.device, dtype=param.dtype) - grads.append(grad) - - qk_logits: dict[int, torch.Tensor] = { - i: - torch.randn(model.config.num_attention_heads, - device=model.device, - dtype=torch.bfloat16) - for i in range(model.config.num_hidden_layers) - } - - return [model, grads, qk_logits] - - -def _create_moe_model(num_experts=8, top_k=2, n_layers=4): - """Create a torchtitan Llama4 MoE model with random gradients.""" - from torchtitan.models.llama4.model.args import TransformerModelArgs - from torchtitan.models.llama4.model.model import Transformer - from torchtitan.models.moe import MoEArgs - - torch.manual_seed(SEED) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(SEED) - - moe_args = MoEArgs( - num_experts=num_experts, - num_shared_experts=1, - top_k=top_k, - score_func="sigmoid", - ) - model_args = TransformerModelArgs( - dim=2048, - n_layers=n_layers, - n_heads=16, - n_kv_heads=8, - vocab_size=32000, - norm_eps=1e-5, - rope_theta=10000, - max_seq_len=4096, - moe_args=moe_args, - interleave_moe_layer_step=1, - ) - model = Transformer(model_args) - model.init_weights() - logger.info(f"Created torchtitan Llama4 MoE model " - f"(num_experts={num_experts}, n_layers={n_layers}, " - f"{len(list(model.parameters()))} parameters)") - - grads = [ - torch.randn_like(param, device=param.device, dtype=param.dtype) - for param in model.parameters() - ] - - return [model, grads] - - -@pytest.fixture(scope="session") -def moe_inputs(): - """MoE model with 8 experts (standard config).""" - return _create_moe_model(num_experts=8, top_k=2) - - -@pytest.fixture(scope="session") -def moe_inputs_few_experts(): - """MoE model with 2 experts (triggers EFSDP Shard(1) mode).""" - return _create_moe_model(num_experts=2, top_k=1) diff --git a/test/optimizer b/test/optimizer deleted file mode 120000 index c7ff828a90e1c2a67535184c5e89724fb52bea24..0000000000000000000000000000000000000000 --- a/test/optimizer +++ /dev/null @@ -1 +0,0 @@ -../torch-ext/optimizer/ \ No newline at end of file diff --git a/test/pytest.ini b/test/pytest.ini deleted file mode 100644 index 11c72fa2e2812b16b1c2e92fb0d78cb4adbda2e5..0000000000000000000000000000000000000000 --- a/test/pytest.ini +++ /dev/null @@ -1,3 +0,0 @@ -[pytest] -log_cli = true -log_cli_level = INFO diff --git a/test/run_test.sh b/test/run_test.sh deleted file mode 100755 index 2c2bd5b362a36dd4facebee1454eb7b4809118f1..0000000000000000000000000000000000000000 --- a/test/run_test.sh +++ /dev/null @@ -1 +0,0 @@ -torchrun --nproc-per-node=8 --local-ranks-filter=0 -m pytest test_muon.py diff --git a/test/run_test_moe.sh b/test/run_test_moe.sh deleted file mode 100755 index e9e5f4fc1e22be35f8071e20c1ce0bae6cb6bd85..0000000000000000000000000000000000000000 --- a/test/run_test_moe.sh +++ /dev/null @@ -1,3 +0,0 @@ -#!/usr/bin/env bash -cd "$(dirname "$0")" -torchrun --nproc-per-node=8 --local-ranks-filter=0 -m pytest test_muon_moe.py "$@" diff --git a/test/test_cpu_memory_peak.py b/test/test_cpu_memory_peak.py deleted file mode 100644 index 7d3df5bf4806c671a149d8823ad5a8e06c1f6d35..0000000000000000000000000000000000000000 --- a/test/test_cpu_memory_peak.py +++ /dev/null @@ -1,537 +0,0 @@ -"""CPU memory peak vs offloaded tensor size verification. - -Compares CPU memory usage with turn_on_cpu_offload() vs no offload to isolate -the actual CPU cost of offloading, separating it from CUDA runtime, -NCCL, and DTensor overhead. - -Run with: - torchrun --nproc-per-node=8 --local-ranks-filter=0 test/test_cpu_memory_peak.py -""" - -import gc -import logging -import os - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Shard, distribute_tensor - -logger = logging.getLogger(__name__) -logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s") - - -def _setup(): - dist.init_process_group(backend="nccl") - rank = dist.get_rank() - torch.cuda.set_device(rank % torch.cuda.device_count()) - return rank, dist.get_world_size() - - -def _make_mesh(world_size): - return dist.init_device_mesh("cuda", (world_size, ), - mesh_dim_names=("dp", )) - - -def get_cpu_rss_bytes(): - """Get current process RSS in bytes from /proc/self/statm.""" - with open("/proc/self/statm") as f: - pages = int(f.read().split()[1]) - return pages * os.sysconf("SC_PAGE_SIZE") - - -def get_pinned_pool_bytes(pool): - """Get total pinned CPU buffer size from CPUOffloadPool.""" - total = 0 - for grp in pool._groups.values(): - cpu_flat = grp["cpu_flat"] - total += cpu_flat.numel() * cpu_flat.element_size() - return total - - -def _run_muon_steps(mesh, dim0, dim1, num_params, num_steps, cpu_offload): - """Run Muon optimizer steps and return final CPU RSS.""" - from optimizer.muon import Muon - from optimizer.newton_schulz import set_ns_compile - - set_ns_compile(False) - torch.manual_seed(42) - gc.collect() - torch.cuda.empty_cache() - - params, names = [], [] - for i in range(num_params): - full = torch.randn(dim0, dim1, device="cuda") - dt = distribute_tensor(full, mesh, [Shard(0)]) - p = torch.nn.Parameter(dt) - params.append(p) - names.append(f"layer.{i}.weight") - - param_groups = [{ - "params": params, - "names": names, - "use_muon": True, - "lr": 0.02, - "weight_decay": 0.01, - "momentum": 0.95, - "nesterov": True, - "ns_steps": 5, - "none_grad": False, - }] - - optim = Muon(params=param_groups, chunk_size=2, warmup_step=1) - if cpu_offload: - optim.turn_on_cpu_offload() - - for step_idx in range(num_steps): - for p in params: - p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"), - mesh, [Shard(0)]) - optim.step() - torch.cuda.synchronize() - - gc.collect() - cpu_rss = get_cpu_rss_bytes() - - pinned_bytes = 0 - if cpu_offload and optim._cpu_offload_pool is not None: - pool = optim._cpu_offload_pool - pinned_bytes = get_pinned_pool_bytes(pool) - - # Cleanup. - del optim, params, param_groups - gc.collect() - torch.cuda.empty_cache() - - set_ns_compile(True) - return cpu_rss, pinned_bytes - - -def test_offload_cpu_cost_isolation(rank, world_size): - """A/B test: measure CPU cost of offload by comparing ON vs OFF.""" - mesh = _make_mesh(world_size) - - dim0, dim1 = 2048, 4096 - num_params = 8 - num_steps = 3 - - if rank == 0: - logger.info("=" * 70) - logger.info("A/B TEST: CPU MEMORY COST OF OFFLOAD (ON vs OFF)") - logger.info("=" * 70) - logger.info("Config: %d params of shape (%d, %d), %d ranks, %d steps", - num_params, dim0, dim1, world_size, num_steps) - logger.info("Local param shape per rank: (%d, %d)", dim0 // world_size, - dim1) - logger.info("-" * 70) - - # Run WITHOUT offload first (baseline). - gc.collect() - torch.cuda.empty_cache() - cpu_before_no_offload = get_cpu_rss_bytes() - cpu_after_no_offload, _ = _run_muon_steps(mesh, - dim0, - dim1, - num_params, - num_steps, - cpu_offload=False) - cpu_growth_no_offload = cpu_after_no_offload - cpu_before_no_offload - - # Run WITH offload. - gc.collect() - torch.cuda.empty_cache() - cpu_before_offload = get_cpu_rss_bytes() - cpu_after_offload, pinned_bytes = _run_muon_steps(mesh, - dim0, - dim1, - num_params, - num_steps, - cpu_offload=True) - cpu_growth_offload = cpu_after_offload - cpu_before_offload - - # Delta = additional CPU cost from offloading. - offload_delta = cpu_growth_offload - cpu_growth_no_offload - - if rank == 0: - logger.info("CPU growth WITHOUT offload: %.2f MB", - cpu_growth_no_offload / 1024**2) - logger.info("CPU growth WITH offload: %.2f MB", - cpu_growth_offload / 1024**2) - logger.info("-" * 70) - logger.info("Pinned buffer size (expected): %.2f MB", - pinned_bytes / 1024**2) - logger.info("Offload delta (WITH - WITHOUT): %.2f MB", - offload_delta / 1024**2) - - if pinned_bytes > 0: - ratio = offload_delta / pinned_bytes - logger.info("Ratio (delta / pinned buffer): %.2fx", ratio) - - if ratio > 1.5: - logger.warning( - "Offload adds %.2f MB CPU memory but pinned buffer is " - "only %.2f MB (%.1f%% overhead beyond expected)", - offload_delta / 1024**2, - pinned_bytes / 1024**2, - (offload_delta - pinned_bytes) / pinned_bytes * 100, - ) - else: - logger.info("Offload CPU cost is within expected range.") - - # Only assert on rank 0 to avoid multi-rank assertion mismatches. - if rank == 0 and pinned_bytes > 0: - ratio = offload_delta / pinned_bytes - assert ratio < 3.0, ( - f"Offload CPU cost ({offload_delta / 1024**2:.2f} MB) is " - f"{ratio:.2f}x the pinned buffer ({pinned_bytes / 1024**2:.2f} MB). " - f"Expected < 3.0x.") - - if rank == 0: - logger.info("PASSED: test_offload_cpu_cost_isolation") - - -def test_cpu_memory_peak_detailed(rank, world_size): - """Detailed per-phase CPU memory tracking for offload.""" - from optimizer.muon import Muon - from optimizer.newton_schulz import set_ns_compile - - set_ns_compile(False) - torch.manual_seed(42) - - mesh = _make_mesh(world_size) - - dim0, dim1 = 2048, 4096 - num_params = 8 - - gc.collect() - torch.cuda.empty_cache() - - if rank == 0: - logger.info("=" * 70) - logger.info("DETAILED PER-PHASE CPU MEMORY TRACKING") - logger.info("=" * 70) - - cpu_0 = get_cpu_rss_bytes() - if rank == 0: - logger.info("[Phase 0] Baseline RSS: %.2f MB", cpu_0 / 1024**2) - - # Create params. - params, names = [], [] - for i in range(num_params): - full = torch.randn(dim0, dim1, device="cuda") - dt = distribute_tensor(full, mesh, [Shard(0)]) - p = torch.nn.Parameter(dt) - params.append(p) - names.append(f"layer.{i}.weight") - - gc.collect() - cpu_1 = get_cpu_rss_bytes() - if rank == 0: - logger.info("[Phase 1] After param creation: %.2f MB (+%.2f MB)", - cpu_1 / 1024**2, (cpu_1 - cpu_0) / 1024**2) - - # Create optimizer. - param_groups = [{ - "params": params, - "names": names, - "use_muon": True, - "lr": 0.02, - "weight_decay": 0.01, - "momentum": 0.95, - "nesterov": True, - "ns_steps": 5, - "none_grad": False, - }] - optim = Muon(params=param_groups, chunk_size=2, warmup_step=1) - optim.turn_on_cpu_offload() - - gc.collect() - cpu_2 = get_cpu_rss_bytes() - if rank == 0: - logger.info("[Phase 2] After optimizer creation: %.2f MB (+%.2f MB)", - cpu_2 / 1024**2, (cpu_2 - cpu_1) / 1024**2) - - # Set grads. - for p in params: - p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"), - mesh, [Shard(0)]) - - gc.collect() - cpu_3 = get_cpu_rss_bytes() - if rank == 0: - logger.info("[Phase 3] After grad creation: %.2f MB (+%.2f MB)", - cpu_3 / 1024**2, (cpu_3 - cpu_2) / 1024**2) - - # Step 1 (creates states + first offload). - optim.step() - torch.cuda.synchronize() - gc.collect() - cpu_4 = get_cpu_rss_bytes() - - pool = optim._cpu_offload_pool - pinned_bytes = get_pinned_pool_bytes(pool) - - if rank == 0: - logger.info( - "[Phase 4] After step 1 (init+offload): %.2f MB (+%.2f MB)", - cpu_4 / 1024**2, (cpu_4 - cpu_3) / 1024**2) - logger.info(" Pinned buffer size: %.2f MB", pinned_bytes / 1024**2) - logger.info(" Step 1 growth vs pinned: %.2f MB extra", - (cpu_4 - cpu_3 - pinned_bytes) / 1024**2) - - # Step 2 (reload + compute + offload). - for p in params: - p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"), - mesh, [Shard(0)]) - optim.step() - torch.cuda.synchronize() - gc.collect() - cpu_5 = get_cpu_rss_bytes() - if rank == 0: - logger.info("[Phase 5] After step 2: %.2f MB (+%.2f MB)", - cpu_5 / 1024**2, (cpu_5 - cpu_4) / 1024**2) - - # Step 3. - for p in params: - p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"), - mesh, [Shard(0)]) - optim.step() - torch.cuda.synchronize() - gc.collect() - cpu_6 = get_cpu_rss_bytes() - if rank == 0: - logger.info("[Phase 6] After step 3: %.2f MB (+%.2f MB)", - cpu_6 / 1024**2, (cpu_6 - cpu_5) / 1024**2) - - # Summary. - total_growth = cpu_6 - cpu_0 - if rank == 0: - logger.info("-" * 70) - logger.info("SUMMARY:") - logger.info(" Total CPU growth: %.2f MB", total_growth / 1024**2) - logger.info(" Pinned buffer: %.2f MB", pinned_bytes / 1024**2) - logger.info(" Overhead: %.2f MB", - (total_growth - pinned_bytes) / 1024**2) - if pinned_bytes > 0: - logger.info(" Ratio: %.2fx", - total_growth / pinned_bytes) - logger.info("") - logger.info(" NOTE: Overhead includes CUDA runtime, NCCL buffers,") - logger.info(" DTensor metadata, and optimizer internals — NOT just") - logger.info(" offload cost. Use A/B test for isolated measurement.") - - set_ns_compile(True) - if rank == 0: - logger.info("PASSED: test_cpu_memory_peak_detailed") - - -def test_offload_cpu_cost_mixed(rank, world_size): - """A/B test for mixed Muon + AdamW offload CPU cost.""" - from optimizer.muon import Muon - from optimizer.newton_schulz import set_ns_compile - - mesh = _make_mesh(world_size) - - muon_dim0, muon_dim1 = 2048, 4096 - num_muon = 8 - adamw_dim = 4096 - num_adamw = 8 - num_steps = 3 - - def run_mixed(cpu_offload): - set_ns_compile(False) - torch.manual_seed(42) - gc.collect() - torch.cuda.empty_cache() - - muon_params, muon_names = [], [] - for i in range(num_muon): - full = torch.randn(muon_dim0, muon_dim1, device="cuda") - dt = distribute_tensor(full, mesh, [Shard(0)]) - p = torch.nn.Parameter(dt) - muon_params.append(p) - muon_names.append(f"layer.{i}.weight") - - adamw_params = [] - for i in range(num_adamw): - full = torch.randn(adamw_dim, device="cuda") - dt = distribute_tensor(full, mesh, [Shard(0)]) - p = torch.nn.Parameter(dt) - adamw_params.append(p) - - param_groups = [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - "lr": 0.02, - "weight_decay": 0.01, - "momentum": 0.95, - "nesterov": True, - "ns_steps": 5, - "none_grad": False, - "adamw_betas": (0.9, 0.95), - "adamw_eps": 1e-8, - }, - { - "params": adamw_params, - "use_muon": False, - "lr": 1e-3, - "weight_decay": 0.01, - "adamw_betas": (0.9, 0.95), - "adamw_eps": 1e-8, - }, - ] - - optim = Muon(params=param_groups, chunk_size=2, warmup_step=1) - if cpu_offload: - optim.turn_on_cpu_offload() - - for step_idx in range(num_steps): - for p in muon_params: - p.grad = distribute_tensor( - torch.randn(muon_dim0, muon_dim1, device="cuda"), mesh, - [Shard(0)]) - for p in adamw_params: - p.grad = distribute_tensor( - torch.randn(adamw_dim, device="cuda"), mesh, [Shard(0)]) - optim.step() - torch.cuda.synchronize() - - gc.collect() - cpu_rss = get_cpu_rss_bytes() - - pinned_bytes = 0 - if cpu_offload and optim._cpu_offload_pool is not None: - pinned_bytes = get_pinned_pool_bytes(optim._cpu_offload_pool) - - del optim, muon_params, adamw_params, param_groups - gc.collect() - torch.cuda.empty_cache() - set_ns_compile(True) - return cpu_rss, pinned_bytes - - if rank == 0: - logger.info("=" * 70) - logger.info("A/B TEST: CPU COST OF MIXED OFFLOAD (Muon + AdamW)") - logger.info("=" * 70) - - gc.collect() - torch.cuda.empty_cache() - cpu_before_no = get_cpu_rss_bytes() - cpu_after_no, _ = run_mixed(False) - growth_no = cpu_after_no - cpu_before_no - - gc.collect() - torch.cuda.empty_cache() - cpu_before_yes = get_cpu_rss_bytes() - cpu_after_yes, pinned_bytes = run_mixed(True) - growth_yes = cpu_after_yes - cpu_before_yes - - delta = growth_yes - growth_no - - if rank == 0: - logger.info("CPU growth WITHOUT offload: %.2f MB", growth_no / 1024**2) - logger.info("CPU growth WITH offload: %.2f MB", - growth_yes / 1024**2) - logger.info("Pinned buffer size: %.2f MB", - pinned_bytes / 1024**2) - logger.info("Offload delta: %.2f MB", delta / 1024**2) - if pinned_bytes > 0: - logger.info("Ratio (delta / pinned): %.2fx", - delta / pinned_bytes) - - if rank == 0 and pinned_bytes > 0: - ratio = delta / pinned_bytes - assert ratio < 3.0, ( - f"Mixed offload CPU cost ({delta / 1024**2:.2f} MB) is " - f"{ratio:.2f}x the pinned buffer ({pinned_bytes / 1024**2:.2f} MB)." - ) - - if rank == 0: - logger.info("PASSED: test_offload_cpu_cost_mixed") - - -def test_pinned_memory_rss_overhead(rank, world_size): - """Isolate: does cudaHostAlloc itself cause 2x RSS overhead?""" - sizes_mb = [8, 16, 32, 64, 128] - - if rank == 0: - logger.info("=" * 70) - logger.info("ISOLATED TEST: PINNED MEMORY RSS OVERHEAD") - logger.info("=" * 70) - - for size_mb in sizes_mb: - numel = size_mb * 1024 * 1024 // 4 # float32 - - # Test 1: pin_memory=True (direct allocation). - gc.collect() - torch.cuda.empty_cache() - rss_before = get_cpu_rss_bytes() - t1 = torch.empty(numel, - dtype=torch.float32, - device="cpu", - pin_memory=True) - rss_after = get_cpu_rss_bytes() - rss_growth_direct = rss_after - rss_before - del t1 - gc.collect() - - # Test 2: .pin_memory() (copy-based). - gc.collect() - torch.cuda.empty_cache() - rss_before2 = get_cpu_rss_bytes() - t2 = torch.empty(numel, dtype=torch.float32, device="cpu").pin_memory() - rss_after2 = get_cpu_rss_bytes() - rss_growth_copy = rss_after2 - rss_before2 - del t2 - gc.collect() - - # Test 3: regular (non-pinned) CPU allocation. - gc.collect() - torch.cuda.empty_cache() - rss_before3 = get_cpu_rss_bytes() - t3 = torch.empty(numel, dtype=torch.float32, device="cpu") - # Touch all pages to ensure RSS reflects actual allocation. - t3.fill_(1.0) - rss_after3 = get_cpu_rss_bytes() - rss_growth_regular = rss_after3 - rss_before3 - del t3 - gc.collect() - - if rank == 0: - logger.info( - "%3d MB: pin_memory=True → RSS +%.1f MB (%.2fx) | " - ".pin_memory() → RSS +%.1f MB (%.2fx) | " - "regular → RSS +%.1f MB (%.2fx)", - size_mb, - rss_growth_direct / 1024**2, - rss_growth_direct / (size_mb * 1024**2) if size_mb > 0 else 0, - rss_growth_copy / 1024**2, - rss_growth_copy / (size_mb * 1024**2) if size_mb > 0 else 0, - rss_growth_regular / 1024**2, - rss_growth_regular / (size_mb * 1024**2) if size_mb > 0 else 0, - ) - - if rank == 0: - logger.info("PASSED: test_pinned_memory_rss_overhead") - - -def main(): - rank, world_size = _setup() - - try: - test_pinned_memory_rss_overhead(rank, world_size) - test_cpu_memory_peak_detailed(rank, world_size) - test_offload_cpu_cost_isolation(rank, world_size) - test_offload_cpu_cost_mixed(rank, world_size) - - if rank == 0: - logger.info("=" * 50) - logger.info("ALL CPU MEMORY PEAK TESTS PASSED") - logger.info("=" * 50) - finally: - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/test/test_cpu_offload.py b/test/test_cpu_offload.py deleted file mode 100644 index 7a2775e8c9f616a93b854559b59110a0dba6b62f..0000000000000000000000000000000000000000 --- a/test/test_cpu_offload.py +++ /dev/null @@ -1,966 +0,0 @@ -"""CPU offloading tests for optimizer states. - -Run with: - torchrun --nproc-per-node=8 --local-ranks-filter=0 test/test_cpu_offload.py - -Tests: - 1. Correctness: turn_on_cpu_offload() produces identical results to no offload - 2. Memory: GPU optimizer state storage is freed after offload - 3. AdamW: moment1/moment2 offloading works correctly -""" - -import copy -import logging - -import pytest -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor, Shard, distribute_tensor - -logger = logging.getLogger(__name__) -logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s") - - -def _setup(): - dist.init_process_group(backend="nccl") - rank = dist.get_rank() - torch.cuda.set_device(rank % torch.cuda.device_count()) - return rank, dist.get_world_size() - - -def _make_mesh(world_size): - return dist.init_device_mesh("cuda", (world_size, ), - mesh_dim_names=("dp", )) - - -def test_correctness(rank, world_size): - """Verify that turn_on_cpu_offload() produces identical parameters as no offload.""" - from optimizer.muon import Muon - from optimizer.newton_schulz import set_ns_compile - - set_ns_compile(False) - torch.manual_seed(42) - - mesh = _make_mesh(world_size) - - dim0, dim1 = 64, 128 - num_params = 4 - num_steps = 3 - - # Pre-generate all data on all ranks (same seed → same values). - full_params = [ - torch.randn(dim0, dim1, device="cuda") for _ in range(num_params) - ] - full_grads = [[ - torch.randn(dim0, dim1, device="cuda") for _ in range(num_params) - ] for _ in range(num_steps)] - - def make_optimizer(cpu_offload): - params, names = [], [] - for i, fp in enumerate(full_params): - dt = distribute_tensor(fp.clone(), mesh, [Shard(0)]) - p = torch.nn.Parameter(dt) - params.append(p) - names.append(f"layer.{i}.weight") - param_groups = [{ - "params": params, - "names": names, - "use_muon": True, - "lr": 0.02, - "weight_decay": 0.01, - "momentum": 0.95, - "nesterov": True, - "ns_steps": 5, - "none_grad": False, - }] - optim = Muon(params=param_groups, chunk_size=2, warmup_step=1) - if cpu_offload: - optim.turn_on_cpu_offload() - return optim, params - - optim_ref, params_ref = make_optimizer(False) - optim_off, params_off = make_optimizer(True) - - for step_idx in range(num_steps): - for i in range(num_params): - g = full_grads[step_idx][i] - params_ref[i].grad = distribute_tensor(g.clone(), mesh, [Shard(0)]) - params_off[i].grad = distribute_tensor(g.clone(), mesh, [Shard(0)]) - - optim_ref.step() - optim_off.step() - - for i in range(num_params): - ref_full = params_ref[i].data.full_tensor() - off_full = params_off[i].data.full_tensor() - torch.testing.assert_close(ref_full, off_full, atol=0, rtol=0) - - if rank == 0: - logger.info("Step %d: correctness OK", step_idx) - - set_ns_compile(True) - if rank == 0: - logger.info("PASSED: test_correctness") - - -def test_memory(rank, world_size): - """Verify that GPU storage is freed after offload.""" - from optimizer.muon import Muon - from optimizer.newton_schulz import set_ns_compile - - set_ns_compile(False) - torch.manual_seed(42) - - mesh = _make_mesh(world_size) - - dim0, dim1 = 512, 1024 - num_params = 8 - - params, names = [], [] - for i in range(num_params): - full = torch.randn(dim0, dim1, device="cuda") - dt = distribute_tensor(full, mesh, [Shard(0)]) - p = torch.nn.Parameter(dt) - p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"), - mesh, [Shard(0)]) - params.append(p) - names.append(f"layer.{i}.weight") - - param_groups = [{ - "params": params, - "names": names, - "use_muon": True, - "lr": 0.02, - "weight_decay": 0.01, - "momentum": 0.95, - "nesterov": True, - "ns_steps": 5, - "none_grad": False, - }] - optim = Muon(params=param_groups, chunk_size=2, warmup_step=1) - optim.turn_on_cpu_offload() - - optim.step() - torch.cuda.synchronize() - - # After step + offload, all momentum buffer GPU storage should be freed. - for p in params: - state = optim.state[p] - if "momentum_buffer" not in state: - continue - buf = state["momentum_buffer"] - local_buf = buf._local_tensor if isinstance(buf, DTensor) else buf - assert local_buf.untyped_storage().size() == 0, ( - f"Expected freed GPU storage after offload, got " - f"{local_buf.untyped_storage().size()} bytes") - - # Verify CPU pool has pinned buffers. - pool = optim._cpu_offload_pool - assert len(pool._managed) > 0, "No tensors tracked by CPU offload pool" - for grp in pool._groups.values(): - assert grp["cpu_flat"].is_pinned(), "CPU buffer must be pinned memory" - - # Run another step to verify reload + compute + offload cycle works. - for p in params: - p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"), - mesh, [Shard(0)]) - optim.step() - torch.cuda.synchronize() - - # Storage should be freed again after second step. - for p in params: - state = optim.state[p] - if "momentum_buffer" not in state: - continue - buf = state["momentum_buffer"] - local_buf = buf._local_tensor if isinstance(buf, DTensor) else buf - assert local_buf.untyped_storage().size() == 0 - - set_ns_compile(True) - if rank == 0: - logger.info("PASSED: test_memory") - - -def test_adamw_offload(rank, world_size): - """Verify AdamW moment1/moment2 are offloaded correctly.""" - from optimizer.muon import Muon - from optimizer.newton_schulz import set_ns_compile - - set_ns_compile(False) - torch.manual_seed(42) - - mesh = _make_mesh(world_size) - - num_steps = 3 - - # Create both Muon (2D) and AdamW (1D) params. - muon_params, muon_names = [], [] - adamw_params, adamw_names = [], [] - - for i in range(4): - full = torch.randn(64, 128, device="cuda") - dt = distribute_tensor(full, mesh, [Shard(0)]) - p = torch.nn.Parameter(dt) - muon_params.append(p) - muon_names.append(f"layer.{i}.weight") - - for i in range(3): - full = torch.randn(128, device="cuda") - dt = distribute_tensor(full, mesh, [Shard(0)]) - p = torch.nn.Parameter(dt) - adamw_params.append(p) - adamw_names.append(f"layer.{i}.bias") - - # Pre-generate grads. - muon_grads = [[torch.randn(64, 128, device="cuda") for _ in range(4)] - for _ in range(num_steps)] - adamw_grads = [[torch.randn(128, device="cuda") for _ in range(3)] - for _ in range(num_steps)] - - def make_optimizer(cpu_offload): - mp = [ - torch.nn.Parameter( - distribute_tensor(p.data.full_tensor().clone(), mesh, - [Shard(0)])) for p in muon_params - ] - ap = [ - torch.nn.Parameter( - distribute_tensor(p.data.full_tensor().clone(), mesh, - [Shard(0)])) for p in adamw_params - ] - param_groups = [ - { - "params": mp, - "names": list(muon_names), - "use_muon": True, - "lr": 0.02, - "weight_decay": 0.01, - "momentum": 0.95, - "nesterov": True, - "ns_steps": 5, - "none_grad": False, - "adamw_betas": (0.9, 0.95), - "adamw_eps": 1e-8, - }, - { - "params": ap, - "use_muon": False, - "lr": 1e-3, - "weight_decay": 0.01, - "adamw_betas": (0.9, 0.95), - "adamw_eps": 1e-8, - }, - ] - optim = Muon(params=param_groups, chunk_size=2, warmup_step=1) - if cpu_offload: - optim.turn_on_cpu_offload() - return optim, mp, ap - - optim_ref, mp_ref, ap_ref = make_optimizer(False) - optim_off, mp_off, ap_off = make_optimizer(True) - - for step_idx in range(num_steps): - for i in range(4): - g = muon_grads[step_idx][i] - mp_ref[i].grad = distribute_tensor(g.clone(), mesh, [Shard(0)]) - mp_off[i].grad = distribute_tensor(g.clone(), mesh, [Shard(0)]) - for i in range(3): - g = adamw_grads[step_idx][i] - ap_ref[i].grad = distribute_tensor(g.clone(), mesh, [Shard(0)]) - ap_off[i].grad = distribute_tensor(g.clone(), mesh, [Shard(0)]) - - optim_ref.step() - optim_off.step() - - # Compare Muon params. - for i in range(4): - ref_full = mp_ref[i].data.full_tensor() - off_full = mp_off[i].data.full_tensor() - torch.testing.assert_close(ref_full, off_full, atol=0, rtol=0) - - # Compare AdamW params. - for i in range(3): - ref_full = ap_ref[i].data.full_tensor() - off_full = ap_off[i].data.full_tensor() - torch.testing.assert_close(ref_full, off_full, atol=0, rtol=0) - - if rank == 0: - logger.info("Step %d: AdamW offload correctness OK", step_idx) - - # Verify AdamW states are offloaded. - for p in ap_off: - state = optim_off.state[p] - for key in ("moment1", "moment2"): - if key not in state: - continue - t = state[key] - local_t = t._local_tensor if isinstance(t, DTensor) else t - assert local_t.untyped_storage().size() == 0, ( - f"AdamW {key} storage not freed after offload") - - set_ns_compile(True) - if rank == 0: - logger.info("PASSED: test_adamw_offload") - - -def test_memory_savings(rank, world_size): - """Measure actual GPU memory difference with and without offload.""" - from optimizer.muon import Muon - from optimizer.newton_schulz import set_ns_compile - - set_ns_compile(False) - - mesh = _make_mesh(world_size) - dim0, dim1 = 1024, 2048 - num_params = 8 - - def run_step(cpu_offload): - torch.cuda.empty_cache() - torch.cuda.reset_peak_memory_stats() - torch.manual_seed(42) - - params, names = [], [] - for i in range(num_params): - full = torch.randn(dim0, dim1, device="cuda") - dt = distribute_tensor(full, mesh, [Shard(0)]) - p = torch.nn.Parameter(dt) - p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"), - mesh, [Shard(0)]) - params.append(p) - names.append(f"layer.{i}.weight") - - param_groups = [{ - "params": params, - "names": names, - "use_muon": True, - "lr": 0.02, - "weight_decay": 0.01, - "momentum": 0.95, - "nesterov": True, - "ns_steps": 5, - "none_grad": False, - }] - optim = Muon(params=param_groups, chunk_size=2, warmup_step=1) - if cpu_offload: - optim.turn_on_cpu_offload() - optim.step() - torch.cuda.synchronize() - - mem = torch.cuda.memory_allocated() - # Clean up to avoid interference. - del optim, params, param_groups - torch.cuda.empty_cache() - return mem - - mem_no_offload = run_step(False) - mem_with_offload = run_step(True) - - if rank == 0: - logger.info("Memory without offload: %.2f MB", - mem_no_offload / 1024**2) - logger.info("Memory with offload: %.2f MB", - mem_with_offload / 1024**2) - saved = mem_no_offload - mem_with_offload - logger.info("Memory saved: %.2f MB", saved / 1024**2) - - assert mem_with_offload < mem_no_offload, ( - f"Expected memory reduction with CPU offload. " - f"Without: {mem_no_offload / 1024**2:.2f} MB, " - f"With: {mem_with_offload / 1024**2:.2f} MB") - - set_ns_compile(True) - if rank == 0: - logger.info("PASSED: test_memory_savings") - - -def test_toggle_correctness(rank, world_size): - """Verify toggling offload on/off between steps produces identical results.""" - from optimizer.muon import Muon - from optimizer.newton_schulz import set_ns_compile - - set_ns_compile(False) - torch.manual_seed(42) - - mesh = _make_mesh(world_size) - - dim0, dim1 = 64, 128 - num_params = 4 - num_steps = 6 - - full_params = [ - torch.randn(dim0, dim1, device="cuda") for _ in range(num_params) - ] - full_grads = [[ - torch.randn(dim0, dim1, device="cuda") for _ in range(num_params) - ] for _ in range(num_steps)] - - def make_optimizer(): - params, names = [], [] - for i, fp in enumerate(full_params): - dt = distribute_tensor(fp.clone(), mesh, [Shard(0)]) - p = torch.nn.Parameter(dt) - params.append(p) - names.append(f"layer.{i}.weight") - param_groups = [{ - "params": params, - "names": names, - "use_muon": True, - "lr": 0.02, - "weight_decay": 0.01, - "momentum": 0.95, - "nesterov": True, - "ns_steps": 5, - "none_grad": False, - }] - optim = Muon(params=param_groups, chunk_size=2, warmup_step=1) - return optim, params - - # Reference: no offload at all. - optim_ref, params_ref = make_optimizer() - - # Toggle: on → step → off → step → on → step ... - optim_toggle, params_toggle = make_optimizer() - - for step_idx in range(num_steps): - # Toggle offload every 2 steps: on for [0,1], off for [2,3], on for [4,5]. - want_on = (step_idx // 2) % 2 == 0 - if want_on and not optim_toggle.cpu_offload: - optim_toggle.turn_on_cpu_offload() - elif not want_on and optim_toggle.cpu_offload: - optim_toggle.turn_off_cpu_offload() - - for i in range(num_params): - g = full_grads[step_idx][i] - params_ref[i].grad = distribute_tensor(g.clone(), mesh, [Shard(0)]) - params_toggle[i].grad = distribute_tensor(g.clone(), mesh, - [Shard(0)]) - - optim_ref.step() - optim_toggle.step() - - for i in range(num_params): - ref_full = params_ref[i].data.full_tensor() - tog_full = params_toggle[i].data.full_tensor() - torch.testing.assert_close(ref_full, tog_full, atol=0, rtol=0) - - if rank == 0: - logger.info( - "Step %d (offload=%s): toggle correctness OK", - step_idx, - optim_toggle.cpu_offload, - ) - - set_ns_compile(True) - if rank == 0: - logger.info("PASSED: test_toggle_correctness") - - -def test_leak(rank, world_size): - """Run many iterations and verify no CPU/GPU memory leak.""" - import os - - from optimizer.muon import Muon - from optimizer.newton_schulz import set_ns_compile - - set_ns_compile(False) - torch.manual_seed(42) - - mesh = _make_mesh(world_size) - - dim0, dim1 = 512, 1024 - num_params = 8 - num_steps = 50 - - params, names = [], [] - for i in range(num_params): - full = torch.randn(dim0, dim1, device="cuda") - dt = distribute_tensor(full, mesh, [Shard(0)]) - p = torch.nn.Parameter(dt) - params.append(p) - names.append(f"layer.{i}.weight") - - param_groups = [{ - "params": params, - "names": names, - "use_muon": True, - "lr": 0.02, - "weight_decay": 0.01, - "momentum": 0.95, - "nesterov": True, - "ns_steps": 5, - "none_grad": False, - }] - optim = Muon(params=param_groups, chunk_size=2, warmup_step=1) - optim.turn_on_cpu_offload() - - def get_cpu_rss_mb(): - """Get current process RSS in MB from /proc/self/statm.""" - with open("/proc/self/statm") as f: - pages = int(f.read().split()[1]) - return pages * os.sysconf("SC_PAGE_SIZE") / (1024**2) - - gpu_after_warmup = None - cpu_after_warmup = None - - for step_idx in range(num_steps): - for p in params: - p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"), - mesh, [Shard(0)]) - - optim.step() - torch.cuda.synchronize() - - gpu_mem = torch.cuda.memory_allocated() - cpu_mem = get_cpu_rss_mb() - - # Record baseline after warmup (step 2 — first step creates states, - # second step does first full offload/reload cycle). - if step_idx == 2: - gpu_after_warmup = gpu_mem - cpu_after_warmup = cpu_mem - - if rank == 0 and step_idx % 10 == 0: - logger.info( - "Step %d: GPU alloc=%.2f MB, CPU RSS=%.2f MB", - step_idx, - gpu_mem / (1024**2), - cpu_mem, - ) - - # Final measurements. - torch.cuda.synchronize() - gpu_final = torch.cuda.memory_allocated() - cpu_final = get_cpu_rss_mb() - - if rank == 0: - logger.info( - "After %d steps: GPU alloc=%.2f MB, CPU RSS=%.2f MB", - num_steps, - gpu_final / (1024**2), - cpu_final, - ) - logger.info( - "Warmup baseline: GPU alloc=%.2f MB, CPU RSS=%.2f MB", - gpu_after_warmup / (1024**2), - cpu_after_warmup, - ) - - # GPU memory should not grow beyond warmup baseline. - assert gpu_final <= gpu_after_warmup, ( - f"GPU memory leak detected! Warmup: {gpu_after_warmup / 1024**2:.2f} MB, " - f"Final: {gpu_final / 1024**2:.2f} MB") - - # CPU RSS should not grow more than 50 MB over warmup (allows for minor - # Python/CUDA runtime overhead but catches real leaks). - cpu_growth = cpu_final - cpu_after_warmup - assert cpu_growth < 50, ( - f"CPU memory leak detected! Growth: {cpu_growth:.2f} MB over " - f"{num_steps - 2} steps (warmup={cpu_after_warmup:.2f} MB, " - f"final={cpu_final:.2f} MB)") - - set_ns_compile(True) - if rank == 0: - logger.info("PASSED: test_leak (GPU stable, CPU growth=%.2f MB)", - cpu_growth) - - -def test_state_dict_save_load(rank, world_size): - """Verify state_dict() works after offload and load_state_dict() resumes correctly. - - Uses torch.distributed.checkpoint (DCP) for serialization, matching - the actual LLM training checkpoint flow. DCP handles DTensors natively - so the roundtrip is bitwise exact. - """ - import shutil - import tempfile - - import torch.distributed.checkpoint as dcp - from optimizer.muon import Muon - from optimizer.newton_schulz import set_ns_compile - - set_ns_compile(False) - torch.manual_seed(42) - - mesh = _make_mesh(world_size) - - dim0, dim1 = 64, 128 - num_muon = 4 - num_adamw = 3 - num_steps = 3 - - # Pre-generate all data. - muon_init = [ - torch.randn(dim0, dim1, device="cuda") for _ in range(num_muon) - ] - adamw_init = [torch.randn(dim1, device="cuda") for _ in range(num_adamw)] - all_grads_muon = [[ - torch.randn(dim0, dim1, device="cuda") for _ in range(num_muon) - ] for _ in range(num_steps * 2)] - all_grads_adamw = [[ - torch.randn(dim1, device="cuda") for _ in range(num_adamw) - ] for _ in range(num_steps * 2)] - - def make_optimizer(cpu_offload): - mp = [ - torch.nn.Parameter( - distribute_tensor(muon_init[i].clone(), mesh, [Shard(0)])) - for i in range(num_muon) - ] - ap = [ - torch.nn.Parameter( - distribute_tensor(adamw_init[i].clone(), mesh, [Shard(0)])) - for i in range(num_adamw) - ] - param_groups = [ - { - "params": mp, - "names": [f"layer.{i}.weight" for i in range(num_muon)], - "use_muon": True, - "lr": 0.02, - "weight_decay": 0.01, - "momentum": 0.95, - "nesterov": True, - "ns_steps": 5, - "none_grad": False, - "adamw_betas": (0.9, 0.95), - "adamw_eps": 1e-8, - }, - { - "params": ap, - "use_muon": False, - "lr": 1e-3, - "weight_decay": 0.01, - "adamw_betas": (0.9, 0.95), - "adamw_eps": 1e-8, - }, - ] - optim = Muon(params=param_groups, chunk_size=2, warmup_step=1) - if cpu_offload: - optim.turn_on_cpu_offload() - return optim, mp, ap - - # --- Run one optimizer for first half, save state, then create TWO - # fresh optimizers: ref loads via deepcopy, resumed loads via DCP. - # Both are fresh → same internal cache state → isolates DCP fidelity. - optim_off, mp_off, ap_off = make_optimizer(True) - - for step_idx in range(num_steps): - for i in range(num_muon): - mp_off[i].grad = distribute_tensor( - all_grads_muon[step_idx][i].clone(), mesh, [Shard(0)]) - for i in range(num_adamw): - ap_off[i].grad = distribute_tensor( - all_grads_adamw[step_idx][i].clone(), mesh, [Shard(0)]) - optim_off.step() - - with pytest.raises( - RuntimeError, - match="turn_off_cpu_offload\\(\\) before checkpoint save"): - optim_off.state_dict() - - optim_off.turn_off_cpu_offload() - sd_off = optim_off.state_dict() - - # Verify state tensors are NOT empty in the state_dict. - for param_states in sd_off["state"].values(): - for key, val in param_states.items(): - if isinstance(val, torch.Tensor) and val.is_floating_point(): - assert val.untyped_storage().size() > 0, ( - f"state_dict() returned empty storage for key '{key}' — " - f"offload reload is broken") - - if rank == 0: - logger.info("state_dict() contains valid (non-empty) tensors") - - # Save state tensors via DCP (matches real LLM training checkpoint flow). - # Flatten state tensors with string keys for DCP compatibility. - flat_state = {} - for param_idx, param_state in sd_off["state"].items(): - for key, val in param_state.items(): - if isinstance(val, torch.Tensor): - flat_state[f"state.{param_idx}.{key}"] = val - - # All ranks must use the same checkpoint directory. - if rank == 0: - ckpt_dir = tempfile.mkdtemp(prefix="cpu_offload_test_") - else: - ckpt_dir = "" - ckpt_dir_list = [ckpt_dir] - dist.broadcast_object_list(ckpt_dir_list, src=0) - ckpt_dir = ckpt_dir_list[0] - try: - dcp.save(flat_state, checkpoint_id=ckpt_dir) - dist.barrier() - - if rank == 0: - logger.info("DCP save completed to %s", ckpt_dir) - - # --- Reference: fresh optimizer, load via deepcopy (no serialization). - optim_ref, mp_ref, ap_ref = make_optimizer(True) - for i in range(num_muon): - mp_ref[i].data = mp_off[i].data.clone() - for i in range(num_adamw): - ap_ref[i].data = ap_off[i].data.clone() - with pytest.raises( - RuntimeError, - match="turn_off_cpu_offload\\(\\) before checkpoint load"): - optim_ref.load_state_dict(copy.deepcopy(sd_off)) - optim_ref.turn_off_cpu_offload() - optim_ref.load_state_dict(copy.deepcopy(sd_off)) - optim_ref.turn_on_cpu_offload() - - # --- Resumed: fresh optimizer, load via DCP. - optim_resumed, mp_resumed, ap_resumed = make_optimizer(True) - for i in range(num_muon): - mp_resumed[i].data = mp_off[i].data.clone() - for i in range(num_adamw): - ap_resumed[i].data = ap_off[i].data.clone() - - flat_target = {k: torch.zeros_like(v) for k, v in flat_state.items()} - dcp.load(flat_target, checkpoint_id=ckpt_dir) - dist.barrier() - - sd_loaded = copy.deepcopy(sd_off) - for param_idx, param_state in sd_loaded["state"].items(): - for key in list(param_state.keys()): - flat_key = f"state.{param_idx}.{key}" - if flat_key in flat_target: - param_state[key] = flat_target[flat_key] - with pytest.raises( - RuntimeError, - match="turn_off_cpu_offload\\(\\) before checkpoint load"): - optim_resumed.load_state_dict(copy.deepcopy(sd_loaded)) - optim_resumed.turn_off_cpu_offload() - optim_resumed.load_state_dict(sd_loaded) - optim_resumed.turn_on_cpu_offload() - - if rank == 0: - logger.info("Both optimizers loaded, starting comparison steps") - - finally: - dist.barrier() - if rank == 0: - shutil.rmtree(ckpt_dir, ignore_errors=True) - - # Second half: reference continues, resumed uses loaded state. - for step_idx in range(num_steps, num_steps * 2): - for i in range(num_muon): - g = all_grads_muon[step_idx][i] - mp_ref[i].grad = distribute_tensor(g.clone(), mesh, [Shard(0)]) - mp_resumed[i].grad = distribute_tensor(g.clone(), mesh, [Shard(0)]) - for i in range(num_adamw): - g = all_grads_adamw[step_idx][i] - ap_ref[i].grad = distribute_tensor(g.clone(), mesh, [Shard(0)]) - ap_resumed[i].grad = distribute_tensor(g.clone(), mesh, [Shard(0)]) - optim_ref.step() - optim_resumed.step() - - # Compare final params: bitwise exact (DCP preserves DTensor identity). - for i in range(num_muon): - ref_full = mp_ref[i].data.full_tensor() - res_full = mp_resumed[i].data.full_tensor() - torch.testing.assert_close(ref_full, res_full, atol=0, rtol=0) - - for i in range(num_adamw): - ref_full = ap_ref[i].data.full_tensor() - res_full = ap_resumed[i].data.full_tensor() - torch.testing.assert_close(ref_full, res_full, atol=0, rtol=0) - - # Verify offload is active on the resumed optimizer. - for p in mp_resumed: - state = optim_resumed.state[p] - if "momentum_buffer" in state: - buf = state["momentum_buffer"] - local_buf = buf._local_tensor if isinstance(buf, DTensor) else buf - assert local_buf.untyped_storage().size() == 0, ( - "Resumed optimizer should have offloaded state after step()") - - set_ns_compile(True) - if rank == 0: - logger.info("PASSED: test_state_dict_save_load") - - -def test_checkpoint_memory(rank, world_size): - """Verify checkpoint APIs require offload to be disabled explicitly.""" - from optimizer.muon import Muon - from optimizer.newton_schulz import set_ns_compile - - set_ns_compile(False) - torch.manual_seed(42) - - mesh = _make_mesh(world_size) - - dim0, dim1 = 512, 1024 - num_params = 8 - - params, names = [], [] - for i in range(num_params): - full = torch.randn(dim0, dim1, device="cuda") - dt = distribute_tensor(full, mesh, [Shard(0)]) - p = torch.nn.Parameter(dt) - p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"), - mesh, [Shard(0)]) - params.append(p) - names.append(f"layer.{i}.weight") - - param_groups = [{ - "params": params, - "names": names, - "use_muon": True, - "lr": 0.02, - "weight_decay": 0.01, - "momentum": 0.95, - "nesterov": True, - "ns_steps": 5, - "none_grad": False, - }] - optim = Muon(params=param_groups, chunk_size=2, warmup_step=1) - optim.turn_on_cpu_offload() - - # Step 1: run a step so offload initializes. - optim.step() - torch.cuda.synchronize() - - mem_after_step = torch.cuda.memory_allocated() - - # Calculate expected state size (momentum buffers, bf16). - state_bytes = 0 - for p in params: - state = optim.state[p] - if "momentum_buffer" in state: - buf = state["momentum_buffer"] - local = buf._local_tensor if isinstance(buf, DTensor) else buf - # Storage is freed, so use the tracked size. - state_bytes += optim._cpu_offload_pool._storage_nbytes[id(buf)] - - if rank == 0: - logger.info( - "After step (offloaded): GPU alloc=%.2f MB, expected state size=%.2f MB", - mem_after_step / 1024**2, - state_bytes / 1024**2, - ) - - with pytest.raises( - RuntimeError, - match="turn_off_cpu_offload\\(\\) before checkpoint save"): - optim.state_dict() - - optim.turn_off_cpu_offload() - torch.cuda.synchronize() - mem_after_turn_off = torch.cuda.memory_allocated() - sd_for_load = copy.deepcopy(optim.state_dict()) - - if rank == 0: - logger.info( - "After turn_off_cpu_offload: GPU alloc=%.2f MB", - mem_after_turn_off / 1024**2, - ) - - assert mem_after_turn_off > mem_after_step, ( - f"turn_off_cpu_offload() should reload states to GPU. " - f"After offload: {mem_after_step / 1024**2:.2f} MB, " - f"After turn_off: {mem_after_turn_off / 1024**2:.2f} MB") - - optim.turn_on_cpu_offload() - torch.cuda.synchronize() - mem_after_turn_on = torch.cuda.memory_allocated() - - if rank == 0: - logger.info("After turn_on_cpu_offload: GPU alloc=%.2f MB", - mem_after_turn_on / 1024**2) - - assert mem_after_turn_on <= mem_after_step + 4 * 1024 * 1024, ( - f"turn_on_cpu_offload() should return memory to offloaded level. " - f"Expected <= {mem_after_step / 1024**2:.2f} MB (+4 MB tolerance), " - f"got {mem_after_turn_on / 1024**2:.2f} MB") - - for p in params: - p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"), - mesh, [Shard(0)]) - optim.step() - torch.cuda.synchronize() - - mem_after_next_step = torch.cuda.memory_allocated() - - if rank == 0: - logger.info( - "After next step (re-offloaded): GPU alloc=%.2f MB", - mem_after_next_step / 1024**2, - ) - - # Allow 4 MB tolerance for CUDA allocator fragmentation. - assert mem_after_next_step <= mem_after_step + 4 * 1024 * 1024, ( - f"Memory should return to offloaded level after step(). " - f"Expected <= {mem_after_step / 1024**2:.2f} MB (+4 MB tolerance), " - f"got {mem_after_next_step / 1024**2:.2f} MB") - - with pytest.raises( - RuntimeError, - match="turn_off_cpu_offload\\(\\) before checkpoint load"): - optim.load_state_dict(copy.deepcopy(sd_for_load)) - - optim.turn_off_cpu_offload() - optim.load_state_dict(sd_for_load) - torch.cuda.synchronize() - - mem_after_load = torch.cuda.memory_allocated() - - if rank == 0: - logger.info( - "After load_state_dict with offload disabled: GPU alloc=%.2f MB", - mem_after_load / 1024**2, - ) - - assert mem_after_load >= mem_after_turn_off, ( - "Loaded optimizer state should stay on GPU while offload is disabled") - - optim.turn_on_cpu_offload() - torch.cuda.synchronize() - - pool = optim._cpu_offload_pool - assert pool._initialized, ( - "Offload pool should be initialized after re-enabling offload") - for grp in pool._groups.values(): - assert grp["cpu_flat"].is_pinned(), "CPU buffer must be pinned" - - # Step 5: verify the loaded optimizer can still step correctly. - for p in params: - p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"), - mesh, [Shard(0)]) - optim.step() - torch.cuda.synchronize() - - mem_final = torch.cuda.memory_allocated() - assert mem_final <= mem_after_step + 4 * 1024 * 1024, ( - f"Final memory should be at offloaded level. " - f"Expected <= {mem_after_step / 1024**2:.2f} MB (+4 MB tolerance), " - f"got {mem_final / 1024**2:.2f} MB") - - set_ns_compile(True) - if rank == 0: - logger.info("PASSED: test_checkpoint_memory") - - -def main(): - rank, world_size = _setup() - - try: - test_correctness(rank, world_size) - test_memory(rank, world_size) - test_adamw_offload(rank, world_size) - test_memory_savings(rank, world_size) - test_toggle_correctness(rank, world_size) - test_leak(rank, world_size) - test_state_dict_save_load(rank, world_size) - test_checkpoint_memory(rank, world_size) - - if rank == 0: - logger.info("=" * 50) - logger.info("ALL CPU OFFLOAD TESTS PASSED") - logger.info("=" * 50) - finally: - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/test/test_muon.py b/test/test_muon.py deleted file mode 100644 index 256cd34480c125c953e92e38a5adb9fcd79aa712..0000000000000000000000000000000000000000 --- a/test/test_muon.py +++ /dev/null @@ -1,513 +0,0 @@ -import copy -import logging -import time -from contextlib import nullcontext - -import pytest -import torch -import torch.distributed as dist -from optimizer.muon import Muon, get_default_muon_param_groups -from optimizer.newton_schulz import set_ns_compile -from torch.distributed.tensor import (DTensor, Replicate, Shard, - distribute_tensor) -from torch.profiler import ProfilerActivity, profile - -from .utils import (ParallelDims, _apply_fsdp, assert_params_equal, - parallelize_motif, parallelize_qk_logits) - -logger = logging.getLogger(__name__) -logging.basicConfig(level=logging.INFO) - - -def apply_muon_step( - model: torch.nn.Module, - parallel_dims: ParallelDims | None, - grads: list[torch.Tensor], - warmup_step: int, - chunk_size: int, - qk_logits: dict[int, torch.Tensor] | None = None, - use_distributed_muon: bool = False, - measure_perf: bool = False, - do_profile: bool = False, - test_name: str | None = None, -) -> tuple[torch.nn.Module, tuple[float, float] | None]: - """ apply single Muon step with optional QK clipping """ - - # 1. Apply gradients to model parameters - assert len(grads) == len(list(model.parameters())) - for grad, param in zip(grads, model.parameters()): - grad = grad.to(param.device) - if isinstance(param.data, DTensor): - unsharded_grad = DTensor.from_local( - grad, - device_mesh=param.data.device_mesh, - placements=[Replicate()] * param.data.device_mesh.ndim, - ) - sharded_grad = unsharded_grad.redistribute( - device_mesh=param.data.device_mesh, - placements=param.data.placements) - param.grad = sharded_grad - else: - param.grad = grad - - # 2. Setup Muon optimizer - params = get_default_muon_param_groups(model) - clip_config = dict({ - "q_indices": - list(range(model.config.num_attention_heads)), - "k_indices": - list(range(model.config.num_attention_heads)), - "head_dim": - model.config.hidden_size // model.config.num_attention_heads, - "threshold": - 0.5 - }) - optim = Muon( - params=params, - clip_config=clip_config if qk_logits is not None else None, - none_grad=False, - warmup_step=warmup_step, - chunk_size=chunk_size, - use_distributed_muon=use_distributed_muon, - ) - - optim.step(qk_logits=qk_logits) - - timing_result: tuple[float, float] | None = None - - if measure_perf: - # extra warm up - optim.step(qk_logits=qk_logits) - - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - - torch.cuda.reset_peak_memory_stats() - start.record() - num_iters = 20 - - if do_profile: - context = profile( - activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], - record_shapes=True) - else: - context = nullcontext() - - with context as prof: - for _i in range(num_iters): - optim.step(qk_logits=qk_logits) - - end.record() - end.synchronize() - - if prof is not None: - date = time.strftime("%Y%m%d_%H%M%S", time.localtime()) - name = test_name or "trace" - rank = dist.get_rank() - prof.export_chrome_trace(f"{name}_{date}_rank{rank}.json") - - peak_memory = torch.cuda.max_memory_allocated() - - elapsed_time_ms = start.elapsed_time(end) / num_iters - - timing_result = (elapsed_time_ms, peak_memory) - - return model, timing_result - - -@pytest.fixture(scope="session") -def sequential_muon_result( - skip_verify, # from conftest.py - inputs # from conftest.py -) -> dict[tuple[bool, bool], torch.nn.Module]: - """Run Muon optimizer to sequential model for baseline results. - - Returns dict keyed by ``(apply_qk_clip, use_compile)``. - """ - if skip_verify: - logger.info("Skipping verification tests as per user request") - return None - - model, grads, qk_logits = inputs - results: dict[tuple[bool, bool], torch.nn.Module] = {} - - for use_compile in [False, True]: - set_ns_compile(use_compile) - - results[(False, use_compile)] = apply_muon_step( - model=copy.deepcopy(model).cuda(), - parallel_dims=None, - grads=grads, - warmup_step=-1, - chunk_size=-1, - qk_logits=None, - )[0].cpu() - - results[(True, use_compile)] = apply_muon_step( - model=copy.deepcopy(model).cuda(), - parallel_dims=None, - grads=grads, - warmup_step=-1, - chunk_size=-1, - qk_logits=qk_logits, - )[0].cpu() - - set_ns_compile(True) # restore default - return results - - -OVERLAP_STEPS = [5] -CHUNK_SIZES = [2] - - -@pytest.mark.parametrize("parallel_dims", [ - pytest.param(ParallelDims(8, 1, 1), id="base"), - pytest.param(ParallelDims(1, 8, 1), id="fsdp"), - pytest.param(ParallelDims(2, 4, 1), id="hsdp"), - pytest.param(ParallelDims(1, 1, 8), id="tp"), - pytest.param(ParallelDims(2, 2, 2), id="hsdp+tp"), - pytest.param(ParallelDims(1, 2, 4), id="fsdp+tp"), -]) -@pytest.mark.parametrize("apply_qk_clip", [False, True]) -@pytest.mark.parametrize("use_distributed_muon", [False]) -@pytest.mark.parametrize("warmup_step", OVERLAP_STEPS) -@pytest.mark.parametrize("chunk_size", CHUNK_SIZES) -@pytest.mark.parametrize("use_compile", [False, True]) -def test_parallel_muon( - request, - sequential_muon_result: dict[tuple[bool, bool], torch.nn.Module], - parallel_dims: ParallelDims, - apply_qk_clip: bool, - use_distributed_muon: bool, - warmup_step: int, - chunk_size: int, - use_compile: bool, - inputs: tuple[torch.nn.Module, list[torch.Tensor], - dict[int, torch.Tensor]], # from conftest.py - measure_perf, # from conftest.py - do_profile, # from conftest.py -) -> None: - if use_distributed_muon and chunk_size != CHUNK_SIZES[0]: - pytest.skip("Distributed Muon does not effected by chunk size") - if use_distributed_muon and warmup_step != OVERLAP_STEPS[0]: - pytest.skip("Distributed Muon does not effected by warmup step") - - set_ns_compile(use_compile) - - model, grads, qk_logits = inputs - - if not apply_qk_clip: - qk_logits = None - - # Deepcopy the model to avoid in-place modification - model = copy.deepcopy(model).cuda() - - parallelized_model = parallelize_motif(model, parallel_dims) - - if qk_logits is not None: - # Deepcopy the qk logits to avoid in-place modification - qk_logits = copy.deepcopy(qk_logits) - qk_logits = parallelize_qk_logits(qk_logits, parallel_dims) - - parallelized_model, timing_result = apply_muon_step( - model=parallelized_model, - parallel_dims=parallel_dims, - grads=grads, - warmup_step=warmup_step, - chunk_size=chunk_size, - qk_logits=qk_logits, - use_distributed_muon=use_distributed_muon, - measure_perf=measure_perf, - do_profile=do_profile, - test_name=request.node.name, - ) - - if measure_perf: - assert timing_result is not None - avg_time_ms, peak_memory = timing_result - logger.info( - f"\nParallel dims: {parallel_dims}, " - f"\nUse distributed Muon: {use_distributed_muon}, " - f"\nApply QK clip: {apply_qk_clip} => " - f"\nChunk Size, Warmup Step, Avg Time (ms), Peak Memory (MB):" - f"\n{chunk_size}, {warmup_step}, {avg_time_ms:.2f}, {peak_memory / (1024**2):.2f}," - ) - - if sequential_muon_result is None: - logger.info("Skipping correctness check as sequential result is None") - elif measure_perf: - logger.info("Skipping correctness check as timing is enabled") - else: - atol = 1e-5 if use_compile else 0 - rtol = 1e-2 if use_compile else 0 - assert_params_equal(parallelized_model, - sequential_muon_result[(apply_qk_clip, - use_compile)], - atol=atol, - rtol=rtol) - - -def test_parallel_muon_empty_shard(init_dist): - """Regression: parallel Muon must handle chunks where some ranks have - empty local shards (dim-0 < world_size). - - With 8-way Shard(0) and dim-0 of size 4, ranks 4-7 get 0-element local - shards. Previously ``_launch_gather`` hit ``assert total_send > 0``. - """ - rank = dist.get_rank() - world_size = dist.get_world_size() - mesh = dist.init_device_mesh("cuda", (world_size, ), - mesh_dim_names=("dp", )) - - set_ns_compile(False) - - # dim-0 = 4 < 8 ranks → ranks 4-7 have empty local shards with Shard(0) - small_dim = 4 - num_params = 4 - torch.manual_seed(42) - - muon_params = [] - muon_names = [] - for i in range(num_params): - full = torch.randn(small_dim, 64, device="cuda") - dt = distribute_tensor(full, mesh, [Shard(0)]) - p = torch.nn.Parameter(dt) - grad_full = torch.randn(small_dim, 64, device="cuda") - p.grad = distribute_tensor(grad_full, mesh, [Shard(0)]) - muon_params.append(p) - muon_names.append(f"layer.{i}.weight") - - param_groups = [{ - "params": muon_params, - "names": muon_names, - "use_muon": True, - "lr": 0.02, - "weight_decay": 0.01, - "momentum": 0.95, - "nesterov": True, - "ns_steps": 5, - "none_grad": False, - }] - - optim = Muon(params=param_groups, chunk_size=1, warmup_step=0) - # Must not raise AssertionError: total_send > 0 - optim.step() - - # Run a second step to verify cached path also works - for p in muon_params: - grad_full = torch.randn(small_dim, 64, device="cuda") - p.grad = distribute_tensor(grad_full, mesh, [Shard(0)]) - optim.step() - - set_ns_compile(True) - logger.info("test_parallel_muon_empty_shard PASSED (rank %d)", rank) - - -@pytest.mark.parametrize("uneven_dim", [ - pytest.param(33, id="33"), - pytest.param(19, id="19"), - pytest.param(11, id="11"), -]) -def test_parallel_muon_uneven_shard(init_dist, uneven_dim): - """Test that parallel Muon produces correct results when parameter - dimensions are not evenly divisible by the number of shard ranks. - - For example, dim=33 with 8 ranks gives 7 ranks with 4 rows and - 1 rank with 5 rows. This exercises the remainder-handling logic - in ``get_slices_of_dtensor`` and the all-to-all pipeline. - """ - rank = dist.get_rank() - world_size = dist.get_world_size() - mesh = dist.init_device_mesh("cuda", (world_size, ), - mesh_dim_names=("dp", )) - - set_ns_compile(False) - torch.manual_seed(42) - - other_dim = 64 - num_params = 3 - - # --- Build sharded params + grads --- - muon_params = [] - muon_names = [] - full_params_snapshot = [] - full_grads = [] - - for i in range(num_params): - full = torch.randn(uneven_dim, other_dim, device="cuda") - full_params_snapshot.append(full.clone()) - dt = distribute_tensor(full, mesh, [Shard(0)]) - p = torch.nn.Parameter(dt) - grad_full = torch.randn(uneven_dim, other_dim, device="cuda") - full_grads.append(grad_full.clone()) - p.grad = distribute_tensor(grad_full, mesh, [Shard(0)]) - muon_params.append(p) - muon_names.append(f"layer.{i}.weight") - - # --- Parallel path (all2all pipeline) --- - param_groups_par = [{ - "params": muon_params, - "names": muon_names, - "use_muon": True, - "lr": 0.02, - "weight_decay": 0.01, - "momentum": 0.95, - "nesterov": True, - "ns_steps": 5, - "none_grad": False, - }] - optim_par = Muon(params=param_groups_par, chunk_size=1, warmup_step=0) - optim_par.step() - - # --- Sequential baseline (base path, no sharding) --- - seq_params = [] - seq_names = [] - for i in range(num_params): - p = torch.nn.Parameter(full_params_snapshot[i].clone()) - p.grad = full_grads[i].clone() - seq_params.append(p) - seq_names.append(f"layer.{i}.weight") - - param_groups_seq = [{ - "params": seq_params, - "names": seq_names, - "use_muon": True, - "lr": 0.02, - "weight_decay": 0.01, - "momentum": 0.95, - "nesterov": True, - "ns_steps": 5, - "none_grad": False, - }] - optim_seq = Muon(params=param_groups_seq) - optim_seq.step() - - # --- Compare: parallel result (gathered) must match sequential --- - for i in range(num_params): - par_full = muon_params[i].data.full_tensor() - seq_full = seq_params[i].data - torch.testing.assert_close(par_full, seq_full, atol=0, rtol=0) - - set_ns_compile(True) - logger.info("test_parallel_muon_uneven_shard (dim=%d) PASSED (rank %d)", - uneven_dim, rank) - - -def test_pp_dp_replicate_no_deadlock(init_dist, inputs): - """PP regression test using real Motif model. - - PP=2, dp_replicate=2, dp_shard=2 on 8 GPUs. Splits the - Motif-2.6B-4layer model across 2 pipeline stages following the - torchtitan pattern (deep copy → delete non-stage layers → per-stage - FSDP). Each stage independently runs Muon optimizer and the result - is verified against a sequential baseline (atol=0, rtol=0). - - Without use_local_synchronization=True in construct_shard_mesh(), - different stages would deadlock on dist.new_group() because they - call it for different parameters. - """ - import re - - import torch.nn as nn - from optimizer.distributed.utils import _ranks_to_dist_cache - - rank = dist.get_rank() - assert dist.get_world_size() == 8 - - set_ns_compile(False) - _ranks_to_dist_cache.clear() - - model_orig, grads_orig, _ = inputs - - # Build name→grad mapping from original model - grad_dict = { - name: grad - for (name, _), grad in zip(model_orig.named_parameters(), grads_orig) - } - - # Full mesh: PP=2, dp_replicate=2, dp_shard=2 - full_mesh = dist.init_device_mesh( - "cuda", - (2, 2, 2), - mesh_dim_names=("pp", "dp_replicate", "dp_shard"), - ) - dp_mesh = full_mesh["dp_replicate", "dp_shard"] - pp_rank = full_mesh.get_local_rank("pp") - - # -- Helpers ---------------------------------------------------------- - def _split_motif(model): - """Split Motif model per PP stage (torchtitan pattern). - - Stage 0: embed_tokens + layers[0:2] - Stage 1: layers[2:4] + norm + output - Non-stage components replaced with nn.Identity (no params). - """ - all_layers = list(model.model.layers) - if pp_rank == 0: - model.model.layers = nn.ModuleList(all_layers[:2]) - model.model.norm = nn.Identity() - if hasattr(model, "output"): - model.output = nn.Identity() - if hasattr(model, "lm_head"): - model.lm_head = nn.Identity() - else: - model.model.layers = nn.ModuleList(all_layers[2:]) - model.model.embed_tokens = nn.Identity() - return model - - layer_offset = 0 if pp_rank == 0 else 2 - - def _remap(name): - """Map stage param name → original param name (layer index offset). - - Also handles weight tying: Motif ties lm_head.weight to - model.embed_tokens.weight, so named_parameters() only lists the - latter. After stage-split, stage 1 loses embed_tokens but keeps - lm_head, so we remap it back. - """ - # Weight tying: lm_head.weight ↔ model.embed_tokens.weight - if name == "lm_head.weight": - return "model.embed_tokens.weight" - - if layer_offset == 0: - return name - - def _replace(m): - return f"layers.{int(m.group(1)) + layer_offset}." - - return re.sub(r"layers\.(\d+)\.", _replace, name) - - def _stage_grads(model): - """Build grads list aligned with stage model parameters.""" - return [grad_dict[_remap(n)] for n, _ in model.named_parameters()] - - # -- Parallel path: split → FSDP → Muon step ------------------------- - par_model = _split_motif(copy.deepcopy(model_orig).cuda()) - _apply_fsdp(par_model, dp_mesh) - par_model, _ = apply_muon_step( - model=par_model, - parallel_dims=None, - grads=_stage_grads(par_model), - warmup_step=5, - chunk_size=2, - qk_logits=None, - ) - - # -- Sequential baseline: split → no FSDP → base Muon ---------------- - seq_model = _split_motif(copy.deepcopy(model_orig).cuda()) - seq_model, _ = apply_muon_step( - model=seq_model, - parallel_dims=None, - grads=_stage_grads(seq_model), - warmup_step=-1, - chunk_size=-1, - qk_logits=None, - ) - - # Correctness: parallel must match sequential exactly - assert_params_equal(par_model, seq_model, atol=0, rtol=0) - - set_ns_compile(True) - logger.info( - "test_pp_dp_replicate_no_deadlock PASSED (rank %d, pp_rank %d)", rank, - pp_rank) diff --git a/test/test_muon_moe.py b/test/test_muon_moe.py deleted file mode 100644 index 15da956acae58c29e53f721ab9c4f183cb7be54c..0000000000000000000000000000000000000000 --- a/test/test_muon_moe.py +++ /dev/null @@ -1,518 +0,0 @@ -import copy -import logging -import time -from contextlib import nullcontext - -import pytest -import torch -import torch.distributed as dist -from optimizer.muon import Muon, get_default_muon_param_groups -from torch.distributed.tensor import (DTensor, Replicate, Shard, - distribute_tensor) -from torch.profiler import ProfilerActivity, profile - -from .utils import ParallelDims, assert_params_equal, parallelize_llama4 - -logger = logging.getLogger(__name__) -logging.basicConfig(level=logging.INFO) - - -def _apply_grads(model, grads): - """Apply gradients to model parameters (with DTensor redistribute).""" - for grad, param in zip(grads, model.parameters()): - grad = grad.to(param.device) - if isinstance(param.data, DTensor): - unsharded_grad = DTensor.from_local( - grad, - device_mesh=param.data.device_mesh, - placements=[Replicate()] * param.data.device_mesh.ndim, - ) - param.grad = unsharded_grad.redistribute( - device_mesh=param.data.device_mesh, - placements=param.data.placements) - else: - param.grad = grad - - -def _restore_grads(model, saved_grads): - """Restore previously saved grads (no redistribute, just reassign).""" - for param, g in zip(model.parameters(), saved_grads): - param.grad = g - - -def apply_muon_step_moe( - model: torch.nn.Module, - parallel_dims: ParallelDims | None, - grads: list[torch.Tensor], - warmup_step: int, - chunk_size: int, - use_distributed_muon: bool = False, - measure_perf: bool = False, - do_profile: bool = False, - test_name: str | None = None, -) -> tuple[torch.nn.Module, tuple[float, float] | None]: - """Apply a single Muon step to an MoE model (no QK clipping).""" - - assert len(grads) == len(list(model.parameters())) - _apply_grads(model, grads) - - params = get_default_muon_param_groups(model, expert_keys=["experts"]) - optim = Muon( - params=params, - clip_config=None, - none_grad=False, - warmup_step=warmup_step, - chunk_size=chunk_size, - use_distributed_muon=use_distributed_muon, - expert_keys=["experts"], - ) - - # Save sharded grads for re-use before step clears 3D grads. - saved_grads = [p.grad for p in model.parameters()] - - optim.step() - - # Second step to exercise expert expand cache hot path. - _restore_grads(model, saved_grads) - optim.step() - - timing_result: tuple[float, float] | None = None - - if measure_perf: - # extra warm up - _restore_grads(model, saved_grads) - optim.step() - - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - - torch.cuda.reset_peak_memory_stats() - start.record() - num_iters = 20 - - if do_profile: - context = profile( - activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], - record_shapes=True) - else: - context = nullcontext() - - with context as prof: - for _i in range(num_iters): - _restore_grads(model, saved_grads) - optim.step() - - end.record() - end.synchronize() - - if prof is not None: - date = time.strftime("%Y%m%d_%H%M%S", time.localtime()) - name = test_name or "trace_moe" - rank = dist.get_rank() - prof.export_chrome_trace(f"{name}_{date}_rank{rank}.json") - - peak_memory = torch.cuda.max_memory_allocated() - elapsed_time_ms = start.elapsed_time(end) / num_iters - timing_result = (elapsed_time_ms, peak_memory) - - return model, timing_result - - -@pytest.fixture(scope="session") -def sequential_moe_result( - skip_verify, - moe_inputs, -) -> torch.nn.Module | None: - """Run Muon optimizer on sequential MoE model for baseline.""" - if skip_verify: - logger.info("Skipping verification tests as per user request") - return None - - model, grads = moe_inputs - - result, _ = apply_muon_step_moe( - model=copy.deepcopy(model).cuda(), - parallel_dims=None, - grads=grads, - warmup_step=-1, - chunk_size=-1, - ) - result = result.cpu() - - return result - - -OVERLAP_STEPS = [5] -CHUNK_SIZES = [2] - - -@pytest.mark.parametrize( - "parallel_dims", - [ - # --- No EP (non-expert only) --- - pytest.param(ParallelDims(8, 1, 1), id="dp8"), - pytest.param(ParallelDims(1, 8, 1), id="fsdp8"), - pytest.param(ParallelDims(2, 4, 1), id="hsdp2x4"), - # --- EP configs --- - # naming: fsdp{dp_shard}_ep{ep} where dp_shard = dp_shard_mod_ep * ep - # dp_shard_mod_ep (= expert FSDP) = dp_shard_degree in our ParallelDims - pytest.param(ParallelDims(1, 1, 1, ep_degree=8), id="fsdp8_ep8"), - pytest.param(ParallelDims(1, 4, 1, ep_degree=2), id="fsdp8_ep2"), - pytest.param(ParallelDims(1, 2, 1, ep_degree=4), id="fsdp8_ep4"), - pytest.param(ParallelDims(2, 2, 1, ep_degree=2), id="hsdp_ep2"), - ]) -@pytest.mark.parametrize("use_distributed_muon", [False]) -@pytest.mark.parametrize("warmup_step", OVERLAP_STEPS) -@pytest.mark.parametrize("chunk_size", CHUNK_SIZES) -def test_parallel_muon_moe( - request, - sequential_moe_result: torch.nn.Module | None, - parallel_dims: ParallelDims, - use_distributed_muon: bool, - warmup_step: int, - chunk_size: int, - moe_inputs: tuple[torch.nn.Module, list[torch.Tensor]], - measure_perf, - do_profile, -) -> None: - model, grads = moe_inputs - - # Deepcopy the model to avoid in-place modification - model = copy.deepcopy(model).cuda() - - parallelized_model = parallelize_llama4(model, parallel_dims) - - parallelized_model, timing_result = apply_muon_step_moe( - model=parallelized_model, - parallel_dims=parallel_dims, - grads=grads, - warmup_step=warmup_step, - chunk_size=chunk_size, - use_distributed_muon=use_distributed_muon, - measure_perf=measure_perf, - do_profile=do_profile, - test_name=request.node.name, - ) - - if measure_perf: - assert timing_result is not None - avg_time_ms, peak_memory = timing_result - logger.info(f"\nParallel dims: {parallel_dims}, " - f"\nAvg Time (ms): {avg_time_ms:.2f}, " - f"Peak Memory (MB): {peak_memory / (1024**2):.2f}") - - if sequential_moe_result is None: - logger.info("Skipping correctness check as sequential result is None") - elif measure_perf: - logger.info("Skipping correctness check as timing is enabled") - else: - assert_params_equal(parallelized_model, sequential_moe_result) - - -# --------------------------------------------------------------------------- -# Few-experts tests: num_experts=2, triggers EFSDP Shard(1) mode -# --------------------------------------------------------------------------- - - -@pytest.fixture(scope="session") -def sequential_moe_result_few_experts( - skip_verify, - moe_inputs_few_experts, -) -> torch.nn.Module | None: - """Run Muon optimizer on sequential MoE model (2 experts) for baseline.""" - if skip_verify: - logger.info("Skipping verification tests as per user request") - return None - - model, grads = moe_inputs_few_experts - - result, _ = apply_muon_step_moe( - model=copy.deepcopy(model).cuda(), - parallel_dims=None, - grads=grads, - warmup_step=-1, - chunk_size=-1, - ) - result = result.cpu() - - return result - - -@pytest.mark.parametrize("parallel_dims", [ - pytest.param(ParallelDims(1, 4, 1, ep_degree=2), id="fsdp8_ep2"), - pytest.param(ParallelDims(2, 2, 1, ep_degree=2), id="hsdp_ep2"), -]) -@pytest.mark.parametrize("use_distributed_muon", [False]) -@pytest.mark.parametrize("warmup_step", OVERLAP_STEPS) -@pytest.mark.parametrize("chunk_size", CHUNK_SIZES) -def test_parallel_muon_moe_few_experts( - request, - sequential_moe_result_few_experts: torch.nn.Module | None, - parallel_dims: ParallelDims, - use_distributed_muon: bool, - warmup_step: int, - chunk_size: int, - moe_inputs_few_experts: tuple[torch.nn.Module, list[torch.Tensor]], - measure_perf, - do_profile, -) -> None: - model, grads = moe_inputs_few_experts - - model = copy.deepcopy(model).cuda() - - parallelized_model = parallelize_llama4(model, parallel_dims) - - parallelized_model, timing_result = apply_muon_step_moe( - model=parallelized_model, - parallel_dims=parallel_dims, - grads=grads, - warmup_step=warmup_step, - chunk_size=chunk_size, - use_distributed_muon=use_distributed_muon, - measure_perf=measure_perf, - do_profile=do_profile, - test_name=request.node.name, - ) - - if measure_perf: - assert timing_result is not None - avg_time_ms, peak_memory = timing_result - logger.info(f"\nParallel dims: {parallel_dims}, " - f"\nAvg Time (ms): {avg_time_ms:.2f}, " - f"Peak Memory (MB): {peak_memory / (1024**2):.2f}") - - if sequential_moe_result_few_experts is None: - logger.info("Skipping correctness check as sequential result is None") - elif measure_perf: - logger.info("Skipping correctness check as timing is enabled") - else: - assert_params_equal(parallelized_model, - sequential_moe_result_few_experts) - - -# --------------------------------------------------------------------------- -# Uneven shard test: mixed expert (3D plain) + non-expert (2D DTensor) -# with dimensions not evenly divisible by shard count. -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize("uneven_dim", [ - pytest.param(33, id="33"), - pytest.param(19, id="19"), -]) -def test_parallel_muon_moe_uneven_shard(init_dist, uneven_dim): - """Test MoE parallel Muon with uneven shard dimensions. - - Mixes non-expert 2D DTensor params (uneven FSDP sharding, parallel - pipeline path) with expert 3D plain-tensor params (batched NS path). - Verifies the combination produces correct results vs sequential baseline. - """ - from optimizer.newton_schulz import set_ns_compile - - rank = dist.get_rank() - world_size = dist.get_world_size() - mesh = dist.init_device_mesh("cuda", (world_size, ), - mesh_dim_names=("dp", )) - - set_ns_compile(False) - torch.manual_seed(42) - - other_dim = 64 - num_experts = 4 - - muon_params = [] - muon_names = [] - full_params = [] - full_grads = [] - - # 2D non-expert params with uneven dims → parallel pipeline - for i in range(2): - full = torch.randn(uneven_dim, other_dim, device="cuda") - full_params.append(full.clone()) - dt = distribute_tensor(full, mesh, [Shard(0)]) - p = torch.nn.Parameter(dt) - g = torch.randn(uneven_dim, other_dim, device="cuda") - full_grads.append(g.clone()) - p.grad = distribute_tensor(g, mesh, [Shard(0)]) - muon_params.append(p) - muon_names.append(f"layers.{i}.weight") - - # 3D expert params (plain tensors) → batched NS path - full = torch.randn(num_experts, uneven_dim, other_dim, device="cuda") - full_params.append(full.clone()) - p = torch.nn.Parameter(full) - g = torch.randn(num_experts, uneven_dim, other_dim, device="cuda") - full_grads.append(g.clone()) - p.grad = g - muon_params.append(p) - muon_names.append("layers.2.experts.w1.weight") - - # --- Parallel path --- - param_groups_par = [{ - "params": muon_params, - "names": muon_names, - "use_muon": True, - "lr": 0.02, - "weight_decay": 0.01, - "momentum": 0.95, - "nesterov": True, - "ns_steps": 5, - "none_grad": False, - }] - optim_par = Muon(params=param_groups_par, - chunk_size=1, - warmup_step=0, - expert_keys=["experts"]) - optim_par.step() - - # --- Sequential baseline --- - seq_params = [] - for fp in full_params: - p = torch.nn.Parameter(fp.clone()) - seq_params.append(p) - - for p, g in zip(seq_params, full_grads): - p.grad = g.clone() - - param_groups_seq = [{ - "params": seq_params, - "names": list(muon_names), - "use_muon": True, - "lr": 0.02, - "weight_decay": 0.01, - "momentum": 0.95, - "nesterov": True, - "ns_steps": 5, - "none_grad": False, - }] - optim_seq = Muon(params=param_groups_seq, expert_keys=["experts"]) - optim_seq.step() - - # --- Compare --- - for i in range(len(muon_params)): - par_data = muon_params[i].data - if isinstance(par_data, DTensor): - par_data = par_data.full_tensor() - torch.testing.assert_close(par_data, - seq_params[i].data, - atol=0, - rtol=0) - - set_ns_compile(True) - logger.info( - "test_parallel_muon_moe_uneven_shard (dim=%d) PASSED (rank %d)", - uneven_dim, rank) - - -def test_pp_dp_replicate_moe_no_deadlock(init_dist, moe_inputs): - """PP regression test using real torchtitan Llama4 MoE model. - - PP=2, dp_replicate=2, dp_shard=2 on 8 GPUs. Splits the Llama4 MoE - model (4 layers, 8 experts) across 2 pipeline stages following the - torchtitan pattern. Uses torchtitan's ``parallelize_llama`` for - realistic FSDP application (same function as real training). - - Each stage independently runs Muon optimizer with expert_keys and - the result is verified against a sequential baseline (atol=0, rtol=0). - - Without use_local_synchronization=True in construct_shard_mesh(), - different stages would deadlock on dist.new_group(). - """ - from optimizer.distributed.utils import _ranks_to_dist_cache - from optimizer.newton_schulz import set_ns_compile - from torchtitan.config import JobConfig - from torchtitan.distributed import ParallelDims as TTParallelDims - from torchtitan.models.llama4.infra.parallelize import parallelize_llama - - rank = dist.get_rank() - assert dist.get_world_size() == 8 - - set_ns_compile(False) - _ranks_to_dist_cache.clear() - - model_orig, grads_orig = moe_inputs - - # Build name→grad mapping from original model - grad_dict = { - name: grad - for (name, _), grad in zip(model_orig.named_parameters(), grads_orig) - } - - # torchtitan ParallelDims with PP=2 (same as real training config) - tt_dims = TTParallelDims( - dp_replicate=2, - dp_shard=2, - cp=1, - tp=1, - pp=2, - ep=1, - etp=1, - world_size=8, - ) - - # Accessing world_mesh triggers build_mesh() (lazy init). - # All ranks participate in init_device_mesh (collective). - pp_rank = tt_dims.world_mesh.get_local_rank("pp") - - job_config = JobConfig() - job_config.training.mixed_precision_param = "float32" - job_config.activation_checkpoint.mode = "none" - job_config.compile.enable = False - job_config.parallelism.disable_loss_parallel = True - - # -- Helpers ---------------------------------------------------------- - def _split_llama4(model): - """Split Llama4 MoE model per PP stage (torchtitan pattern). - - Stage 0: tok_embeddings + layers["0"], ["1"] - Stage 1: layers["2"], ["3"] + norm + output - ModuleDict preserves keys → param names unchanged. - torchtitan model natively supports None modules in forward(). - """ - if pp_rank == 0: - for key in ["2", "3"]: - if key in model.layers: - del model.layers[key] - model.norm = None - model.output = None - else: - for key in ["0", "1"]: - if key in model.layers: - del model.layers[key] - model.tok_embeddings = None - return model - - def _stage_grads(model): - """Build grads list aligned with stage model parameters.""" - return [grad_dict[n] for n, _ in model.named_parameters()] - - # -- Parallel path: split → parallelize_llama → Muon step ------------- - par_model = _split_llama4(copy.deepcopy(model_orig).cuda()) - parallelize_llama(par_model, tt_dims, job_config) - - par_model, _ = apply_muon_step_moe( - model=par_model, - parallel_dims=None, - grads=_stage_grads(par_model), - warmup_step=5, - chunk_size=2, - ) - - # -- Sequential baseline: split → no parallelization → base Muon ------ - seq_model = _split_llama4(copy.deepcopy(model_orig).cuda()) - - seq_model, _ = apply_muon_step_moe( - model=seq_model, - parallel_dims=None, - grads=_stage_grads(seq_model), - warmup_step=-1, - chunk_size=-1, - ) - - # Correctness: parallel must match sequential exactly - assert_params_equal(par_model, seq_model, atol=0, rtol=0) - - set_ns_compile(True) - logger.info( - "test_pp_dp_replicate_moe_no_deadlock PASSED (rank %d, pp_rank %d)", - rank, pp_rank) diff --git a/test/test_normalize_fqn.py b/test/test_normalize_fqn.py deleted file mode 100644 index db020fc4df25905e5ab49e0d33b0a2ee144a0521..0000000000000000000000000000000000000000 --- a/test/test_normalize_fqn.py +++ /dev/null @@ -1,98 +0,0 @@ -"""Unit tests for FQN normalization (no GPU / distributed required).""" - -from optimizer.core import default_is_muon, is_expert_param, normalize_fqn -from optimizer.qk_clip import parse_qk_layer - - -class TestNormalizeFqn: - - def test_passthrough(self): - assert normalize_fqn("model.layers.3.attn.q_proj.weight") == \ - "model.layers.3.attn.q_proj.weight" - - def test_strip_orig_mod(self): - assert normalize_fqn("model._orig_mod.layers.3.attn.q_proj.weight") == \ - "model.layers.3.attn.q_proj.weight" - - def test_strip_checkpoint_wrapped(self): - name = "model.layers.0._checkpoint_wrapped_module.moe.experts.w1.weight" - assert normalize_fqn(name) == \ - "model.layers.0.moe.experts.w1.weight" - - def test_strip_both(self): - name = "model._orig_mod.layers.0._checkpoint_wrapped_module.attn.q_proj.weight" - assert normalize_fqn(name) == \ - "model.layers.0.attn.q_proj.weight" - - def test_strip_nested_orig_mod(self): - name = "_orig_mod._orig_mod.layers.0.mlp.gate_proj.weight" - assert normalize_fqn(name) == \ - "layers.0.mlp.gate_proj.weight" - - -class TestParseQkLayerWithWrappers: - - def test_plain_name(self): - assert parse_qk_layer("model.layers.3.attn.q_proj.weight") == ( - "q_proj", 3) - - def test_orig_mod(self): - assert parse_qk_layer("model._orig_mod.layers.3.attn.wq.weight") == ( - "wq", 3) - - def test_checkpoint_wrapped(self): - name = "model.layers.5._checkpoint_wrapped_module.self_attn.k_proj.weight" - assert parse_qk_layer(name) == ("k_proj", 5) - - def test_both_wrappers(self): - name = "_orig_mod.model._checkpoint_wrapped_module.layers.7.attn.wk.weight" - assert parse_qk_layer(name) == ("wk", 7) - - def test_non_qk_still_none(self): - name = "model._orig_mod.layers.2.attn.v_proj.weight" - assert parse_qk_layer(name) == (None, -1) - - -class TestExpertKeyMatching: - """Verify expert_keys uses component-level matching, not substring.""" - - class FakeParam: - - def __init__(self, ndim): - self.ndim = ndim - - def test_experts_matches(self): - name = "model.layers.0.moe.experts.w1.weight" - assert default_is_muon(name, - self.FakeParam(3), - expert_keys=["experts"]) - - def test_shared_experts_does_not_match(self): - name = "model.layers.0.moe.shared_experts.w1.weight" - # shared_experts has ndim=2, which is muon-eligible on its own. - # But it must NOT be recognized as expert (ndim-1 would make it 1D → False). - assert default_is_muon(name, - self.FakeParam(2), - expert_keys=["experts"]) - - def test_shared_experts_3d_not_treated_as_expert(self): - # 3D shared_experts: if wrongly matched as expert, ndim-1=2 → True (same result). - # Verify by checking that a 2D shared_experts is NOT downgraded to 1D. - name = "model.layers.0.moe.shared_experts.gate_proj.weight" - # 2D param: if expert-matched → ndim-1=1 → False. Must stay True. - assert default_is_muon(name, - self.FakeParam(2), - expert_keys=["experts"]) - - def test_multi_component_key_matches(self): - name = "model.layers.0.moe.experts.w1.weight" - assert is_expert_param(name, expert_keys=["experts.w1"]) - - def test_multi_component_key_no_false_positive(self): - # "experts.w2" should not match "experts.w1" - name = "model.layers.0.moe.experts.w1.weight" - assert not is_expert_param(name, expert_keys=["experts.w2"]) - - def test_multi_component_key_shared_experts(self): - name = "model.layers.0.moe.shared_experts.w1.weight" - assert not is_expert_param(name, expert_keys=["experts.w1"]) diff --git a/test/utils.py b/test/utils.py deleted file mode 100644 index 150e932467fbee99061e325f3e998fa3f1cfbb1d..0000000000000000000000000000000000000000 --- a/test/utils.py +++ /dev/null @@ -1,286 +0,0 @@ -from dataclasses import dataclass - -import torch -import torch.distributed as dist -from torch.distributed.fsdp import fully_shard -from torch.distributed.tensor import DeviceMesh, DTensor, Replicate, Shard -from torch.distributed.tensor.parallel import (ColwiseParallel, - PrepareModuleInput, - RowwiseParallel, - SequenceParallel, - parallelize_module) - - -@dataclass -class ParallelDims: - dp_replicate_degree: int - dp_shard_degree: int - tp_degree: int - ep_degree: int = 1 - - def __str__(self) -> str: - s = (f"dp_replicate-{self.dp_replicate_degree}_" - f"dp_shard-{self.dp_shard_degree}_" - f"tp-{self.tp_degree}") - if self.ep_degree > 1: - s += f"_ep-{self.ep_degree}" - return s - - -def _construct_device_mesh(parallel_dims: ParallelDims) -> DeviceMesh: - """Constructs a DeviceMesh based on the given parallel dimensions. - - Args: - parallel_dims (ParallelDims): The parallelism configuration. - - Returns: - DeviceMesh: The constructed device mesh. - """ - world_size = dist.get_world_size() - expected_devices = (parallel_dims.dp_replicate_degree * - parallel_dims.dp_shard_degree * - parallel_dims.ep_degree * parallel_dims.tp_degree) - if world_size < expected_devices: - raise ValueError( - f"Not enough devices: found {world_size}, " - f"but expected at least {expected_devices}. ({parallel_dims})") - - degrees = [ - parallel_dims.dp_replicate_degree, parallel_dims.dp_shard_degree, - parallel_dims.ep_degree, parallel_dims.tp_degree - ] - dim_names = ["dp_replicate", "dp_shard", "ep", "tp"] - - mesh_shape = [] - mesh_dim_names = [] - for degree, dim_name in zip(degrees, dim_names): - if degree > 1: - mesh_shape.append(degree) - mesh_dim_names.append(dim_name) - - device_mesh = dist.init_device_mesh("cuda", - mesh_shape, - mesh_dim_names=mesh_dim_names) - - return device_mesh - - -def _apply_tp( - model: torch.nn.Module, - tp_mesh: DeviceMesh, -): - """Apply tensor parallelism.""" - - # Layer names must match Motif model definition - # https://huggingface.co/Motif-Technologies/Motif-2.6B/blob/main/modeling_motif.py - - assert type(model).__name__ == "MotifForCausalLM" - - # 1. Parallelize the embedding and shard its outputs (which are the first - # transformer block's inputs) - # 2. Parallelize the root norm layer over the sequence dim - # 3. Parallelize the final linear output layer - - parallelize_module( - model, - tp_mesh, - { - # This below separate tie_weights and make difficult to compare - # the answer with non-tensor-parallel version. - # TODO(jeesoo): check correctness for training semantic - - #"model.embed_tokens": - #RowwiseParallel( - # input_layouts=Replicate(), - # output_layouts=Shard(1), - #), - "model.norm": - SequenceParallel(), - "output": - ColwiseParallel( - input_layouts=Shard(1), - output_layouts=Shard(-1), # loss_parallel - use_local_output=False, - ), - }, - ) - - # Apply tensor + sequence parallelism to every transformer block - for transformer_block in model.model.layers: - layer_plan = { - "input_layernorm": - SequenceParallel(), - "post_attention_layernorm": - SequenceParallel(), - "self_attn": - PrepareModuleInput( - # x, freqs_cis, attention_mask, position_ids, qk_clip - input_layouts=(Shard(1), Replicate(), None, None, None), - desired_input_layouts=(Replicate(), Replicate(), None, None, - None), - ), - "self_attn.q_proj": - ColwiseParallel(), - "self_attn.k_proj": - ColwiseParallel(), - "self_attn.v_proj": - ColwiseParallel(), - "self_attn.o_proj": - RowwiseParallel(output_layouts=Shard(1)), - "mlp": - PrepareModuleInput( - input_layouts=(Shard(1), ), - desired_input_layouts=(Replicate(), ), - ), - "mlp.gate_proj": - ColwiseParallel(), - "mlp.down_proj": - RowwiseParallel(output_layouts=Shard(1)), - "mlp.up_proj": - ColwiseParallel(), - } - - parallelize_module( - module=transformer_block, - device_mesh=tp_mesh, - parallelize_plan=layer_plan, - ) - - -def _apply_fsdp( - model: torch.nn.Module, - dp_mesh: DeviceMesh, -): - for layer in model.model.layers: - fully_shard(layer, mesh=dp_mesh) - layer.reshard() - fully_shard(model, mesh=dp_mesh) - model.reshard() - - -def parallelize_llama4(model: torch.nn.Module, - parallel_dims: ParallelDims) -> torch.nn.Module: - """Parallelize the torchtitan Llama4 MoE model using torchtitan's - ``parallelize_llama`` directly. - """ - from torchtitan.config import JobConfig - from torchtitan.distributed import ParallelDims as TTParallelDims - from torchtitan.models.llama4.infra.parallelize import parallelize_llama - - world_size = dist.get_world_size() - - # Map our simple ParallelDims to torchtitan's ParallelDims. - # In torchtitan, EP borrows from dp_shard. - tt_dp_shard = parallel_dims.dp_shard_degree * parallel_dims.ep_degree - - tt_dims = TTParallelDims( - dp_replicate=parallel_dims.dp_replicate_degree, - dp_shard=tt_dp_shard, - cp=1, - tp=parallel_dims.tp_degree, - pp=1, - ep=parallel_dims.ep_degree, - etp=1, - world_size=world_size, - ) - - # Minimal JobConfig with test-appropriate settings. - job_config = JobConfig() - job_config.training.mixed_precision_param = "float32" - job_config.activation_checkpoint.mode = "none" - job_config.compile.enable = False - job_config.parallelism.disable_loss_parallel = True - - parallelize_llama(model, tt_dims, job_config) - return model - - -def parallelize_motif(model: torch.nn.Module, - parallel_dims: ParallelDims) -> torch.nn.Module: - """Parallelize the Motif model according to the given parallel dimensions. - - Args: - model (torch.nn.Module): The Motif model to be parallelized. - parallel_dims (ParallelDims): The parallelism configuration. - - Returns: - torch.nn.Module: The parallelized Motif model. - """ - - mesh = _construct_device_mesh(parallel_dims) - - if parallel_dims.tp_degree > 1: - _apply_tp(model, mesh["tp"]) - - if parallel_dims.dp_shard_degree > 1: - if parallel_dims.dp_replicate_degree > 1: - dp_dim_names = ("dp_replicate", "dp_shard") - else: - dp_dim_names = ("dp_shard", ) - _apply_fsdp(model, mesh[dp_dim_names]) - - return model - - -def parallelize_qk_logits( - qk_logits: dict[int, torch.Tensor], - parallel_dims: ParallelDims, -) -> dict[int, torch.Tensor]: - """Parallelize the QK logits according to the given parallel dimensions. - - Args: - qk_logits (dict[int, torch.Tensor]): The QK logits to be parallelized. - parallel_dims (ParallelDims): The parallelism configuration. - - Returns: - dict[int, torch.Tensor]: The parallelized QK logits. - """ - - mesh = _construct_device_mesh(parallel_dims) - - if parallel_dims.tp_degree > 1: - tp_rank = mesh["tp"].get_local_rank() - placements = [ - Shard(0) if dim_name == "tp" else Replicate() - for dim_name in mesh.mesh_dim_names - ] - for layer_idx, logits in qk_logits.items(): - assert logits.size(0) % parallel_dims.tp_degree == 0 - local_logits = logits.chunk(parallel_dims.tp_degree, - dim=0)[tp_rank].contiguous() - - qk_logits[layer_idx] = DTensor.from_local( - local_tensor=local_logits, - device_mesh=mesh, - placements=placements, - ) - - return qk_logits - - -def assert_params_equal(actual: torch.nn.Module, - expected: torch.nn.Module, - atol: float = 0, - rtol: float = 0) -> None: - """Asserts that the parameters of two models are equal. - - Args: - actual (torch.nn.Module): The actual model. - expected (torch.nn.Module): The expected model. - atol: Absolute tolerance. - rtol: Relative tolerance. - Returns: - None - """ - - def get_full_param(param: torch.nn.Parameter) -> torch.Tensor: - if isinstance(param.data, DTensor): - return param.data.full_tensor() - return param.data - - for (name_p, p), (name_s, s) in zip(actual.named_parameters(), - expected.named_parameters()): - p = get_full_param(p.cuda()) - s = get_full_param(s.cuda()) - - torch.testing.assert_close(p, s, atol=atol, rtol=rtol) diff --git a/torch-ext/optimizer/adamw.py b/torch-ext/optimizer/adamw.py deleted file mode 100644 index b5a95816a9f5b9e1889eaadae65373bfbced809a..0000000000000000000000000000000000000000 --- a/torch-ext/optimizer/adamw.py +++ /dev/null @@ -1,271 +0,0 @@ -import logging -from collections import defaultdict -from typing import cast - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -def fused_adamw( - params: list[torch.Tensor], - grads: list[torch.Tensor], - exp_avgs: list[torch.Tensor], - exp_avg_sqs: list[torch.Tensor], - max_exp_avg_sqs: list[torch.Tensor], - state_steps: list[torch.Tensor], - amsgrad: bool, - beta1: float, - beta2: float, - lr: float | torch.Tensor, - weight_decay: float, - eps: float, - maximize: bool, -) -> None: - if not params: - return - - # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer - # treating it as a scalar. - lr_dict: dict | None = ({ - lr.device: lr - } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else None) - grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype( - [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, - state_steps] # type: ignore[list-item] - ) - for (device, _), ( - ( - device_params_, - device_grads_, - device_exp_avgs_, - device_exp_avg_sqs_, - device_max_exp_avg_sqs, - device_state_steps_, - ), - _, - ) in grouped_tensors.items(): - device_params = cast(list[torch.Tensor], device_params_) - device_grads = cast(list[torch.Tensor], device_grads_) - device_exp_avgs = cast(list[torch.Tensor], device_exp_avgs_) - device_exp_avg_sqs = cast(list[torch.Tensor], device_exp_avg_sqs_) - device_state_steps = cast(list[torch.Tensor], device_state_steps_) - - if lr_dict is not None and device not in lr_dict: - lr_dict[device] = lr.to( - device=device, non_blocking=True) # type: ignore[union-attr] - lr = lr_dict[device] - torch._foreach_add_(device_state_steps, 1) - func = torch._fused_adamw_ - func( - device_params, - device_grads, - device_exp_avgs, - device_exp_avg_sqs, - device_max_exp_avg_sqs, # type: ignore[arg-type] - device_state_steps, - amsgrad=amsgrad, - lr=lr, # type: ignore[arg-type] - beta1=beta1, - beta2=beta2, - weight_decay=weight_decay, - eps=eps, - maximize=maximize, - ) - - -def _to_local(t): - """Unwrap DTensor to local tensor for fused ops.""" - return t._local_tensor if isinstance(t, DTensor) else t - - -# --------------------------------------------------------------------------- -# Caches for eliminating per-step Python overhead. -# -# Placement grouping and tensor list assembly are identical every step -# (params don't change placement, moment/step tensors are the same objects -# after initialisation). We cache them keyed by id() of the param list -# stored in param_groups (stable across steps). -# -# Only gradients change each step and must be collected fresh. -# --------------------------------------------------------------------------- - -# id(group["params"]) → dict[placement_key, list[param]] -_placement_cache: dict[int, dict[tuple, list]] = {} - -# id(placement_group_list) → (params_local, moment1, moment2, state_steps) -_tensor_cache: dict[int, tuple[list, list, list, list]] = {} - - -def _step_adamw_params_slow(optimizer_state, params, group): - """Uncached fallback for the rare case where some params lack grads.""" - params_with_grads = [] - grads = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - g = p.grad - if g is None: - continue - state = optimizer_state[p] - params_with_grads.append(_to_local(p)) - grads.append(_to_local(g)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(g) - state["moment2"] = torch.zeros_like(g) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - if not params_with_grads: - return - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - fused_adamw( - params_with_grads, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw_params(optimizer_state, params, group): - """Run fused AdamW on a list of parameters sharing the same placement. - - After the first call, cached tensor lists (params_local, moment1, - moment2, state_steps) are reused — only gradients are collected fresh. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - params: List of parameters to update. - group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay. - """ - # Collect grads — the only thing that changes each step. - with record_function("adamw::collect_grads"): - grads = [] - for p in params: - g = p.grad - if g is None: - # Rare: fall back to slow path that filters per-param. - _step_adamw_params_slow(optimizer_state, params, group) - return - grads.append(_to_local(g)) - - tensor_key = id(params) - if tensor_key not in _tensor_cache: - with record_function("adamw::init_tensor_cache"): - params_local = [] - moment1 = [] - moment2 = [] - state_steps = [] - - for p in params: - state = optimizer_state[p] - params_local.append(_to_local(p)) - if "step" not in state: - state["step"] = torch.zeros((), - dtype=torch.float32, - device=p.device) - state["moment1"] = torch.zeros_like(p.grad) - state["moment2"] = torch.zeros_like(p.grad) - moment1.append(_to_local(state["moment1"])) - moment2.append(_to_local(state["moment2"])) - if not isinstance(state["step"], torch.Tensor): - state["step"] = torch.tensor(state["step"], - dtype=torch.float32, - device=p.device) - state_steps.append(state["step"]) - - _tensor_cache[tensor_key] = (params_local, moment1, moment2, - state_steps) - - params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key] - - lr = group["lr"] - beta1, beta2 = group["adamw_betas"] - eps = group["adamw_eps"] - weight_decay = group["weight_decay"] - - with record_function("adamw::fused_adamw"): - fused_adamw( - params_local, - grads, - moment1, - moment2, - [], - state_steps, - amsgrad=False, - beta1=beta1, - beta2=beta2, - lr=lr, - weight_decay=weight_decay, - eps=eps, - maximize=False, - ) - - -def step_adamw(optimizer_state, group): - """Dispatch AdamW step, grouping parameters by type and placement. - - Placement grouping is cached after the first call since params never - change their placement between steps. - - Args: - optimizer_state: The optimizer's state dict (self.state in Muon). - group: Parameter group dict. - """ - params = group["params"] - placement_key = id(params) - - if placement_key not in _placement_cache: - with record_function("adamw::group_by_placement"): - placement_to_params: dict[tuple, - list[torch.Tensor]] = defaultdict(list) - for p in params: - match p: - case DTensor(): - logger.debug( - "[AdamW] DTensor param: shape=%s, placements=%s, " - "mesh=%s, grad=%s", p.shape, p.placements, - p.device_mesh.mesh_dim_names, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple( - [p.placements, p.device_mesh])].append(p) - case torch.Tensor(): - logger.debug( - "[AdamW] plain param: shape=%s, grad=%s", p.shape, - p.grad.shape if p.grad is not None else None) - placement_to_params[tuple([torch.Tensor, - None])].append(p) - - logger.debug("[AdamW] %d placement groups, %d total params", - len(placement_to_params), len(params)) - - _placement_cache[placement_key] = dict(placement_to_params) - - for group_params in _placement_cache[placement_key].values(): - step_adamw_params(optimizer_state, group_params, group) diff --git a/torch-ext/optimizer/async_utils.py b/torch-ext/optimizer/async_utils.py deleted file mode 100644 index a45c530ac9cad88e3555ec1047a6aa59f225347e..0000000000000000000000000000000000000000 --- a/torch-ext/optimizer/async_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import logging -from typing import Generator - -logger = logging.getLogger(__name__) - - -class _Task: - """Internal: wraps a generator, advances one yield at a time.""" - - def __init__(self, generator: Generator[None, None, None], index: int): - self._generator = generator - self._index = index - self._steps_completed = 0 - self.step() # run to first yield - - def step(self) -> bool: - try: - next(self._generator) - self._steps_completed += 1 - logger.debug("pipeline[%d] completed stage %d", self._index, - self._steps_completed) - return True - except StopIteration: - logger.debug("pipeline[%d] finished after %d stages", self._index, - self._steps_completed) - return False - - def close(self): - self._generator.close() - - -def run_pipeline( - pipelines: Generator[Generator[None, None, None], None, None], - max_concurrent: int, -) -> None: - """Run generator-based pipelines with bounded concurrency. - - Each pipeline is a generator that yields at stage boundaries. - The runtime interleaves pipelines so communication and computation - overlap across chunks. - """ - if max_concurrent <= 0: - raise ValueError(f"max_concurrent must be > 0, got {max_concurrent}") - - have_new = True - task_index = 0 - previous_tasks: list[_Task] = [] - - try: - while have_new or previous_tasks: - running_tasks: list[_Task] = [] - - # Admit one new pipeline per iteration (staggered admission). - # Admitting one at a time ensures that while chunk N does NS - # compute on the default stream, chunk N+1's NCCL all-to-all - # runs concurrently on the NCCL stream — creating real - # communication/computation overlap on the GPU. - if have_new and len(previous_tasks) < max_concurrent: - try: - gen = next(pipelines) - task = _Task(gen, task_index) - task_index += 1 - running_tasks.append(task) - except StopIteration: - have_new = False - - # Advance every previously-yielded task by one step. - for task in previous_tasks: - if task.step(): - running_tasks.append(task) - - previous_tasks = running_tasks - except BaseException: - # Clean up all in-flight generators to release GPU resources. - for task in previous_tasks: - task.close() - raise diff --git a/torch-ext/optimizer/core.py b/torch-ext/optimizer/core.py deleted file mode 100644 index c69d515afef305ad0ed66374095fa2d2468d99cc..0000000000000000000000000000000000000000 --- a/torch-ext/optimizer/core.py +++ /dev/null @@ -1,219 +0,0 @@ -import logging -import math -from dataclasses import dataclass -from typing import List - -import torch -from torch.distributed import ProcessGroup -from torch.distributed.tensor import DTensor - -# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into -# parameter FQNs. Activation checkpointing similarly inserts -# "_checkpoint_wrapped_module". Strip these so name-based matching (skip_keys, -# expert_keys, QK layer parsing) works regardless of wrapper nesting. -_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) - -logger = logging.getLogger(__name__) - - -def normalize_fqn(name: str) -> str: - """Strip torch.compile / checkpoint wrapper components from a parameter FQN.""" - return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS) - - -@dataclass -class _muon_state: - worker_rank: int - process_group: ProcessGroup - rank_indices: dict[int, tuple] # local_rank -> per-dim indices - rank_numels: dict[int, int] # local_rank -> numel - name: str - qk_clip_state: torch.Tensor | None = None - - -def _batch_momentum( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update (no nesterov).""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - - -def _batch_momentum_nesterov( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, -) -> None: - """Batched momentum update with nesterov correction.""" - torch._foreach_mul_(momentum_bufs, momentum) - torch._foreach_add_(momentum_bufs, grads) - nesterov_terms = torch._foreach_mul(momentum_bufs, momentum) - torch._foreach_add_(grads, nesterov_terms) - - -_compiled_momentum: dict[bool, callable] = {} -_use_momentum_compile = True - - -def set_momentum_compile(enabled: bool): - """Toggle torch.compile for batched momentum.""" - global _use_momentum_compile - _use_momentum_compile = enabled - - -def batch_pre_ortho( - grads: List[torch.Tensor], - momentum_bufs: List[torch.Tensor], - momentum: torch.Tensor, - nesterov: bool, -) -> None: - """Batched momentum update on lists of plain tensors. - - Mirrors dion's ``muon_update_pre_orthogonalize``. - Inputs must be plain CUDA tensors (not DTensor). - Modifies ``momentum_bufs`` and (for nesterov) ``grads`` in-place. - - When compile is enabled, uses separately compiled functions for - nesterov=True/False to avoid graph breaks from the branch. - """ - fn = _batch_momentum_nesterov if nesterov else _batch_momentum - if _use_momentum_compile: - if nesterov not in _compiled_momentum: - _compiled_momentum[nesterov] = torch.compile(fn) - fn = _compiled_momentum[nesterov] - fn(grads, momentum_bufs, momentum) - - -def _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay): - """Weight-decay + update on plain tensors. - - Not compiled: per-param @torch.compile caused ~0.25ms TorchDynamo cache - lookup per call × 256+ params = massive overhead. The pipeline path uses - batched _foreach_* ops instead; this function remains for base() and - distributed_muon(). - """ - p_data.mul_(1 - lr * weight_decay) - p_data.add_(u_data, alpha=-adjusted_lr) - - -def update_p(p, u, lr, adjusted_lr, weight_decay): - """Apply weight decay and orthogonalized update to parameter. - - Args: - p: Parameter (torch.nn.Parameter or DTensor). - u: Orthogonalized update tensor. - lr: Base learning rate. - adjusted_lr: Size-adjusted learning rate. - weight_decay: Weight decay coefficient. - """ - # Unwrap Parameter -> underlying data tensor. - p_data = p.data if isinstance(p, torch.nn.Parameter) else p - # Unwrap DTensor -> local CUDA tensor for compiled kernel. - if isinstance(p_data, DTensor): - p_data = p_data._local_tensor - u_data = u._local_tensor if isinstance(u, DTensor) else u - _update_p_impl(p_data, u_data, lr, adjusted_lr, weight_decay) - - -def adjust_lr_for_muon(lr, param_shape): - """Scale learning rate based on parameter matrix dimensions. - - Args: - lr: Base learning rate. - param_shape: Shape of the parameter tensor. - - Returns: - Adjusted learning rate. - """ - A, B = param_shape[:2] - # We adjust the learning rate and weight decay based on the size of the parameter matrix - # as described in the paper - adjusted_ratio = 0.2 * math.sqrt(max(A, B)) - adjusted_lr = lr * adjusted_ratio - return adjusted_lr - - -def _match_key(parts, key): - """Check if key matches as contiguous components in parts. - - Single-component keys (e.g. "experts") match any single component. - Multi-component keys (e.g. "experts.w1") match as a contiguous subsequence. - """ - key_parts = key.split(".") - key_len = len(key_parts) - if key_len == 1: - return key in parts - return any(parts[i:i + key_len] == key_parts - for i in range(len(parts) - key_len + 1)) - - -def is_expert_param(name, expert_keys): - """Check if a parameter name matches any expert key (component-level).""" - if not expert_keys: - return False - parts = normalize_fqn(name).split(".") - return any(_match_key(parts, key) for key in expert_keys) - - -def default_is_muon(name, x, expert_keys=None): - normalized = normalize_fqn(name) - parts = normalized.split(".") - skip_keys = [ - "embed_tokens", - "lm_head", - "tok_embeddings", - "output", - "mhc_attn", - "mhc_ffn", - "lambda_proj", - ] - if any(key in parts for key in skip_keys): - logger.info( - "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d", - normalized, name, x.ndim) - return False - effective_ndim = x.ndim - is_expert = is_expert_param(name, expert_keys) - if is_expert: - effective_ndim -= 1 - result = effective_ndim >= 2 - logger.info( - "[is_muon] %s (orig: %s): ndim=%d, expert=%s, effective_ndim=%d → %s", - normalized, name, x.ndim, is_expert, effective_ndim, - "Muon" if result else "AdamW") - return result - - -def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None): - if is_muon_func is None: - is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys) - - muon_params, muon_names = [], [] - non_muon_params, non_muon_names = [], [] - - for n, p in model.named_parameters(): - if not p.requires_grad: - continue - if is_muon_func(n, p): - muon_params.append(p) - muon_names.append(n) - else: - non_muon_params.append(p) - non_muon_names.append(n) - - logger.info("[param_groups] expert_keys=%s, Muon=%d, AdamW=%d", - expert_keys, len(muon_names), len(non_muon_names)) - - return [ - { - "params": muon_params, - "names": muon_names, - "use_muon": True, - }, - { - "params": non_muon_params, - "use_muon": False, - }, - ] diff --git a/torch-ext/optimizer/cpu_offload.py b/torch-ext/optimizer/cpu_offload.py deleted file mode 100644 index 29cacd1e06d20fbe44b2adb751f2de21467d84dd..0000000000000000000000000000000000000000 --- a/torch-ext/optimizer/cpu_offload.py +++ /dev/null @@ -1,364 +0,0 @@ -"""CPU offloading for optimizer states. - -Manages a pinned CPU memory pool and async CUDA streams to offload -optimizer state tensors (momentum buffers, Adam moments) to CPU between -optimizer steps, freeing GPU memory. - -All tracked tensors are packed into a single flat pinned CPU buffer -(per dtype). D2H and H2D copies are performed per-tensor directly -between individual GPU tensors and their slice of the CPU flat buffer -— no GPU staging buffer is allocated, so there is **no temporary GPU -memory spike** during offload or reload. - -Individual tensor storages are freed after offload via -``untyped_storage().resize_(0)``, preserving tensor identity so -downstream caches remain valid. -""" - -import logging -from collections import defaultdict - -import torch -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -logger = logging.getLogger(__name__) - - -class CPUOffloadPool: - """Pinned CPU memory pool for async optimizer state offloading. - - Tracked tensors are grouped by dtype. Each group gets a single flat - pinned CPU buffer. D2H / H2D copies are per-tensor (into slices of - the flat buffer) to avoid allocating a GPU staging buffer. - """ - - def __init__(self): - self._managed: list[torch.Tensor] = [] - self._storage_nbytes: dict[int, int] = {} # id(t) → bytes - # Optional tag → managed-indices map for group-wise reload - # (e.g. per-layer lockstep reload driven by backward hooks). - self._tag_to_indices: dict[str, list[int]] = {} - - # Per-dtype group: populated on first offload. - # dtype → dict with keys: - # "indices" : list[int] managed-list indices - # "offsets" : list[tuple[int,int]] (start, numel) in flat buf - # "total" : int total numel - # "cpu_flat" : Tensor pinned CPU buffer - self._groups: dict[torch.dtype, dict] = {} - - self._offload_stream: torch.cuda.Stream | None = None - self._reload_stream: torch.cuda.Stream | None = None - self._reload_event: torch.cuda.Event | None = None - self._device: torch.device | None = None - self._initialized: bool = False - self._logged: bool = False - - # ------------------------------------------------------------------ - @staticmethod - def _local(t: torch.Tensor) -> torch.Tensor: - """Unwrap DTensor to its local CUDA tensor.""" - return t._local_tensor if isinstance(t, DTensor) else t - - def _ensure_stream(self): - if self._offload_stream is None: - self._offload_stream = torch.cuda.Stream(device=self._device) - - def _ensure_reload_stream(self): - if self._reload_stream is None: - least_priority, _ = torch.cuda.Stream.priority_range() - self._reload_stream = torch.cuda.Stream( - device=self._device, - priority=least_priority, - ) - logger.info( - "[CPUOffload] reload stream created with priority=%d " - "(range: %d..%d)", - least_priority, - *torch.cuda.Stream.priority_range(), - ) - - # ------------------------------------------------------------------ - def track(self, tensor: torch.Tensor, tag: str | None = None): - """Register a GPU tensor for CPU offloading. Idempotent. - - If ``tag`` is given, the tensor's managed index is recorded under - that tag so callers can trigger a partial reload via - :meth:`reload_group`. - """ - tid = id(tensor) - if tid in self._storage_nbytes: - return - local = self._local(tensor) - if self._device is None: - self._device = local.device - storage = local.untyped_storage() - # Skip tensors with empty storage (e.g. empty FSDP shards) - if storage.size() == 0: - return - self._storage_nbytes[tid] = storage.size() - idx = len(self._managed) - self._managed.append(tensor) - if tag is not None: - self._tag_to_indices.setdefault(tag, []).append(idx) - - # ------------------------------------------------------------------ - def _init_buffers(self): - """Build per-dtype flat buffers on first offload.""" - # Group managed tensors by dtype. - dtype_map: dict[torch.dtype, list[tuple[int, int]]] = defaultdict(list) - for idx, t in enumerate(self._managed): - local = self._local(t) - dtype_map[local.dtype].append((idx, local.numel())) - - total_cpu_bytes = 0 - for dtype, entries in dtype_map.items(): - offsets: list[tuple[int, int]] = [] - indices: list[int] = [] - off = 0 - for idx, n in entries: - indices.append(idx) - offsets.append((off, n)) - off += n - cpu_flat = torch.empty(off, - dtype=dtype, - device="cpu", - pin_memory=True) - self._groups[dtype] = { - "indices": indices, - "offsets": offsets, - "total": off, - "cpu_flat": cpu_flat, - } - total_cpu_bytes += off * cpu_flat.element_size() - - self._initialized = True - logger.info( - "[CPUOffload] Pool initialized: %d tensors, %d dtype group(s), " - "%.2f MB pinned CPU memory", - len(self._managed), - len(self._groups), - total_cpu_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def offload(self): - """Per-tensor async D2H into CPU flat buffer, then free GPU storage.""" - if not self._managed: - return - if not self._initialized: - self._init_buffers() - self._ensure_stream() - - # Offload stream waits for compute to finish. - compute_event = torch.cuda.current_stream(self._device).record_event() - self._offload_stream.wait_event(compute_event) - - offloaded_bytes = 0 - - # Per-tensor D2H copies directly into CPU flat buffer slices. - # No GPU staging buffer → no temporary GPU memory spike. - with torch.cuda.stream(self._offload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - cpu_flat[off:off + n].copy_(local.reshape(-1), - non_blocking=True) - - offloaded_bytes += grp["total"] * cpu_flat.element_size() - - # Wait for all D2H copies to land, then free GPU storage. - self._offload_stream.synchronize() - for t in self._managed: - storage = self._local(t).untyped_storage() - if storage.size() != 0: - storage.resize_(0) - else: - raise RuntimeError( - f"Tensor storage is already freed (size=0) before offload. " - f"This indicates a double-free or external interference. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}") - - if not self._logged: - logger.info( - "[CPUOffload] Offloaded %.2f MB (GPU → CPU)", - offloaded_bytes / (1024**2), - ) - - # ------------------------------------------------------------------ - def reload(self): - """Per-tensor H2D from CPU flat buffer. - - Storage re-allocation (``resize_``) runs on the current (default) - stream. H2D copies run on a dedicated ``_reload_stream``. - - Call :meth:`wait_reload` before consuming the reloaded tensors. - """ - if not self._managed or not self._initialized: - return - self._ensure_reload_stream() - - reloaded_bytes = 0 - - # Re-allocate all GPU storages with per-tensor profiling. - with record_function("CPUOffload::resize_storages"): - for i, t in enumerate(self._managed): - local = self._local(t) - storage = local.untyped_storage() - if storage.size() != 0: - raise RuntimeError( - f"Storage should have been freed (size=0) before reload, " - f"but got size={storage.size()}. " - f"Tensor shape: {t.shape}, dtype: {t.dtype}") - nbytes = self._storage_nbytes[id(t)] - with record_function(f"resize_[{i}]_{nbytes // 1024}KB"): - storage.resize_(nbytes) - - # Reload stream waits for the resize_ ops to finish. - alloc_event = torch.cuda.current_stream(self._device).record_event() - self._reload_stream.wait_event(alloc_event) - - # Per-tensor H2D copies on the reload stream. - with record_function("CPUOffload::h2d_copies"): - with torch.cuda.stream(self._reload_stream): - for dtype, grp in self._groups.items(): - indices = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices): - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off:off + n], - non_blocking=True) - - reloaded_bytes += grp["total"] * cpu_flat.element_size() - - self._reload_event = self._reload_stream.record_event() - - if not self._logged: - logger.info( - "[CPUOffload] Reloaded %.2f MB (CPU → GPU, async)", - reloaded_bytes / (1024**2), - ) - self._logged = True - - def reload_group(self, tag: str, sync_streams: tuple = ()): - """Reload only the managed tensors registered under ``tag``. - - Intended for layer-lockstep overlap: backward frees a layer's - activations, then the backward hook calls ``reload_group`` with - that layer's tag so the H2D copy reuses the freshly-freed memory - from the default stream's allocator pool. - - ``sync_streams`` is an optional iterable of CUDA streams whose - currently-queued work must complete before the H2D memcpy runs. - This is used to avoid allocator cross-stream reuse races under - ``expandable_segments``: if a just-freed block was last used on - FSDP's all-gather stream, making the reload stream wait on that - stream guarantees FIFO ordering between the block's prior use - and our H2D write. - """ - if not self._managed or not self._initialized: - return - indices = self._tag_to_indices.get(tag) - if not indices: - return - self._ensure_reload_stream() - - # Sync reload_stream with the supplied streams (e.g. FSDP AG - # streams) before we queue any H2D: ensures past uses of any - # allocator block we're about to reuse are fully drained. - for s in sync_streams: - if s is not None: - self._reload_stream.wait_stream(s) - - idx_set = set(indices) - - with record_function(f"CPUOffload::group_resize[{tag}]"): - for i in indices: - t = self._managed[i] - local = self._local(t) - storage = local.untyped_storage() - if storage.size() == 0: - storage.resize_(self._storage_nbytes[id(t)]) - - alloc_event = torch.cuda.current_stream(self._device).record_event() - self._reload_stream.wait_event(alloc_event) - - with record_function(f"CPUOffload::group_h2d[{tag}]"): - with torch.cuda.stream(self._reload_stream): - for dtype, grp in self._groups.items(): - indices_grp = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices_grp): - if mgd_idx not in idx_set: - continue - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off:off + n], - non_blocking=True) - - self._reload_event = self._reload_stream.record_event() - - def reload_untagged(self): - """Reload managed tensors that were not registered under any tag. - - Useful when a subset of params (e.g. MoE experts) is driven via - per-tag layer-lockstep hooks while the remainder should still be - reloaded before optimizer.step() in a single bulk call. - """ - if not self._managed or not self._initialized: - return - tagged: set[int] = set() - for idx_list in self._tag_to_indices.values(): - tagged.update(idx_list) - untagged = [i for i in range(len(self._managed)) if i not in tagged] - if not untagged: - return - self._ensure_reload_stream() - - idx_set = set(untagged) - - with record_function("CPUOffload::untagged_resize"): - for i in untagged: - t = self._managed[i] - local = self._local(t) - storage = local.untyped_storage() - if storage.size() == 0: - storage.resize_(self._storage_nbytes[id(t)]) - - alloc_event = torch.cuda.current_stream(self._device).record_event() - self._reload_stream.wait_event(alloc_event) - - with record_function("CPUOffload::untagged_h2d"): - with torch.cuda.stream(self._reload_stream): - for dtype, grp in self._groups.items(): - indices_grp = grp["indices"] - offsets = grp["offsets"] - cpu_flat = grp["cpu_flat"] - - for i, mgd_idx in enumerate(indices_grp): - if mgd_idx not in idx_set: - continue - local = self._local(self._managed[mgd_idx]) - off, n = offsets[i] - local.reshape(-1).copy_(cpu_flat[off:off + n], - non_blocking=True) - - self._reload_event = self._reload_stream.record_event() - - def wait_reload(self): - """Block the current (default) stream until reload H2D completes.""" - if self._reload_event is not None: - torch.cuda.current_stream(self._device).wait_event( - self._reload_event) - self._reload_event = None diff --git a/torch-ext/optimizer/distributed/utils.py b/torch-ext/optimizer/distributed/utils.py deleted file mode 100644 index 683096bcd50489e7a09959753fdbc29ba0f3b6a6..0000000000000000000000000000000000000000 --- a/torch-ext/optimizer/distributed/utils.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import (Placement, Shard, - _StridedShard) - - -def _is_shard(placement: Placement) -> bool: - """Check if a placement is a shard type (Shard or _StridedShard). - - In PyTorch 2.10+, _StridedShard no longer inherits from Shard, so - ``placement.is_shard()`` returns False for _StridedShard. This helper - handles both old and new hierarchies. - """ - return isinstance(placement, (Shard, _StridedShard)) - - -def get_slices_of_dtensor( - target: DTensor | torch.Tensor, - local_rank: int, - shard_mesh: DeviceMesh, - shard_placements: tuple[Placement], -) -> tuple[slice | torch.Tensor, ...]: - """ - Get per-dimension indices for a given rank's shard of the target tensor. - - Uses ``Shard.local_shard_size_and_offset`` and - ``_StridedShard.local_shard_size_and_offset`` for correct handling of - both contiguous and strided (non-contiguous) sharding. - - Args: - target (DTensor | torch.Tensor): The target tensor (for its shape). - local_rank (int): The local rank within the shard group. - shard_mesh (DeviceMesh): The shard mesh (only shard dimensions). - shard_placements (tuple[Placement]): The shard placements. - - Returns: - A tuple of indices (one per tensor dim). Each element is either: - - A ``slice`` (for contiguous or unsharded dims) - - A 1-D ``torch.LongTensor`` of indices (for strided sharding) - """ - - # find the global rank of the local rank in the shard mesh - rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank] - - rank_coords = (shard_mesh.mesh == rank).nonzero() - - assert len(rank_coords) == 1 - rank_coords = tuple(rank_coords[0].tolist()) - - assert len(rank_coords) == len(shard_placements) - - # Track per-shard-dim indices. - # None means "not yet sharded on this dim". - dim_indices: dict[int, torch.Tensor] = {} - - # Caution: Assuming replicate-to-shard of the shard mesh goes with - # left-to-right sharding. This is ensured by the sorting logic of - # construct_shard_mesh function. - for mesh_dim_idx, (rank_coord, placement) in enumerate( - zip(rank_coords, shard_placements)): - assert _is_shard(placement) - - num_chunks = shard_mesh.mesh.shape[mesh_dim_idx] - shard_dim = placement.dim - - # Current effective size on this dim (may already be sub-sharded) - if shard_dim in dim_indices: - curr_size = len(dim_indices[shard_dim]) - else: - curr_size = target.size()[shard_dim] - - # Compute indices for this level of sharding - if isinstance(placement, _StridedShard): - _shard_size, offsets = _StridedShard.local_shard_size_and_offset( - placement, - curr_size, - num_chunks, - rank_coord, - return_first_offset=False) - new_indices = torch.tensor(offsets, dtype=torch.long) - else: - shard_size, offset = Shard.local_shard_size_and_offset( - curr_size, num_chunks, rank_coord) - new_indices = torch.arange(offset, - offset + shard_size, - dtype=torch.long) - - # Compose with previous indices on this dim - if shard_dim in dim_indices: - dim_indices[shard_dim] = dim_indices[shard_dim][new_indices] - else: - dim_indices[shard_dim] = new_indices - - # Build result tuple - result: list[slice | torch.Tensor] = [] - for d in range(len(target.size())): - if d not in dim_indices: - result.append(slice(None)) - else: - indices = dim_indices[d] - # Convert contiguous indices to slice for efficiency - if len(indices) > 0: - start = indices[0].item() - expected = torch.arange(start, - start + len(indices), - dtype=torch.long) - if torch.equal(indices, expected): - result.append(slice(start, start + len(indices))) - else: - result.append(indices) - else: - result.append(slice(0, 0)) - - return tuple(result) - - -_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, - ProcessGroup]] = dict() - - -def construct_shard_mesh( - placements: tuple[Placement], - mesh: DeviceMesh, -) -> tuple[DeviceMesh, ProcessGroup, tuple[Placement, ...]]: - """Construct shard sub-mesh and ProcessGroup for all-to-all communication. - - Given a DTensor's placements and device mesh, extracts the "shard group" - — the set of ranks that together hold all shards of the same replica — - and creates a ProcessGroup for all-to-all among them. - - Steps: - 1. Sort placements: Replicate first, then Shard by (dim, granularity). - 2. Permute the mesh tensor to match the sorted order. - 3. Collapse Replicate dims → list of shard sub-meshes (one per replica). - 4. Create/retrieve a cached ProcessGroup for the current rank's sub-mesh. - - Example — 8 GPUs, mesh shape (2, 2, 2), - placements ``[Shard(0), Replicate, _StridedShard(0)]``:: - - Step 1 — Sort: [Replicate, _StridedShard(0), Shard(0)] - Permutation: [1, 2, 0] - - Step 2 — Permute mesh dims by [1, 2, 0]: - Original: Permuted: - [[[0,1],[2,3]], [[[0,2],[1,3]], - [[4,5],[6,7]]] [[4,6],[5,7]]] - - Step 3 — Unbind replicate dim (dim 0), giving 2 shard sub-meshes: - sub-mesh 0 = [[0,2],[1,3]] (replica group 0) - sub-mesh 1 = [[4,6],[5,7]] (replica group 1) - shard_placements = (_StridedShard(0), Shard(0)) - - Step 4 — Rank 0 → ProcessGroup([0,1,4,5]) - Rank 2 → ProcessGroup([2,3,6,7]) - - Returns: - ``(shard_mesh, process_group, shard_placements)`` - """ - my_rank = dist.get_rank() - assert mesh.mesh.device.type == 'cpu' - - # -- Fast path: 1D all-shard mesh → reuse existing PG. ---------------- - # Reuses the mesh's existing ProcessGroup directly, avoiding the - # overhead of dist.new_group(). The standard path below also handles - # subset calls safely via use_local_synchronization=True, but this - # fast path is still beneficial for the common 1D shard case. - if mesh.ndim == 1 and len(placements) == 1 and _is_shard(placements[0]): - key = (*mesh.mesh.shape, *mesh.mesh.flatten().tolist()) - if key not in _ranks_to_dist_cache: - _ranks_to_dist_cache[key] = (mesh, mesh.get_group()) - return (*_ranks_to_dist_cache[key], tuple(placements)) - - mesh_tensor = mesh.mesh.clone() - - # -- Step 1: Sort placements (Replicate first, then Shard by dim). ------ - # _StridedShard comes BEFORE regular Shard on the same dim so that - # get_slices_of_dtensor applies the outer sharding first, matching - # DTensor's left-to-right (outer-to-inner) composition order. - def _sort_key(item): - index, placement = item - assert not placement.is_partial(), "Partial placement not supported" - if placement.is_replicate(): - return (-1, 0, index) - assert _is_shard(placement), f"Unsupported: {type(placement)}" - split = (-1 / placement.split_factor if isinstance( - placement, _StridedShard) else 0) - return (placement.dim, split, index) - - indexed = sorted(enumerate(placements), key=_sort_key) - perm, sorted_placements = zip(*indexed) - - # -- Step 2: Permute mesh to match sorted placement order. -------------- - sorted_mesh = mesh_tensor.permute(perm) - - # -- Step 3: Collapse replicate dims → list of shard sub-meshes. -------- - # E.g. mesh (2, 3, 4, 4) with [R, R, S(0), S(1)] → 6 sub-meshes of (4, 4) - num_rep = sum(1 for p in sorted_placements if p.is_replicate()) - if num_rep > 0: - if num_rep > 1: - sorted_mesh = sorted_mesh.flatten(0, num_rep - 1) - shard_meshes = list(torch.unbind(sorted_mesh, dim=0)) - else: - shard_meshes = [sorted_mesh] - shard_placements = sorted_placements[num_rep:] - assert len(shard_placements) == len(set(shard_placements)) - - # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. -- - # Each rank only creates the group it belongs to, using - # use_local_synchronization=True so that only group members need to - # coordinate. This avoids deadlocks when different PP stages call - # construct_shard_mesh for different parameters. - def _cache_key(t: torch.Tensor) -> tuple: - return (*t.shape, *t.flatten().tolist()) - - my_key = None - for sm in shard_meshes: - if (my_rank == sm).any().item(): - key = _cache_key(sm) - assert my_key is None, "Rank appears in multiple shard groups" - my_key = key - if key not in _ranks_to_dist_cache: - pg = dist.new_group(sm.flatten().tolist(), - use_local_synchronization=True) - _ranks_to_dist_cache[key] = ( - DeviceMesh(device_type="cuda", mesh=sm), - pg, - ) - - return (*_ranks_to_dist_cache[my_key], shard_placements) diff --git a/torch-ext/optimizer/matmul_transpose_triton.py b/torch-ext/optimizer/matmul_transpose_triton.py deleted file mode 100644 index 792de23d82c3fb45fe33d397ab9b76a0787259d0..0000000000000000000000000000000000000000 --- a/torch-ext/optimizer/matmul_transpose_triton.py +++ /dev/null @@ -1,122 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Tianyang Lin -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -import torch -import triton -import triton.language as tl - - -def get_autotune_config(): - return [ - triton.Config( - { - 'BLOCK_SIZE_M': blk_m, - 'BLOCK_SIZE_K': blk_k, - 'GROUP_SIZE_M': grp_sz - }, - num_stages=n_stages, - num_warps=n_warps) for blk_m in [32, 64, 128] - for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5] - for n_warps in [4, 8] - ] - - -@triton.autotune( - configs=get_autotune_config(), - key=['M', 'K'], - restore_value=['y'], -) -@triton.jit -def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr): - """ - Core kernel jit function of matmul_transpose that computes y = x @ x.T - The code is a simple adaptation from the triton `matmul` tutorial: - https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - """ - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - if pid_m > pid_n: - return - - offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_k = tl.arange(0, BLOCK_SIZE_K) - # we use a & b ptrs to denote different rows of x. - a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk) - b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - b = tl.load(b_ptrs, - mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, - other=0.0) - accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator) - a_ptrs += BLOCK_SIZE_K * stride_xk - b_ptrs += BLOCK_SIZE_K * stride_xk - # use dtype.element_ty to accommodate different input datatypes as in cpp templates - # https://github.com/triton-lang/triton/issues/2252 - c = accumulator.to(x.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M) - tl.store(c_ptrs, c, mask=c_mask) - - # transpose and copy - if pid_m < pid_n: - ct_ptrs = y + stride_ym * offs_cn[:, - None] + stride_yn * offs_cm[None, :] - ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M) - tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask) - - -@torch.library.custom_op("muon::matmul_transpose_assign", - mutates_args=("d_out", )) -def matmul_transpose_assign(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """Compute d_out = d_in @ d_in.T using an optimized Triton kernel.""" - d_in = d_in.contiguous() - M, K = d_in.shape - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv( - M, META['BLOCK_SIZE_M']), ) - with torch.cuda.device(d_in.device.index): - mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1), - d_out.stride(0), d_out.stride(1)) - - -@matmul_transpose_assign.register_fake -def _(d_in: torch.Tensor, d_out: torch.Tensor) -> None: - """FakeTensor impl: d_out is already allocated, mutation is declared.""" - pass diff --git a/torch-ext/optimizer/muon.py b/torch-ext/optimizer/muon.py index 443caa4aa2c89c13763c84a65d320f056a06c29f..0d614d55d721efac406c147b4f62e6c703a91107 100644 --- a/torch-ext/optimizer/muon.py +++ b/torch-ext/optimizer/muon.py @@ -1,136 +1,138 @@ -import logging -import types -from collections import defaultdict -from typing import Any +import math +from dataclasses import dataclass import torch import torch.distributed as dist -from torch.distributed.tensor import DTensor, Replicate, Shard -from torch.profiler import record_function - -from .adamw import _placement_cache, _tensor_cache, step_adamw -from .async_utils import run_pipeline -from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho, - get_default_muon_param_groups, is_expert_param, update_p) -from .cpu_offload import CPUOffloadPool -from .distributed.utils import (_is_shard, construct_shard_mesh, - get_slices_of_dtensor) -from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO, - _zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5, - zeropower_via_newtonschulz5_batched) -from .pipeline import muon_chunk_pipeline, prelaunch_first_gather -from .qk_clip import compute_scales, get_qk_clip_info, qk_clip - -logger = logging.getLogger(__name__) - - -def _expand_expert_params(names, params, expert_keys): - """Expand expert params by splitting on dim 0 (expert dimension). - - Params whose name matches any key in ``expert_keys`` are treated as - expert-parallel tensors. Their outermost dimension is the expert - dimension: an ``(E, out, in)`` tensor becomes ``E`` separate 2D - ``nn.Parameter`` views so that in-place updates propagate back to - the original storage. - - Non-expert params with ``ndim > 2`` trigger an ``AssertionError`` — - if they are expert params, their key must be added to ``expert_keys``. - - The grad must already be set on each expert param (e.g. after momentum). - - For DTensor expert params, placements that shard on dim 0 (expert dim) - are consumed by the split. Non-dim-0 shard placements (e.g. TP) are - preserved: each 2D slice is wrapped as a DTensor on the corresponding - submesh so the parallel pipeline handles the TP communication. +from torch.distributed._tensor import DTensor + + +# This code snippet is a modified version adapted from the following GitHub repositories: +# https://github.com/KellerJordan/Muon/blob/master/muon.py +@torch.no_grad() +def _zeropower_via_newtonschulz5(G, steps): """ - expanded_names = [] - expanded_params = [] - - for n, p in zip(names, params): - is_expert = is_expert_param(n, expert_keys) - is_dtensor = isinstance(p.data, DTensor) - - if is_expert: - if is_dtensor: - logger.debug( - "[expand_expert] %s: expert DTensor, shape=%s, " - "placements=%s, mesh=%s, local_shape=%s", n, p.shape, - p.placements, p.device_mesh.mesh_dim_names, - p.to_local().shape) - else: - logger.debug( - "[expand_expert] %s: expert plain tensor, shape=%s", n, - p.data.shape) - - if not is_expert: - assert p.data.ndim <= 2, ( - f"Param {n} has ndim={p.data.ndim} but does not match " - f"expert_keys={expert_keys}. If this is an expert param, " - f"add its key to expert_keys.") - expanded_names.append(n) - expanded_params.append(p) - continue - - g = p.grad - assert g is not None, ( - f"Expert param {n} must have grad set before expansion") - - tp_mesh = None - tp_placements_2d = None - - if is_dtensor: - local_data = p.to_local() - local_grad = g.to_local() if isinstance(g, DTensor) else g - - # Find non-dim-0 shard placements (e.g. TP sharding). - # After splitting on dim 0, Shard(k) becomes Shard(k-1). - tp_dim_indices = [] - tp_placements_2d = [] - for i, pl in enumerate(p.placements): - if _is_shard(pl) and pl.dim != 0: - tp_dim_indices.append(i) - tp_placements_2d.append(Shard(pl.dim - 1)) - - if tp_dim_indices: - tp_dim_names = tuple(p.device_mesh.mesh_dim_names[i] - for i in tp_dim_indices) - if len(tp_dim_names) == 1: - tp_mesh = p.device_mesh[tp_dim_names[0]] - else: - tp_mesh = p.device_mesh[tp_dim_names] + Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a + quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose + of minimizing steps, it turns out to be empirically effective to keep increasing the slope at + zero even beyond the point where the iteration no longer converges all the way to one everywhere + on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T + where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model + performance at all relative to UV^T, where USV^T = G is the SVD. + """ + assert len(G.shape) == 2 + a, b, c = (3.4445, -4.7750, 2.0315) + X = G # no manual typecast + if G.size(0) > G.size(1): + X = X.T + # Ensure spectral norm is at most 1 + X = X / (X.norm() + 1e-7) + X = X.bfloat16() + # Perform the NS iterations + for _ in range(steps): + A = X @ X.T + # B = ( + # b * A + c * A @ A + # ) + B = torch.addmm(A, A, A, alpha=c, beta=b) + # X = a * X + B @ X + X = torch.addmm(X, B, X, alpha=1.0, beta=a) + + if G.size(0) > G.size(1): + X = X.T + return X.to(G.dtype) + + +@dataclass +class _muon_state: + # TODO: use Optional + worker_rank: int | None = None + gathered_grad: torch.Tensor | None = None + computed_u: torch.Tensor | None = None + gather_event: torch.cuda.Event | None = None + compute_event: torch.cuda.Event | None = None + + +@torch.no_grad() +def _gather(p, state, rank, comm_stream, none_grad): + g = p.grad + mesh = g.device_mesh + + if rank == state.worker_rank: + gather_list = [torch.empty_like(g.to_local()) for _ in range(mesh.mesh.numel())] + else: + gather_list = None + + with torch.cuda.stream(comm_stream): + torch.distributed.gather( + g.to_local(), + dst=state.worker_rank, + gather_list=gather_list, + group=mesh.get_group(), + ) + if rank == state.worker_rank: + if state.gathered_grad is not None: + raise RuntimeError( + "Gather event already exists, which should not happen." + ) + state.gathered_grad = torch.cat(gather_list, dim=0) + state.gather_event = torch.cuda.Event() + state.gather_event.record() else: - local_data = p.data - local_grad = g - - # Expand: split dim 0, reshape each slice to 2D. - num_local_experts = local_data.shape[0] - for i in range(num_local_experts): - slice_data = local_data[i] - slice_grad = local_grad[i] - - if tp_mesh is not None: - # Wrap as DTensor on TP submesh so the pipeline handles - # TP communication (gather/scatter across TP ranks). - dt_data = DTensor.from_local(slice_data, - device_mesh=tp_mesh, - placements=tp_placements_2d) - dt_grad = DTensor.from_local(slice_grad, - device_mesh=tp_mesh, - placements=tp_placements_2d) - expert_param = torch.nn.Parameter(dt_data, requires_grad=False) - expert_param.grad = dt_grad - else: - expert_param = torch.nn.Parameter(slice_data, - requires_grad=False) - expert_param.grad = slice_grad + state.gathered_grad = None + state.gather_event = None + if none_grad: + p.grad = None + + +@torch.no_grad() +def _compute_u(state, steps, rank, compute_stream): + with torch.cuda.stream(compute_stream): + if rank == state.worker_rank: + if state.gather_event is None: + raise RuntimeError("Gather event must be set before compute.") + compute_stream.wait_event(state.gather_event) + u = _zeropower_via_newtonschulz5(state.gathered_grad, steps) + state.computed_u = u + state.compute_event = torch.cuda.Event() + state.compute_event.record() + # Clear the gathered gradient to free memory + state.gathered_grad = None + else: + state.computed_u = None + state.compute_event = None - expanded_names.append(f"{n}[{i}]") - expanded_params.append(expert_param) - p.grad = None # allow expert grad storage to be freed after pipeline +@torch.no_grad() +def _scatter(p, state, lr, wd, rank, comm_stream): + u = state.computed_u + mesh = p.device_mesh - return expanded_names, expanded_params + with torch.cuda.stream(comm_stream): + if rank == state.worker_rank: + if state.compute_event is None: + raise RuntimeError("Compute event must be set before scatter.") + comm_stream.wait_event(state.compute_event) + scatter_list = list(torch.split(u, p.size(0) // mesh.mesh.numel(), dim=0)) + else: + scatter_list = None + + u = torch.empty_like(p.to_local()) + torch.distributed.scatter( + u, + scatter_list=scatter_list, + src=state.worker_rank, + group=mesh.get_group(), + ) + if rank == state.worker_rank: + # Clear u to free memory + state.computed_u = None + u = DTensor.from_local( + u, + placements=p.placements, + device_mesh=mesh, + ) + p.data.mul_(1 - lr * wd) + p.data.add_(u, alpha=-lr) class Muon(torch.optim.Optimizer): @@ -147,109 +149,71 @@ class Muon(torch.optim.Optimizer): - We believe it may not work well for finetuning pretrained models, but we haven't tested this. Arguments: - model: The model to be optimized by Muon. - is_muon_func: A function that takes a parameter and its name, and returns whether the parameter should be optimized by Muon. + muon_params: The parameters to be optimized by Muon. lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) momentum: The momentum used by the internal SGD. (0.95 is a good default) nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) - weight_decay: The weight decay for Muon and AdamW. - Parameters that are {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW instead. + adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are + {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. adamw_lr: The learning rate for the internal AdamW. adamw_betas: The betas for the internal AdamW. adamw_eps: The epsilon for the internal AdamW. - none_grad: Whether to set p.grad to None after gathering the gradients. This can save memory. - debug: Whether to print debug information. - clip_info : Configuration for QK clipping. Expected keys: - - "q_indices" (list[int]): Indices of query heads to consider. - - "k_indices" (list[int]): Indices of key heads to consider. - - "head_dim" (int): Dimensionality of each attention head. - - "threshold" (float): Threshold value; heads whose QK logits exceed - this value will be scaled down. - Default is: - { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100 - } - warmup_step : How many all2all gather, compute operations are launched in advance - before the corresponding all2all scatter steps begin. - A higher warmup_step increases memory usage but can improve - performance by overlapping communication. - Parallel muon only. - chunk_size : Batch size of parameters to process in each - all2all gather/compute/scatter step. - Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified. - use_distributed_muon: Use distributed muon by Liu et al. (2024). - For testing purpose only. - expert_keys: List of strings to identify expert-parallel parameters. - If any key appears in a parameter's name, its outermost - dimension is treated as the expert dimension and expanded - into per-expert 2D params for Muon. For example, - ``expert_keys=["experts"]`` matches any param whose name - contains "experts". 3D+ params not matched by any key - will raise an error. + adamw_wd: The weight decay for the internal AdamW. """ - def __init__(self, - params, - lr=1e-3, - momentum=0.95, - nesterov=True, - ns_steps=5, - weight_decay=0.1, - adamw_betas=(0.9, 0.95), - adamw_eps=1e-8, - none_grad=True, - debug=False, - clip_config=None, - warmup_step=5, - chunk_size=-1, - use_distributed_muon=False, - expert_keys=None): + def __init__( + self, + model, + is_muon_func, + lr=1e-3, + momentum=0.95, + nesterov=True, + ns_steps=5, + adamw_wd=0.1, + adamw_betas=(0.9, 0.95), + adamw_eps=1e-8, + none_grad=True, + debug=False, + ): defaults = dict( lr=lr, - weight_decay=weight_decay, + wd=adamw_wd, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps, adamw_betas=adamw_betas, adamw_eps=adamw_eps, none_grad=none_grad, - use_muon=True, ) - error_message = "The key 'use_muon' is not set in parameter group {idx}. Assuming all parameters in the group will use muon optimization, which may lead to unexpected behavior." - instruction_code = "\n\n please follow this code snippet \n```optimizer = get_kernel('motif-technologies/optimizer')\n\n\nparams = optimizer.muon.get_default_muon_param_groups(model)\n\noptim = optimizer.Muon(params, ...)```" - if isinstance(params, types.GeneratorType): - raise ValueError(error_message.format(idx=0) + instruction_code) - for _idx, param_group in enumerate(params): - if param_group.get("use_muon", None) is None: - raise ValueError( - error_message.format(idx=_idx) + instruction_code) - super().__init__(params, defaults) + super().__init__(model.parameters(), defaults) + self.is_muon_func = is_muon_func + self.model = model + + if not dist.is_initialized(): + raise RuntimeError( + "Muon optimizer requires distributed training to be initialized." + ) + self.rank = dist.get_rank() + + self.comm_stream = torch.cuda.Stream() + self.compute_stream = torch.cuda.Stream() self.debug = debug - self.clip_config = clip_config if clip_config is not None else { - "q_indices": [], - "k_indices": [], - "head_dim": 128, - "threshold": 100, - } - self.warmup_step = warmup_step - self.chunk_size = chunk_size - self.use_distributed_muon = use_distributed_muon - self.expert_keys = expert_keys - self.cpu_offload = False - self.manual_offload = False - self._cpu_offload_pool: CPUOffloadPool | None = None - self._offload_initialized = False - # id(param) -> tag, consumed by _register_states_for_offload so the - # offload pool can do group-wise reload (e.g. per-layer lockstep). - self._param_tags: dict[int, str] = {} - self._parallel_cache: dict[tuple[str, ...], dict] = {} - self._expert_expand_cache: dict[tuple[int, ...], dict] = {} + + def __setstate__(self, state): + # Sort parameters into those for which we will use Muon, and those for which we will not + super().__setstate__(state) + for name, p in self.model.named_parameters(): + if self.is_muon_func(p, name): + # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer + assert p.ndim == 2, p.ndim + self.state[p]["use_muon"] = True + self.state[p]["orig_shape"] = p.shape + else: + # Do not use Muon for parameters in adamw_params + self.state[p]["use_muon"] = False def _calc_flops(self, G, steps): assert len(G.shape) == 2 @@ -259,19 +223,15 @@ class Muon(torch.optim.Optimizer): return steps * ((M**3) * 2 + (M**2 * N) * 4 + M * N * 2 + M**2 * 3) - def get_shard_mesh(self, p): - """ - Get the shard mesh for a parameter p on the given rank. - """ - assert isinstance( - p, DTensor), "Parallel Muon only supports DTensor parameters." - - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - p.placements, p.device_mesh) + def adjust_lr_for_muon(self, lr, param_shape): + A, B = param_shape[:2] + # We adjust the learning rate and weight decay based on the size of the parameter matrix + # as describted in the paper + adjusted_ratio = 0.2 * math.sqrt(max(A, B)) + adjusted_lr = lr * adjusted_ratio + return adjusted_lr - return shard_mesh, shard_pg, shard_placements - - def init_state_and_assign_params(self, names, params, group, qk_logits): + def init_state_and_assign_params(self, params, group): param_to_state = {} param_to_flops = {} @@ -287,831 +247,209 @@ class Muon(torch.optim.Optimizer): total_flops += flops if self.debug: - logger.debug("Total TFLOPs for Muon: %.2f TFLOPs", - total_flops / 1e12) - - paired = list(zip(names, params)) - - paired_sorted = sorted(paired, - key=lambda x: param_to_flops[id(x[1])], - reverse=True) + print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", flush=True) - names_sorted, params_sorted = zip(*paired_sorted) - ordered_names = list(names_sorted) - ordered_params = list(params_sorted) + ordered_params = sorted( + params, key=lambda p: param_to_flops[id(p)], reverse=True + ) round_robin = 0 - mesh = ordered_params[0].device_mesh - placements = ordered_params[0].placements + mesh = None + for p in ordered_params: + if mesh is None: + mesh = p.device_mesh + if mesh.ndim != 1: + raise NotImplementedError( + "Muon requires a 1D mesh for distributed training yet." + ) + elif mesh != p.device_mesh: + raise ValueError("All parameters must be on the same mesh.") - shard_mesh, shard_pg, shard_placements = self.get_shard_mesh( - ordered_params[0]) - shard_mesh_flattened = shard_mesh.mesh.flatten() - num_ranks = dist.get_world_size(group=shard_pg) + param_to_state[id(p)] = _muon_state() + param_to_state[id(p)].worker_rank = mesh.mesh[round_robin].item() - for n, p in zip(ordered_names, ordered_params): - if mesh != p.device_mesh: - raise ValueError("All parameters must be on the same mesh.") - if placements != p.placements: - raise ValueError("All parameters must have same placements.") - - worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks - round_robin = (round_robin + 1) % len(shard_mesh_flattened) - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - # Precompute per-rank indices and numels for all-to-all. - rank_indices: dict[int, tuple] = {} - rank_numels: dict[int, int] = {} - for r in range(num_ranks): - indices = get_slices_of_dtensor(p, r, shard_mesh, - shard_placements) - rank_indices[r] = indices - numel = 1 - for idx, dim_size in zip(indices, p.shape): - if isinstance(idx, slice): - start, stop, step = idx.indices(dim_size) - numel *= max(0, (stop - start + (step - 1)) // step) - else: - numel *= len(idx) - rank_numels[r] = numel - - param_to_state[id(p)] = _muon_state( - worker_rank=worker_rank, - process_group=shard_pg, - rank_indices=rank_indices, - rank_numels=rank_numels, - name=n, - qk_clip_state=qk_clip_state, - ) + round_robin = (round_robin + 1) % mesh.mesh.numel() return param_to_state, ordered_params - def base(self, names, params, group, lr, weight_decay, qk_logits): - # Momentum is already applied by _step_muon before this method. - for n, p in zip(names, params): + def base(self, params, group, lr, wd, momentum): + # generate weight updates in distributed fashion + for p in params: g = p.grad if g is None: continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + assert g is not None + + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf - u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE), - steps=group["ns_steps"]) - - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits) - - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - def distributed_muon( - self, - names: list[str], - params: list[torch.nn.Parameter], - group: dict[str, Any], - lr: float, - weight_decay: float, - qk_logits: list[torch.Tensor | DTensor] | None, - ): - """Batched Distributed Muon — for testing/correctness verification only. - - Uses all-gather to reconstruct full tensors, computes Newton-Schulz on - the full grad, then slices back to local shards. This is simpler but - slower than the parallel pipeline (all2all) path, so it serves as a - reference implementation for verifying correctness. - """ - with record_function("distributed_muon"): - # Momentum is already applied by _step_muon before this method. - ns_steps = group["ns_steps"] - - # Separate plain tensors (no communication) from DTensors. - plain_names, plain_params = [], [] - dtensor_names, dtensor_params = [], [] - for n, p in zip(names, params): - if p.grad is None: - continue - if isinstance(p.data, DTensor): - dtensor_names.append(n) - dtensor_params.append(p) - else: - plain_names.append(n) - plain_params.append(p) - - # Process plain tensors per-param (no communication). - for n, p in zip(plain_names, plain_params): - u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE), - steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - update_p(p, u, lr, adjusted_lr, weight_decay) - - qk_clip_state = get_qk_clip_info(self.clip_config, n, - qk_logits) - scales_full = compute_scales( - p, qk_clip_state) if qk_clip_state is not None else None - if scales_full is not None: - qk_clip(p, scales_full, qk_clip_state) - - if not dtensor_params: - return - - # Group DTensors by (placements, mesh) for batched all-gather. - placement_groups: dict[tuple, - tuple[list, - list]] = defaultdict(lambda: ([], [])) - for n, p in zip(dtensor_names, dtensor_params): - key = (p.placements, p.device_mesh) - placement_groups[key][0].append(n) - placement_groups[key][1].append(p) - - logger.info( - "distributed_muon: %d placement groups, %d total dtensors", - len(placement_groups), len(dtensor_params)) - - for (placements, mesh), (grp_names, - grp_params) in placement_groups.items(): - shard_mesh, shard_pg, shard_placements = construct_shard_mesh( - placements, mesh) - rank = dist.get_rank(shard_pg) - world_size = dist.get_world_size(shard_pg) - - logger.info(" group: %d params, placements=%s, world_size=%d", - len(grp_params), placements, world_size) - - # Separate params that can be batched (all shard dims evenly - # divisible) from those needing per-param full_tensor - # (e.g. MoE gate weights with fewer rows than shard ranks). - # all_gather_into_tensor requires equal buffer sizes across - # ranks, so uneven splits must use DTensor full_tensor(). - batch_names, batch_params = [], [] - single_names, single_params = [], [] - for n, p in zip(grp_names, grp_params): - even = all(p.shape[pl.dim] % - shard_mesh.mesh.shape[dim_idx] == 0 - for dim_idx, pl in enumerate(shard_placements)) - if even: - batch_names.append(n) - batch_params.append(p) - else: - single_names.append(n) - single_params.append(p) - - # Process uneven-split params per-param via full_tensor(). - for n, p in zip(single_names, single_params): - with record_function("distributed_muon::newton_schulz"): - g_full = p.grad.full_tensor().to(COMM_DTYPE) - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - if not batch_params: - continue + u = _zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) - logger.info(" batched=%d, single=%d", len(batch_params), - len(single_params)) - - # Concat all local grad shards into a single flat buffer. - with record_function("distributed_muon::gather"): - grad_locals = [ - p.grad.to_local().to(COMM_DTYPE).flatten() - for p in batch_params - ] - numels = [g.numel() for g in grad_locals] - grad_concat = torch.cat(grad_locals) - del grad_locals - - # Single all-gather (replaces N separate full_tensor). - grad_gathered = torch.empty( - grad_concat.numel() * world_size, - dtype=COMM_DTYPE, - device="cuda", - ) - dist.all_gather_into_tensor(grad_gathered, - grad_concat, - group=shard_pg) - - total_numel = grad_concat.numel() - del grad_concat - - # Precompute per-param offsets within the concat buffer. - offsets = [] - off = 0 - for ne in numels: - offsets.append(off) - off += ne - - # Per-param: reconstruct full grad → NS → local update. - for i, (n, p) in enumerate(zip(batch_names, batch_params)): - with record_function("distributed_muon::newton_schulz"): - g_full = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - for r in range(world_size): - r_start = r * total_numel + offsets[i] - shard = grad_gathered[r_start:r_start + numels[i]] - indices = get_slices_of_dtensor( - p, r, shard_mesh, shard_placements) - g_full[indices] = shard.reshape( - g_full[indices].shape) - - u_full = _zeropower_via_newtonschulz5(g_full, - steps=ns_steps) - del g_full - - with record_function("distributed_muon::update"): - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - p._local_tensor.mul_(1 - lr * weight_decay) - local_indices = get_slices_of_dtensor( - p, rank, shard_mesh, shard_placements) - u_local = u_full[local_indices] - p._local_tensor.add_(u_local, alpha=-adjusted_lr) - del u_full - - qk_clip_state = get_qk_clip_info( - self.clip_config, n, qk_logits) - scales_full = compute_scales( - p, qk_clip_state - ) if qk_clip_state is not None else None - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = local_indices[0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - def _setup_parallel(self, names, params, group, qk_logits): - """Compute (or retrieve cached) parallel pipeline metadata. - - Returns: - (ordered_params, param_to_state, rank, chunk_size) - """ - cache_key = tuple(names) + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) - if cache_key not in self._parallel_cache: - # First call: compute metadata and populate cache. - param_to_state, ordered_params = self.init_state_and_assign_params( - names, params, group, qk_logits) + # apply weight decay + p.data.mul_(1 - lr * wd) - shard_pg = param_to_state[id(ordered_params[0])].process_group - rank = dist.get_rank(group=shard_pg) + # apply update + p.data.add_(u, alpha=-adjusted_lr) - if self.chunk_size == -1: - shard_ranks = dist.get_world_size(shard_pg) - chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO - elif self.chunk_size > 0: - chunk_size = self.chunk_size - else: - raise ValueError( - "chunk_size must be -1 or a positive integer.") - - ordered_names = [ - param_to_state[id(p)].name for p in ordered_params - ] - name_to_state = { - param_to_state[id(p)].name: param_to_state[id(p)] - for p in ordered_params - } - self._parallel_cache[cache_key] = { - 'ordered_names': ordered_names, - 'name_to_state': name_to_state, - 'rank': rank, - 'chunk_size': chunk_size, - } + def _update_g(self, p, g, group, momentum): + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) else: - # Cached path: rebuild param_to_state with current id(p) keys. - cache = self._parallel_cache[cache_key] - rank = cache['rank'] - chunk_size = cache['chunk_size'] - - name_to_param = dict(zip(names, params)) - ordered_params = [name_to_param[n] for n in cache['ordered_names']] - - param_to_state = {} - for p, n in zip(ordered_params, cache['ordered_names']): - cached_state = cache['name_to_state'][n] - param_to_state[id(p)] = _muon_state( - worker_rank=cached_state.worker_rank, - process_group=cached_state.process_group, - rank_indices=cached_state.rank_indices, - rank_numels=cached_state.rank_numels, - name=n, - qk_clip_state=get_qk_clip_info(self.clip_config, n, - qk_logits), - ) - - return ordered_params, param_to_state, rank, chunk_size - - def parallel(self, - names, - params, - group, - lr, - weight_decay, - qk_logits, - prelaunch_gather=None): + g = buf + return g + + def _update_p(self, p, u, lr, wd): + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + # apply weight decay + p.data.mul_(1 - lr * wd) + # apply update + p.data.add_(u, alpha=-adjusted_lr) + + def parallel(self, params, group, lr, wd, momentum): """ Perform a parallel optimization step using Muon. - - Parameters are chunked and each chunk is processed by a - :func:`muon_chunk_pipeline` generator. :func:`run_pipeline` - interleaves multiple chunks so that communication and computation - overlap across chunks (the same overlap previously achieved by the - warmup + main-loop index scheduling). - - If ``prelaunch_gather`` is provided, it is passed to the first - chunk's generator to skip re-launching the already in-flight - A2A gather. """ - # Momentum is already applied by _step_muon before this method. - - ordered_params, param_to_state, rank, chunk_size = ( - self._setup_parallel(names, params, group, qk_logits)) - - def pipelines(): - first = True - for start in range(0, len(ordered_params), chunk_size): - chunk = ordered_params[start:start + chunk_size] - if chunk: - kwargs = dict( - params=chunk, - param_to_state=param_to_state, - rank=rank, - ns_steps=group["ns_steps"], - lr=lr, - weight_decay=weight_decay, - none_grad=group["none_grad"], - ) - if first and prelaunch_gather is not None: - kwargs['prelaunch_gather'] = prelaunch_gather - first = False - yield muon_chunk_pipeline(**kwargs) - - with record_function("muon::pipeline"): - run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1) - - def _step_muon(self, group, qk_logits=None): - params = group["params"] - lr = group["lr"] - weight_decay = group["weight_decay"] - momentum = group["momentum"] - names = group["names"] - - # Apply momentum to all params before routing/expansion. - # Batched using _foreach_* ops (compiled, fullgraph=True). - with record_function("muon::momentum"): - active_params = [p for p in params if p.grad is not None] - if active_params: - # Ensure momentum buffers exist (avoid zeros_like when already present). - for p in active_params: - if "momentum_buffer" not in self.state[p]: - self.state[p]["momentum_buffer"] = torch.zeros_like( - p.grad) - - # Extract local tensors for compiled batch function. - local_grads = [ - p.grad._local_tensor - if isinstance(p.grad, DTensor) else p.grad - for p in active_params - ] - local_bufs = [ - self.state[p]["momentum_buffer"]._local_tensor - if isinstance(self.state[p]["momentum_buffer"], DTensor) - else self.state[p]["momentum_buffer"] - for p in active_params - ] - - # Wrap momentum as tensor for torch.compile. - batch_pre_ortho(local_grads, local_bufs, - torch.tensor(momentum), group["nesterov"]) - - # For non-nesterov, the result is the momentum buffer. - if not group["nesterov"]: - for p in active_params: - p.grad = self.state[p]["momentum_buffer"] - - # Identify batched experts for deferred NS. - # Detection is cheap (condition checks only); actual NS compute is - # deferred so it can overlap with the first chunk's A2A gather. - deferred_expert_work = [] - if self.expert_keys: - batched_expert_indices = [] - for i, (n, p) in enumerate(zip(names, params)): - if not (is_expert_param(n, self.expert_keys) - and p.grad is not None): - continue - # Eligible: plain tensor, or DTensor with no non-dim-0 shards. - if isinstance(p.data, DTensor): - has_tp = any( - _is_shard(pl) and pl.dim != 0 for pl in p.placements) - if has_tp: - continue - batched_expert_indices.append(i) - - if batched_expert_indices: - # Save refs for deferred NS; free grads from param list. - for i in batched_expert_indices: - p = params[i] - g = p.grad - local_g = (g._local_tensor - if isinstance(g, DTensor) else g) - local_data = (p.data._local_tensor if isinstance( - p.data, DTensor) else p.data) - deferred_expert_work.append((local_data, local_g)) - p.grad = None - - # Remove batched experts from lists before expansion. - keep = sorted( - set(range(len(params))) - set(batched_expert_indices)) - names = [names[i] for i in keep] - params = [params[i] for i in keep] - - def _run_deferred_expert_ns(): - """Execute deferred batched expert NS.""" - if not deferred_expert_work: - return - with record_function("muon::batched_expert_ns"): - ns_steps = group["ns_steps"] - for local_data, local_g in deferred_expert_work: - u = zeropower_via_newtonschulz5_batched( - local_g.to(COMM_DTYPE), steps=ns_steps) - adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:]) - local_data.mul_(1 - lr * weight_decay) - local_data.add_(u, alpha=-adjusted_lr) - - # Expand expert params by splitting on dim 0. - logger.debug("[_step_muon] before expand: %d params, expert_keys=%s", - len(params), self.expert_keys) - if self.expert_keys: - cache_key = tuple(id(p) for p in params) - cache = self._expert_expand_cache.get(cache_key) - - if cache is None: - # Cold path: full expansion + build cache metadata. - exp_names, exp_params = _expand_expert_params( - names, params, self.expert_keys) - - # Build per-expert-group info for hot-path grad updates. - grad_info = [] - exp_idx = 0 - for orig_idx, (n, p) in enumerate(zip(names, params)): - if not is_expert_param(n, self.expert_keys): - exp_idx += 1 - continue - - is_dt = isinstance(p.data, DTensor) - num_experts = (p.to_local() if is_dt else p.data).shape[0] - - # Detect TP mesh from the first expanded expert param. - tp_mesh = None - tp_pls = None - sample = exp_params[exp_idx] - if isinstance(sample.data, DTensor): - tp_mesh = sample.data.device_mesh - tp_pls = list(sample.data.placements) - - grad_info.append((orig_idx, num_experts, exp_idx, is_dt, - tp_mesh, tp_pls)) - exp_idx += num_experts - - self._expert_expand_cache[cache_key] = { - 'names': exp_names, - 'params': exp_params, - 'grad_info': grad_info, - } - names, params = exp_names, exp_params - else: - # Hot path: reuse cached params, only update expert grads. - for (orig_idx, num_experts, exp_start, is_dt, tp_mesh, - tp_pls) in cache['grad_info']: - p = params[orig_idx] - g = p.grad - local_grad = (g.to_local() - if is_dt and isinstance(g, DTensor) else g) - for i in range(num_experts): - expert_p = cache['params'][exp_start + i] - sg = local_grad[i] - if tp_mesh is not None: - expert_p.grad = DTensor.from_local( - sg, device_mesh=tp_mesh, placements=tp_pls) - else: - expert_p.grad = sg - p.grad = None - - names = cache['names'] - params = cache['params'] - else: - names, params = _expand_expert_params(names, params, - self.expert_keys) - logger.debug("[_step_muon] after expand: %d params", len(params)) - - param_dtensors = [] - name_dtensors = [] - - param_tensors = [] - name_tensors = [] - - # distributed_muon is a reference implementation for testing only. - # The parallel pipeline (all2all) path below is the production path. - if self.use_distributed_muon: - _run_deferred_expert_ns() - self.distributed_muon(names=names, - params=params, - group=group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits) - return - - for n, p in zip(names, params): - if p is None or p.grad is None: + for p in params: + g = p.grad + if g is None: continue - if isinstance(p.data, DTensor): - if all( - isinstance(placement, Replicate) - for placement in p.placements): - logger.debug( - "[route] %s → base (DTensor all-Replicate), " - "shape=%s, placements=%s", n, p.shape, p.placements) - param_tensors.append(p) - name_tensors.append(n) - else: - logger.debug( - "[route] %s → parallel (DTensor), shape=%s, " - "placements=%s, mesh=%s", n, p.shape, p.placements, - p.device_mesh.mesh_dim_names) - param_dtensors.append(p) - name_dtensors.append(n) - elif isinstance(p.data, torch.Tensor): - logger.debug("[route] %s → base (plain tensor), shape=%s", n, - p.data.shape) - param_tensors.append(p) - name_tensors.append(n) - else: - raise TypeError(f"Unsupported parameter type: {type(p.data)}") + if g.ndim > 2: + g = g.view(g.size(0), -1) - logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, " - f"{len(param_tensors)} Tensors → base") + # Update g in the local rank + g = self._update_g( + p, + g, + group, + momentum=momentum, + ) + p.grad = g - def group_dtensors(dtensors, names): - # To support different placements, we group parameters by placements - # and run parallel Muon on each group. + param_to_state, ordered_params = self.init_state_and_assign_params( + params, group + ) - placement_to_params = defaultdict(lambda: ([], [])) + def enqueue_gathers(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _gather(p, state, self.rank, self.comm_stream, group["none_grad"]) - assert len(dtensors) == len(names) - for p, n in zip(dtensors, names): - placement_to_params[tuple([p.placements, - p.device_mesh])][0].append(n) - placement_to_params[tuple([p.placements, - p.device_mesh])][1].append(p) - return placement_to_params + def enqueue_computes(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + _compute_u(state, group["ns_steps"], self.rank, self.compute_stream) - if len(param_dtensors) > 0: - if not dist.is_initialized(): - raise RuntimeError( - "Parallel Muon requires torch.distributed to be initialized." - ) + def enqueue_scatters(start_idx, chunk_size): + for p in ordered_params[start_idx : start_idx + chunk_size]: + state = param_to_state[id(p)] + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream) - dtensor_group = group_dtensors(param_dtensors, name_dtensors) - - # Pre-launch the first chunk's A2A gather so that the NCCL - # communication overlaps with the (deferred) batched expert NS - # compute on the default CUDA stream. - prelaunch = None - if deferred_expert_work: - first_names, first_params = next(iter(dtensor_group.values())) - ordered, pts, rnk, csz = self._setup_parallel( - first_names, first_params, group, qk_logits) - first_chunk = ordered[:csz] - if first_chunk: - prelaunch = prelaunch_first_gather(first_chunk, pts, rnk, - group["none_grad"]) - - _run_deferred_expert_ns() - - first_group = True - for _, (names, params) in dtensor_group.items(): - pg = prelaunch if first_group else None - first_group = False - self.parallel( - names, - params, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - prelaunch_gather=pg, - ) - else: - _run_deferred_expert_ns() + chunk_size = params[0].device_mesh.mesh.numel() - if len(param_tensors) > 0: - self.base( - name_tensors, - param_tensors, - group, - lr=lr, - weight_decay=weight_decay, - qk_logits=qk_logits, - ) + # Wait grad update + self.comm_stream.wait_stream(torch.cuda.current_stream()) - def _register_states_for_offload(self): - """Register all optimizer state tensors with the CPU offload pool. + enqueue_gathers(0, chunk_size) + for i in range(0, len(params) + chunk_size - 1, chunk_size): + enqueue_computes(i, chunk_size) + enqueue_gathers(i + chunk_size, chunk_size) + enqueue_scatters(i, chunk_size) - Called once after the first step when states have been lazily created. - Offloads all param states (momentum buffers for Muon, moment1/moment2 - for AdamW) to free GPU memory between steps. - """ - pool = self._cpu_offload_pool - tracked = 0 - for group in self.param_groups: - for p in group["params"]: - if p not in self.state: - continue - state = self.state[p] - tag = self._param_tags.get(id(p)) - if group.get("use_muon", False): - if "momentum_buffer" in state: - pool.track(state["momentum_buffer"], tag=tag) - tracked += 1 - else: - if "moment1" in state: - pool.track(state["moment1"], tag=tag) - if "moment2" in state: - pool.track(state["moment2"], tag=tag) - tracked += 1 - logger.info("[CPUOffload] Registered %d param states for offload", - tracked) - - @torch.no_grad - def step(self, closure=None, qk_logits=None): + torch.cuda.current_stream().wait_stream(self.comm_stream) + + def step(self, closure=None): """Perform a single optimization step. Args: closure (Callable, optional): A closure that reevaluates the model and returns the loss. - qk_logits (dict[int, Tensor], optional): A dictionary mapping layer indices - to 1D tensors of shape (num_heads,), representing the maximum - QK logits across all tokens, computed as - (1 / sqrt(head_dim)) * (Q @ K^T). """ loss = None if closure is not None: with torch.enable_grad(): loss = closure() - # H2D: reload optimizer states from CPU before computation. - if not self.manual_offload: - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload() - self._cpu_offload_pool.wait_reload() + for group in self.param_groups: + ############################ + # Muon # + ############################ - logger.debug("[Muon.step] expert_keys=%s, %d param groups", - self.expert_keys, len(self.param_groups)) + params = [p for p in group["params"] if self.state[p]["use_muon"]] + lr = group["lr"] + wd = group["wd"] + momentum = group["momentum"] - for i, group in enumerate(self.param_groups): - if group["use_muon"]: - logger.debug("[Muon.step] group %d: use_muon=True, %d params", - i, len(group["params"])) - self._step_muon(group, qk_logits=qk_logits) + if isinstance(params[0].data, DTensor): + self.parallel( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) else: - logger.debug( - "[Muon.step] group %d: use_muon=False (AdamW), %d params", - i, len(group["params"])) - step_adamw(self.state, group) - - # D2H: offload optimizer states to CPU after computation. - if not self.manual_offload: - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - return loss - - # ------------------------------------------------------------------ - # CPU offload public helpers - # ------------------------------------------------------------------ - - def reload_group(self, tag: str, sync_streams: tuple = ()): - """Reload optimizer states registered under ``tag``. - - Tags are set via :meth:`set_param_tags` before the first step. - ``sync_streams`` forwards to :meth:`CPUOffloadPool.reload_group` - so callers (e.g. FSDP pre/post-hook patches) can make the reload - stream wait on collective streams before its H2D runs. - """ - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload_group(tag, sync_streams=sync_streams) - - def reload_untagged(self): - """Reload all optimizer states not attached to any tag.""" - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.reload_untagged() + self.base( + params, + group, + lr=lr, + wd=wd, + momentum=momentum, + ) - def set_param_tags(self, param_tags: dict[int, str]) -> None: - """Attach an ``id(param) -> tag`` mapping for group-wise reload. + ############################ + # AdamW backup # + ############################ - Must be called before the first ``step()`` (i.e. before - :meth:`_register_states_for_offload`) so the pool receives tags - when states are first registered. - """ - self._param_tags = dict(param_tags) - - def wait_reload(self): - """Block the default stream until the async reload completes.""" - if self.cpu_offload and self._offload_initialized: - self._cpu_offload_pool.wait_reload() - - def offload(self): - """Offload optimizer states from GPU to CPU (D2H).""" - if self.cpu_offload: - if not self._offload_initialized: - if self._cpu_offload_pool is None: - self._cpu_offload_pool = CPUOffloadPool() - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_on_cpu_offload(self): - """Enable CPU offload for optimizer states.""" - if self.cpu_offload: - return - logger.info("[Muon] turn_on_cpu_offload") - self.cpu_offload = True - if not self.state: - return - self._cpu_offload_pool = CPUOffloadPool() - self._offload_initialized = False - self._register_states_for_offload() - self._offload_initialized = True - self._cpu_offload_pool.offload() - - def turn_off_cpu_offload(self): - """Disable CPU offload and keep optimizer states resident on GPU.""" - if not self.cpu_offload: - return - logger.info("[Muon] turn_off_cpu_offload") - if self._offload_initialized: - self._cpu_offload_pool.reload() - self._cpu_offload_pool.wait_reload() - torch.cuda.current_stream().synchronize() - self._cpu_offload_pool = None - self._offload_initialized = False - self.cpu_offload = False - - # ------------------------------------------------------------------ - # Checkpoint support for cpu_offload - # ------------------------------------------------------------------ - - def state_dict(self) -> dict: - if self.cpu_offload: - raise RuntimeError( - "Muon.state_dict() requires turn_off_cpu_offload() before checkpoint save." - ) - return super().state_dict() + params = [p for p in group["params"] if not self.state[p]["use_muon"]] + lr = group["lr"] + beta1, beta2 = group["adamw_betas"] + eps = group["adamw_eps"] + weight_decay = group["wd"] - def load_state_dict(self, state_dict: dict) -> None: - if self.cpu_offload: - raise RuntimeError( - "Muon.load_state_dict() requires turn_off_cpu_offload() before checkpoint load." - ) - super().load_state_dict(state_dict) + for p in params: + g = p.grad + if g is None: + continue + state = self.state[p] + if "step" not in state: + state["step"] = 0 + state["moment1"] = torch.zeros_like(g) + state["moment2"] = torch.zeros_like(g) + state["step"] += 1 + step = state["step"] + buf1 = state["moment1"] + buf2 = state["moment2"] + buf1.lerp_(g, 1 - beta1) + buf2.lerp_(g.square(), 1 - beta2) + + g = buf1 / (eps + buf2.sqrt()) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + scale = bias_correction1 / bias_correction2**0.5 + p.data.mul_(1 - lr * weight_decay) + p.data.add_(g, alpha=-lr / scale) - # Invalidate adamw.py's module-level tensor caches so that - # the next step rebuilds them with the newly loaded state tensors. - _placement_cache.clear() - _tensor_cache.clear() + return loss diff --git a/torch-ext/optimizer/newton_schulz.py b/torch-ext/optimizer/newton_schulz.py deleted file mode 100644 index c79e97a4f1a4f80f95907edbd208610dafe869f6..0000000000000000000000000000000000000000 --- a/torch-ext/optimizer/newton_schulz.py +++ /dev/null @@ -1,237 +0,0 @@ -from itertools import repeat -from math import inf, sqrt - -import numpy as np -import torch - -from .matmul_transpose_triton import matmul_transpose_assign - -COMM_DTYPE = torch.bfloat16 -DEFAULT_CHUNK_SIZE_RATIO = 4 - - -def _optimal_quintic(l, u, max_iter=1000): - """ - Use the simplified Remez algorithm to find the optimal odd quintic approximant - to the constant function x -> 1 over the interval [l, u]. - - Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum - approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the - two interior equioscillation nodes q, r until convergence. Returns the - closed-form equioscillating solution when l ≈ u. - - Raises ValueError if any intermediate value (a, b, c, E, q, r) is non-finite - (NaN or inf). Raises RuntimeError if convergence is not reached within - max_iter iterations. - """ - assert 0 <= l <= u - if 1 - 5e-6 <= l / u: - return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5) - q = (3 * l + u) / 4 - r = (l + 3 * u) / 4 - E = inf - for _ in range(max_iter): - old_E = E - LHS = np.array([ - [l, l**3, l**5, 1], - [q, q**3, q**5, -1], - [r, r**3, r**5, 1], - [u, u**3, u**5, -1], - ]) - a, b, c, E = np.linalg.solve(LHS, np.ones(4)) - if not np.all(np.isfinite([a, b, c, E])): - raise ValueError( - f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}" - ) - q, r = np.sqrt( - (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / - (10 * c)) - if not np.all(np.isfinite([q, r])): - raise ValueError( - f"_optimal_quintic: non-finite node update q={q}, r={r}") - if abs(old_E - E) <= 1e-15: - break - else: - raise RuntimeError( - f"_optimal_quintic: did not converge after {max_iter} iterations") - return float(a), float(b), float(c) - - -def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0): - """ - Compute the Polar Express coefficient series for `num_iters` quintic iterations. - - Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that - compose to map singular values from [l, 1] toward 1. At each step: - 1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion` - prevents near-zero singular values from stalling by raising the effective - lower bound; if it is active (cushion*u > l), the coefficients are - rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u]. - 2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the - last iteration, providing numerical headroom at the cost of a slightly slower - final convergence step. - 3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1). - - Returns a list of (a, b, c) tuples, one per iteration. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - u = 1 - assert 0 <= l <= u - safety_factor = 1 + safety_factor_eps - coefficients = [] - for iter in range(num_iters): - a, b, c = _optimal_quintic(max(l, cushion * u), u) - if cushion * u > l: - pl = a * l + b * l**3 + c * l**5 - pu = a * u + b * u**3 + c * u**5 - rescaler = 2 / (pl + pu) - a *= rescaler - b *= rescaler - c *= rescaler - if iter < num_iters - 1: - a /= safety_factor - b /= safety_factor**3 - c /= safety_factor**5 - coefficients.append((a, b, c)) - l = a * l + b * l**3 + c * l**5 - u = 2 - l - return coefficients - - -# Precomputed Polar Express coefficients (a, b, c) for 10 quintic Newton-Schulz -# iterations. Each tuple is the minimax-optimal (Remez/equioscillation) odd quintic -# approximant to x->1 over the current singular-value interval, computed once at -# import time and reused across all optimizer steps. -# -# Contrast with the former hardcoded NS coefficients (5 fixed tuples): -# - Former: empirically tuned to maximize slope at zero; did not converge -# singular values to 1, yielding US'V^T with S' ~ Uniform(0.5, 1.5) instead -# of the true polar factor UV^T. -# - Polar Express: analytically optimal per step, adapting to the shrinking -# singular-value interval [l, u] as iterations progress; converges all -# singular values to 1, producing the exact polar factor UV^T. -_coeffs_list = _optimal_composition(l=1e-3, - num_iters=10, - safety_factor_eps=1e-2, - cushion=0.02) - - -# This code is adapted from: -# KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py) -# NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress) -# matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon) -@torch.no_grad() -def _zeropower_via_newtonschulz5(G, steps): - """ - Compute the polar factor of G via the Polar Express method. - - Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c) - are the Polar Express coefficients from `_coeffs_list`. Each step is the - optimal odd quintic approximant to x -> 1 over the current singular-value - interval, minimizing the maximum approximation error (Remez / minimax criterion). - The composition maps singular values from [l, 1] to near 1, producing the - polar factor (orthogonal factor in the polar decomposition G = UP). - - `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2, - cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated. - - Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and - Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932 - """ - assert len(G.shape) == 2 - assert G.dtype == COMM_DTYPE - X = G # no manual typecast - - if G.size(0) > G.size(1): - X = X.T - - X = X / (X.norm() + 1e-7) - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list))) - buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device) - # Perform the NS iterations - for a, b, c in hs: - matmul_transpose_assign(X, buf1) - matmul_transpose_assign(buf1, buf2) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.addmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(0) > G.size(1): - X = X.T - - return X - - -@torch.no_grad() -def _zeropower_via_newtonschulz5_batched(G, steps): - """Batched polar factor computation for 3D (E, out, in) tensors. - - Same algorithm as ``_zeropower_via_newtonschulz5`` but uses - ``torch.bmm`` / ``torch.baddbmm`` instead of the 2D Triton kernel, - processing all E expert matrices in a single batched call. - """ - assert len(G.shape) == 3 - assert G.dtype == COMM_DTYPE - X = G - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - # Per-expert Frobenius norm. - X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) - - hs = _coeffs_list[:steps] + list( - repeat(_coeffs_list[-1], steps - len(_coeffs_list))) - for a, b, c in hs: - buf1 = torch.bmm(X, X.transpose(-2, -1)) - buf2 = torch.bmm(buf1, buf1.transpose(-2, -1)) - buf1.mul_(b).add_(buf2, alpha=c) - X = torch.baddbmm(X, buf1, X, alpha=1.0, beta=a) - - if G.size(1) > G.size(2): - X = X.transpose(-2, -1) - - return X - - -_ns_per_shape: dict[tuple[int, ...], callable] = {} -_use_compile = True - - -def set_ns_compile(enabled: bool): - """Toggle torch.compile for Newton-Schulz iteration.""" - global _use_compile - _use_compile = enabled - - -def zeropower_via_newtonschulz5(G, steps=5): - if not _use_compile: - return _zeropower_via_newtonschulz5(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() - - -def zeropower_via_newtonschulz5_batched(G, steps=5): - """Compile-cached batched Newton-Schulz for 3D expert tensors.""" - if not _use_compile: - return _zeropower_via_newtonschulz5_batched(G, steps) - key = G.shape - if key not in _ns_per_shape: - _ns_per_shape[key] = torch.compile( - _zeropower_via_newtonschulz5_batched, - options={ - "triton.cudagraphs": True, - "shape_padding": False - }) - torch.compiler.cudagraph_mark_step_begin() - return _ns_per_shape[key](G, steps).clone() diff --git a/torch-ext/optimizer/pipeline.py b/torch-ext/optimizer/pipeline.py deleted file mode 100644 index c0c2d515856182d8d15ad27dd4e4e093b29397d6..0000000000000000000000000000000000000000 --- a/torch-ext/optimizer/pipeline.py +++ /dev/null @@ -1,468 +0,0 @@ -import logging -from typing import Generator - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.profiler import record_function - -from .core import _muon_state, adjust_lr_for_muon -from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5 -from .qk_clip import compute_scales - -logger = logging.getLogger(__name__) - -# ====================================================================== -# Stage helpers -# ====================================================================== - - -def _launch_gather( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Allocate gather buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_gather``). - gathered_grads: ``{id(p): empty_tensor}`` for owned params, - ``None`` for non-owned. - recv_counts: Per-source-rank element counts. - """ - # Allocate gathered-grad buffers - gathered_grads: dict[int, torch.Tensor | None] = {} - for p in params: - state = param_to_state[id(p)] - if rank == state.worker_rank: - gathered_grads[id(p)] = torch.empty(p.shape, - dtype=COMM_DTYPE, - device="cuda") - else: - gathered_grads[id(p)] = None - - # Build send buffer – batch grad copies via torch.cat - # (1-2 fused kernels vs N individual narrow().copy_() calls). - send_counts = [0] * num_ranks - for p in params: - state = param_to_state[id(p)] - send_counts[state.worker_rank] += state.rank_numels[rank] - - total_send = sum(send_counts) - if total_send > 0: - # Group grad slices by destination rank in a single pass. - dst_to_grads = [[] for _ in range(num_ranks)] - for p in params: - state = param_to_state[id(p)] - n = state.rank_numels[rank] - if n > 0: - g = p.grad.to_local() - dst_to_grads[state.worker_rank].append(g.reshape(-1)) - - # Flatten in dst order and cat once. - all_slices = [] - for dst in range(num_ranks): - all_slices.extend(dst_to_grads[dst]) - send_buf = torch.cat(all_slices) - if send_buf.dtype != COMM_DTYPE: - send_buf = send_buf.to(COMM_DTYPE) - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - total += state.rank_numels[src] - recv_counts[src] = total - - recv_buf = torch.empty(sum(recv_counts), dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - logger.debug(f"send_buf size: {send_buf.numel()}, " - f"recv_buf size: {recv_buf.numel()}, " - f"recv_counts: {recv_counts}, " - f"send_counts: {send_counts}, " - f"process_group: {str(process_group)}") - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, gathered_grads, recv_counts - - -def _complete_gather( - recv_buf: torch.Tensor, - recv_counts: list[int], - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - param_to_state: dict[int, _muon_state], - rank: int, -) -> None: - """Reconstruct gathered grads from the recv buffer (in-place).""" - off = 0 - for src in range(len(recv_counts)): - if recv_counts[src] == 0: - continue - - block = recv_counts[src] - inner_off = 0 - for p in owned_params: - state = param_to_state[id(p)] - assert state.worker_rank == rank - - indices = state.rank_indices[src] - - shard_view = gathered_grads[id(p)][indices] - n = shard_view.numel() - if n == 0: - continue - - sg = recv_buf.narrow(0, off + inner_off, n) - sg = sg.reshape(shard_view.shape) - gathered_grads[id(p)][indices] = sg - - inner_off += n - assert inner_off == block - off += block - - -def _compute_ns( - owned_params: list[DTensor], - gathered_grads: dict[int, torch.Tensor | None], - ns_steps: int, -) -> dict[int, torch.Tensor | None]: - """Run Newton-Schulz orthogonalization on owned parameters. - - Returns: - computed_us: ``{id(p): orthogonalized_update}`` for owned params. - """ - computed_us: dict[int, torch.Tensor | None] = {} - for p in owned_params: - u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps) - gathered_grads[id(p)] = None # free gathered grad - computed_us[id(p)] = u - return computed_us - - -def _launch_scatter( - params: list[DTensor], - owned_params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - num_ranks: int, - process_group: dist.ProcessGroup, - computed_us: dict[int, torch.Tensor | None], -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor], list[int]]: - """Allocate scatter buffers, build send/recv, and launch async all-to-all. - - Returns: - work: Async operation handle. - recv_buf: Flat receive buffer (needed by ``_complete_scatter``). - scattered_us: Empty dict, populated by ``_complete_scatter`` with - zero-copy views into ``recv_buf``. - recv_counts: Per-source-rank element counts. - """ - # scattered_us is populated by _complete_scatter with zero-copy views - # into recv_buf, avoiding N empty_like allocations + N copy_ calls. - # Pre-seed entries for params whose local shard is empty (rank_numels == 0) - # so _update_params can iterate all params without KeyError. - scattered_us: dict[int, torch.Tensor] = {} - for p in params: - if param_to_state[id(p)].rank_numels[rank] == 0: - scattered_us[id(p)] = torch.empty_like(p.to_local(), - dtype=COMM_DTYPE) - - # Build send buffer – batch via torch.cat - # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls). - send_counts = [0] * num_ranks - if owned_params: - for p in owned_params: - state = param_to_state[id(p)] - for dst_rank in range(num_ranks): - send_counts[dst_rank] += state.rank_numels[dst_rank] - - total_send = sum(send_counts) - if total_send > 0: - # Cache u_full conversions to avoid redundant .to() per dst_rank. - u_fulls = {} - for p in owned_params: - u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous() - - # Collect slices in dst order (matches all-to-all send layout). - all_slices = [] - for dst_rank in range(num_ranks): - for p in owned_params: - state = param_to_state[id(p)] - su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten() - if su.numel() > 0: - all_slices.append(su) - - send_buf = torch.cat(all_slices) if all_slices else torch.empty( - 0, dtype=COMM_DTYPE, device="cuda") - else: - send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda") - - # Build recv buffer - recv_counts = [0] * num_ranks - for src in range(num_ranks): - total = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - total += state.rank_numels[rank] - recv_counts[src] = total - - recv_total = sum(recv_counts) - recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda") - - # Launch async all-to-all - work = dist.all_to_all_single( - recv_buf, - send_buf, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=process_group, - async_op=True, - ) - - return work, recv_buf, scattered_us, recv_counts - - -def _complete_scatter( - recv_buf: torch.Tensor, - recv_counts: list[int], - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], -) -> None: - """Populate scattered_us with zero-copy views into recv_buf. - - Instead of pre-allocating tensors and copying, we assign views directly - from ``recv_buf``. This eliminates N ``empty_like`` + N ``copy_`` calls. - The underlying storage of ``recv_buf`` is kept alive through the views - until ``scattered_us`` is cleared after ``_update_params``. - """ - off = 0 - for src in range(len(recv_counts)): - block = recv_counts[src] - if block == 0: - continue - - inner_off = 0 - for p in params: - state = param_to_state[id(p)] - if state.worker_rank != src: - continue - n = state.rank_numels[rank] - if n == 0: - continue - - scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off, - n).view_as(p.to_local()) - - inner_off += n - - assert inner_off == block - off += block - - -def _update_params( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - scattered_us: dict[int, torch.Tensor], - lr: float, - weight_decay: float, -) -> None: - """Apply weight decay, Muon update, and optional QK clipping. - - Uses batched ``_foreach_mul_`` for weight decay and batched - ``_foreach_add_`` for the Muon update, grouping parameters by - adjusted_lr to minimize kernel launches while preserving float32 - precision for the alpha scaling. - """ - if not params: - return - - # Batched weight decay: p *= (1 - lr * wd) — single fused kernel. - p_locals = [p._local_tensor for p in params] - torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay) - - # Group params by adjusted_lr so _foreach_add_ can use a single - # alpha per group (preserves float32 precision for alpha scaling). - lr_groups: dict[float, tuple[list, list]] = {} - for p in params: - adjusted_lr = adjust_lr_for_muon(lr, p.shape) - if adjusted_lr not in lr_groups: - lr_groups[adjusted_lr] = ([], []) - lr_groups[adjusted_lr][0].append(p._local_tensor) - lr_groups[adjusted_lr][1].append(scattered_us[id(p)]) - - for adjusted_lr, (p_group, u_group) in lr_groups.items(): - torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr) - - # QK clipping – applied directly on the local tensor to - # avoid DTensor sharding-propagation issues with _StridedShard. - for p in params: - state = param_to_state[id(p)] - if state.qk_clip_state is None: - continue - scales_full = compute_scales(p, state.qk_clip_state) - if scales_full is not None: - ratio = p.shape[0] // scales_full.shape[0] - idx0 = state.rank_indices[rank][0] - if isinstance(idx0, slice): - start = idx0.start or 0 - idx0 = torch.arange(start, - idx0.stop, - device=scales_full.device) - row_scales = scales_full[idx0 // ratio] - p._local_tensor.mul_(row_scales.view(-1, 1)) - - -# ====================================================================== -# Pre-launch helper for overlapping first chunk's gather with other work. -# ====================================================================== - - -@torch.no_grad() -def prelaunch_first_gather( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - none_grad: bool, -) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]: - """Launch the first chunk's A2A gather early for overlap with other compute. - - Call this *before* expensive GPU work (e.g. batched expert NS) so that - the NCCL all-to-all runs concurrently on the NCCL stream while the - default stream executes compute. - - Returns the same 4-tuple that ``_launch_gather`` produces, which should - be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - with record_function("muon::prelaunch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - return work, recv_buf, gathered_grads, recv_counts - - -# ====================================================================== -# Main generator – thin orchestrator that wires stages together. -# ====================================================================== - - -@torch.no_grad() -def muon_chunk_pipeline( - params: list[DTensor], - param_to_state: dict[int, _muon_state], - rank: int, - ns_steps: int, - lr: float, - weight_decay: float, - none_grad: bool, - prelaunch_gather: tuple | None = None, -) -> Generator[None, None, None]: - """Process one chunk of parameters through the full Muon pipeline. - - Stages: gather -> compute (Newton-Schulz) -> scatter -> update. - - Each ``yield`` lets :func:`run_pipeline` interleave other chunks so - that communication and computation overlap across chunks. Async - communication is launched via ``async_op=True`` and completed after - the yield with ``work.wait()``. - - Overlap happens because :func:`run_pipeline` admits one new chunk - per iteration (staggered admission). While chunk *N* does NS - compute on the default CUDA stream, chunk *N+1*'s async all-to-all - runs concurrently on the NCCL stream — no separate ``comm_stream`` - is required. - - If ``prelaunch_gather`` is provided, the gather was already launched - by :func:`prelaunch_first_gather` and we skip launching it again. - - Yields exactly **2** times: - - 1. After launching async all-to-all gather (or immediately if pre-launched). - 2. After launching async all-to-all scatter. - """ - process_group = param_to_state[id(params[0])].process_group - num_ranks = dist.get_world_size(group=process_group) - owned_params = [ - p for p in params if param_to_state[id(p)].worker_rank == rank - ] - - if prelaunch_gather is not None: - # Gather was pre-launched; none_grad already handled by caller. - work, recv_buf, gathered_grads, recv_counts = prelaunch_gather - else: - # Normal path: launch async gather. - with record_function("muon::launch_gather"): - work, recv_buf, gathered_grads, recv_counts = _launch_gather( - params, owned_params, param_to_state, rank, num_ranks, - process_group) - - if none_grad: - for p in params: - p.grad = None - - yield # --- YIELD 1: other chunks can launch their gather --- - - with record_function("muon::wait_gather"): - work.wait() - _complete_gather(recv_buf, recv_counts, owned_params, gathered_grads, - param_to_state, rank) - del recv_buf - - # Stage 3: Newton-Schulz orthogonalization. - with record_function("muon::newton_schulz"): - computed_us = _compute_ns(owned_params, gathered_grads, ns_steps) - gathered_grads.clear() - - # Stages 4-5: launch async scatter. - with record_function("muon::launch_scatter"): - work, recv_buf, scattered_us, recv_counts = _launch_scatter( - params, owned_params, param_to_state, rank, num_ranks, - process_group, computed_us) - computed_us.clear() - - yield # --- YIELD 2: other chunks can launch their scatter --- - - with record_function("muon::wait_scatter"): - work.wait() - _complete_scatter(recv_buf, recv_counts, params, param_to_state, rank, - scattered_us) - del recv_buf - - # Stage 6: apply parameter updates. - with record_function("muon::update_params"): - _update_params(params, param_to_state, rank, scattered_us, lr, - weight_decay) - scattered_us.clear() diff --git a/torch-ext/optimizer/qk_clip.py b/torch-ext/optimizer/qk_clip.py deleted file mode 100644 index 2aba711b3004b7f09e7141da7ef834bd61cc2430..0000000000000000000000000000000000000000 --- a/torch-ext/optimizer/qk_clip.py +++ /dev/null @@ -1,198 +0,0 @@ -import logging -import math -from dataclasses import dataclass - -import torch -from torch.distributed.tensor import DTensor - -from .core import normalize_fqn - -logger = logging.getLogger(__name__) - - -def parse_qk_layer(name: str) -> tuple[str | None, int]: - """ - Parse a parameter name to check if it is a query/key projection layer - and return (kind, layer_index). - - Supported kinds: - MHA/GQA: 'wq', 'wk', 'q_proj', 'k_proj' - MLA: 'wq_b' (Q up-proj), 'wkv_b' (KV up-proj) - - Returns: - (kind, layer_idx) or (None, -1) if not matched. - - Example: - 'model.3.attn.wq.weight' -> ('wq', 3) - 'model.5.attn.wk.weight' -> ('wk', 5) - 'model.2.attn.q_proj.weight' -> ('q_proj', 2) - 'model.7.attn.k_proj.weight' -> ('k_proj', 7) - 'model.1.attn.wq_b.weight' -> ('wq_b', 1) - 'model.0.attn.wkv_b.weight' -> ('wkv_b', 0) - 'model.4.attn.v_proj.weight' -> (None, -1) - """ - parts = normalize_fqn(name).split('.') - if len(parts) < 3: - return None, -1 - - kind = parts[-2] - - layer_idx = -1 - for part in reversed(parts): - if part.isdigit(): - layer_idx = int(part) - break - - if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'): - return kind, layer_idx - - return None, -1 - - -@dataclass -class QKClipInfo: - """Per-parameter dynamic info computed from config + runtime logits.""" - kind: str | None # 'wq'/'q_proj'/'wq_b' or 'wk'/'k_proj'/'wkv_b' or None - indices: list[int] # which heads to consider for clipping - head_dim: int # from config (qk_head_dim for MLA wq_b) - threshold: float # from config - logit: torch.Tensor | None - - # MLA-specific fields - is_mla: bool = False - qk_nope_head_dim: int = 0 - qk_rope_head_dim: int = 0 - v_head_dim: int = 0 - - -def get_qk_clip_info(clip_config, n, qk_logits): - """Extract QK clipping info for a named parameter. - - Args: - clip_config: QK clipping configuration dict (or None). - MHA/GQA keys: head_dim, threshold, q_indices, k_indices - MLA extra keys: is_mla=True, qk_nope_head_dim, qk_rope_head_dim, v_head_dim - n: Parameter name string. - qk_logits: Dict mapping layer indices to logit tensors (or None). - - Returns: - QKClipInfo instance with clipping configuration for this parameter. - """ - if clip_config is None: - return None - - head_dim = clip_config.get('head_dim') - threshold = clip_config.get('threshold') - kind, layer_idx = parse_qk_layer(n) - is_mla = clip_config.get('is_mla', False) - - logit, indices = None, [] - if qk_logits is not None and kind is not None: - logit = qk_logits[layer_idx] - if isinstance(logit, DTensor): - # In TP settings, qk_logits may be DTensor - # We convert it to full tensor here for simplicity - logit = logit.full_tensor() - - if kind in ('wq_b', 'wq', 'q_proj'): - indices = clip_config.get('q_indices', []) or [] - elif kind in ('wkv_b', 'wk', 'k_proj'): - indices = clip_config.get('k_indices', []) or [] - - if is_mla: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - is_mla=True, - qk_nope_head_dim=clip_config['qk_nope_head_dim'], - qk_rope_head_dim=clip_config['qk_rope_head_dim'], - v_head_dim=clip_config['v_head_dim'], - ) - else: - return QKClipInfo( - kind=kind, - indices=indices, - head_dim=head_dim, - threshold=threshold, - logit=logit, - ) - - -def compute_scales(p, qk_clip_state): - """Compute per-head scaling factors for QK clipping. - - Returns scales tensor (√γ per head) if any head exceeds threshold, else None. - For MLA wkv_b, effective row stride is qk_nope_head_dim + v_head_dim. - """ - kind = qk_clip_state.kind - indices = qk_clip_state.indices - head_dim = qk_clip_state.head_dim - threshold = qk_clip_state.threshold - logit = qk_clip_state.logit - - # Check if any head exceeds threshold before allocating. - head_scales = {} - for logit_idx, head_idx in enumerate(indices): - v_ele = float(logit[logit_idx]) - if v_ele > threshold: - new_scale = math.sqrt(threshold / v_ele) - if head_idx not in head_scales or new_scale < head_scales[head_idx]: - head_scales[head_idx] = new_scale - logger.info( - f"[{kind}] Head {head_idx} exceeded threshold " - f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}" - ) - - if not head_scales: - return None - - # For MLA wkv_b, each KV head spans qk_nope_head_dim + v_head_dim rows - if qk_clip_state.is_mla and kind == 'wkv_b': - effective_head_dim = qk_clip_state.qk_nope_head_dim + qk_clip_state.v_head_dim - else: - effective_head_dim = head_dim - - H_global = p.shape[0] // effective_head_dim - scales_full = torch.ones(H_global, device=p.data.device) - for head_idx, scale in head_scales.items(): - scales_full[head_idx] = scale - return scales_full - - -def qk_clip(p, scales, info): - """Apply per-head scaling to a Q/K projection weight matrix. - - Args: - p: Parameter (nn.Parameter or raw tensor). - scales: [n_heads] tensor, each element = √γ_h. - info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions. - - MLA sub-region scaling per Algorithm 1 (MuonClip): - wq_b: q_nope rows → √γ, q_pe rows → γ - wkv_b: k_nope rows → √γ, v rows → unchanged - """ - W = p.data if isinstance(p, torch.nn.Parameter) else p - - if not info.is_mla: - # MHA/GQA: uniform √γ applied to all rows in each head - W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1)) - return - - # MLA: vectorized sub-region scaling within each head - if info.kind == 'wq_b': - qk_nope = info.qk_nope_head_dim - qk_head_dim = qk_nope + info.qk_rope_head_dim - W_3d = W.view(-1, qk_head_dim, W.shape[1]) # [H, qk_head_dim, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # q_nope → √γ - W_3d[:, qk_nope:, :].mul_((scales * scales).view(-1, 1, - 1)) # q_pe → γ - - elif info.kind == 'wkv_b': - qk_nope = info.qk_nope_head_dim - kv_stride = qk_nope + info.v_head_dim - W_3d = W.view(-1, kv_stride, W.shape[1]) # [H, kv_stride, in_dim] - W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1)) # k_nope → √γ - # v rows: not touched (k_R shared rotary unchanged)