test

.github/actionlint.yaml

@@ -1,3 +0,0 @@
-self-hosted-runner:
-  labels:
-    - docker-builder-01

.github/workflows/build-and-commit.yml

@@ -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

.github/workflows/pre-commit.yml

@@ -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

.github/workflows/push-to-hf.yml

@@ -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

.gitignore

@@ -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/

.pre-commit-config.yaml

@@ -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]

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,72 +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,
 )
 ```
-
-## 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)

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

build/torch210-cxx11-cu126-x86_64-linux/distributed/utils.py

@@ -1,175 +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 get_slices_of_dtensor(
-    target: DTensor | torch.Tensor,
-    local_rank: int,
-    shard_mesh: DeviceMesh,
-    shard_placements: tuple[Placement],
-) -> tuple[slice]:
-    """
-    Get the slice of local tensor for a given rank from a tensor.
-    Args:
-        target (DTensor | torch.Tensor): The target tensor.
-        rank (int): The local rank of the shard group.
-        shard_mesh (DeviceMesh): The shard mesh. It consists of global ranks.
-        shard_placements (tuple[Placement]): The shard placements.
-    """
-
-    slices: list[slice] = [slice(0, dim_size) for dim_size in target.size()]
-
-    # 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)
-
-    # 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 i, (rank_coord,
-            placement) in enumerate(zip(rank_coords, shard_placements)):
-        assert isinstance(placement, Shard)
-
-        num_ranks = shard_mesh.mesh.shape[i]
-
-        dim = placement.dim
-        dim_size = (slices[dim].stop - slices[dim].start)
-
-        if dim_size % num_ranks != 0:
-            raise NotImplementedError(
-                f"Dimension size {dim_size} is not divisible "
-                f"by number of ranks {num_ranks} for shard "
-                f"placement on dim {dim}. (shape: {target.shape})")
-
-        shard_size = dim_size // num_ranks
-
-        start = slices[dim].start + rank_coord * shard_size
-        end = start + shard_size
-
-        assert start < end <= slices[dim].stop
-
-        slices[dim] = slice(start, end)
-
-    return tuple(slices)
-
-
-_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh,
-                                                  ProcessGroup]] = dict()
-
-
-def construct_shard_mesh(
-    placements: tuple[Placement],
-    mesh: DeviceMesh,
-) -> (DeviceMesh, ProcessGroup, tuple[Placement]):
-    """
-    Construct Shard Mesh and Placements for unsharding.
-    It removes Replicate placements and constructs a new Mesh and ProcessGroup.
-    """
-    my_rank = dist.get_rank()
-
-    assert mesh.mesh.device.type == 'cpu'
-
-    # Copy mesh to avoid modifying the original mesh
-    mesh = mesh.mesh.clone()
-
-    # 1. Sort placements. Replicate first, then Shard by dim ascending.
-
-    # For Shard, strided shard comes after regular shard on the same dim
-    # to preserve left-to-right order of replicate-to-shard.
-    # This is because that strided shard is using stride to represent
-    # more fine-grained sharding on the same dim.
-    # Please check the URL below for _StridedShard.
-    # https://github.com/pytorch/pytorch/blob/v2.8.0/torch/distributed/tensor/placement_types.py#L366
-
-    def placement_sort_key(
-        placement_with_index: tuple[float, Placement]
-    ) -> tuple[int, float, int]:  # (dim, split factor, original index)
-        index, placement = placement_with_index
-        is_replicate = placement.is_replicate()
-        is_shard = placement.is_shard()
-        is_partial = placement.is_partial()
-
-        assert is_replicate or is_shard, f"Unsupported placement type: {type(placement)}"
-        assert not is_partial, "Partial placement is not supported."
-
-        if is_replicate:
-            return (-1.0, 0, index)
-        elif is_shard:
-            if isinstance(placement, _StridedShard):
-                return (placement.dim, 1 / placement.split_factor, index)
-            return (placement.dim, 0, index)
-        else:
-            raise TypeError(f"Unknown placement type: {type(placement)}")
-
-    placements_with_index: list[tuple[int,
-                                      Placement]] = list(enumerate(placements))
-    placements_with_index = sorted(placements_with_index,
-                                   key=placement_sort_key)
-
-    sorted_indices, sorted_placements = zip(*placements_with_index)
-
-    # 2. Permute mesh according to sorted placements.
-    sorted_mesh = mesh.permute(sorted_indices)
-
-    # 3. Collect list of shard meshes by removing replicate dims
-    # For example, (2, 3, 4, 4) with placements [R, R, S(0), S(1)]
-    # shard_meshes should be list with 2 * 3 = 6 shard meshes of shape (4, 4)
-    num_replicates = sum(1 for p in sorted_placements if p.is_replicate())
-
-    # merge replicate dims
-    # shard_meshes became a list of shard meshes with a length of replicate degree
-    if num_replicates > 0:
-        sorted_mesh = sorted_mesh.flatten(
-            0, num_replicates - 1) if num_replicates > 1 else sorted_mesh
-        shard_meshes = list(torch.unbind(sorted_mesh, dim=0))
-    else:
-        shard_meshes = [sorted_mesh]
-    shard_placements = sorted_placements[num_replicates:]
-
-    # assume all shard placements are different
-    assert len(shard_placements) == len(set(shard_placements))
-
-    # 4. Construct ProcessGroups
-    # Caution: all groups should be created in the same order in all processes,
-    # even though each process only needs its own group.
-
-    # To use tensor as dict key, convert it to tuple
-    def tensor_to_tuple(t):
-        if isinstance(t, torch.Tensor):
-            t = t.tolist()
-        if isinstance(t, list):
-            return tuple(tensor_to_tuple(x) for x in t)
-        return t
-
-    my_shard_mesh_as_tuple = None
-    for shard_mesh in shard_meshes:
-        assert isinstance(shard_mesh, torch.Tensor)
-        shard_mesh_as_tuple = tensor_to_tuple(shard_mesh)
-
-        if (my_rank == shard_mesh).any().item():
-            assert my_shard_mesh_as_tuple is None
-            my_shard_mesh_as_tuple = shard_mesh_as_tuple
-
-        # update global cache
-        if shard_mesh_as_tuple not in _ranks_to_dist_cache:
-            shard_process_group = dist.new_group(shard_mesh.flatten().tolist())
-            _ranks_to_dist_cache[shard_mesh_as_tuple] = (
-                DeviceMesh(device_type="cuda", mesh=shard_mesh),
-                shard_process_group,
-            )
-
-    my_shard_mesh, my_shard_process_group = _ranks_to_dist_cache[
-        my_shard_mesh_as_tuple]
-
-    return my_shard_mesh, my_shard_process_group, shard_placements

build/torch210-cxx11-cu126-x86_64-linux/matmul_transpose_triton.py

@@ -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

build/torch210-cxx11-cu126-x86_64-linux/metadata.json

@@ -1 +0,0 @@
-{"python-depends":[]}

build/torch210-cxx11-cu126-x86_64-linux/muon.py

@@ -1,1268 +0,0 @@
-import logging
-import math
-import types
-from collections import defaultdict
-from dataclasses import dataclass
-from typing import Any, cast
-
-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, Replicate
-from torch.distributed.tensor.placement_types import Placement
-
-from .distributed.utils import construct_shard_mesh, get_slices_of_dtensor
-from .matmul_transpose_triton import matmul_transpose_assign
-
-logger = logging.getLogger(__name__)
-
-COMM_DTYPE = torch.bfloat16
-DEFAULT_CHUNK_SIZE_RATIO = 4
-
-
-# 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
-    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
-    assert G.dtype == COMM_DTYPE
-    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)
-    # 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)
-
-    if G.size(0) > G.size(1):
-        X = X.T
-    return X
-
-
-@dataclass
-class _muon_state:
-    # TODO: use Optional
-    worker_rank: int
-    process_group: ProcessGroup
-    shard_mesh: DeviceMesh
-    shard_placements: tuple[Placement, ...]
-    name: str
-    qk_clip_state: torch.Tensor | 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
-
-
-def numel_for_rank(
-    param: DTensor,
-    local_rank: int,
-    state: _muon_state,
-) -> int:
-    slices = get_slices_of_dtensor(
-        param,
-        local_rank,
-        state.shard_mesh,
-        state.shard_placements,
-    )
-
-    numel = 1
-    for s, dim in zip(slices, param.shape):
-        start, stop, step = s.indices(dim)
-        length = max(0, (stop - start + (step - 1)) // step)
-        numel *= length
-
-    return numel
-
-
-@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:
-                state.gathered_grad = torch.empty(p.shape,
-                                                  dtype=COMM_DTYPE,
-                                                  device="cuda")
-            else:
-                state.gathered_grad = None
-
-        alloc_event = torch.cuda.Event()
-        alloc_event.record(compute_stream)
-        return alloc_event
-
-
-@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 = numel_for_rank(p, rank, state)
-            g = p.grad
-            g = g.to_local().to(COMM_DTYPE).contiguous()
-            assert g.numel() == shard_elems
-            per_dst[dst].append(g.view(-1))
-            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 += numel_for_rank(p, src, state)
-            recv_counts[src] = total
-
-        recv_total = sum(recv_counts)
-        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
-
-        #All2All
-        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)}")
-        dist.all_to_all_single(
-            recv_buf,
-            send_buf,
-            output_split_sizes=recv_counts,
-            input_split_sizes=send_counts,
-            group=process_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
-        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
-
-                # get the slice of the full dtensor corresponding to rank src.
-                slices = get_slices_of_dtensor(state.gathered_grad, src,
-                                               state.shard_mesh,
-                                               state.shard_placements)
-
-                dst = state.gathered_grad[slices]
-                assert dst._base is state.gathered_grad
-
-                n = dst.numel()
-                assert n > 0
-
-                sg = recv_buf.narrow(0, off + inner_off, n)
-                sg = sg.reshape_as(dst)
-                dst.copy_(sg)
-
-                inner_off += n
-            off += block
-
-        for p in params:
-            state = param_to_state[id(p)]
-            if state.worker_rank == rank:
-                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
-
-
-@torch.no_grad()
-def _compute_u(p, state, steps, rank, compute_stream):
-    """
-    On worker_rank, compute the orthogonalized update using Newton-Schulz iteration.
-    """
-    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()
-        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 _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()
-
-                offset = 0
-                for dst in range(num_ranks):
-                    # get the slice of the full tensor corresponding to rank dst.
-                    slices = get_slices_of_dtensor(u_full, dst,
-                                                   state.shard_mesh,
-                                                   state.shard_placements)
-                    su = u_full[slices].flatten()
-
-                    n = su.numel()
-                    assert n > 0
-
-                    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)
-        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 += numel_for_rank(p, rank, state)
-            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,
-        )
-
-        # 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 = numel_for_rank(p, rank, state)
-                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
-            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 state.qk_clip_state is not None else None
-        if scales_full is not None:
-            # Have to slice scales_full among dim 0
-            weight_slices = get_slices_of_dtensor(p, rank, state.shard_mesh,
-                                                  state.shard_placements)
-            ratio = p.shape[0] // scales_full.shape[0]
-            scales_slice = slice(
-                None if weight_slices[0].start is None else
-                weight_slices[0].start // ratio,
-                None if weight_slices[0].stop is None else
-                weight_slices[0].stop // ratio,
-                None,
-            )
-
-            scales_local = scales_full[scales_slice]
-            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: str | None  # '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: torch.Tensor | None
-
-
-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.
-        {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.
-        small_param_numel_threshold: Threshold for classifying parameters as small and falling back to distributed Muon
-    """
-
-    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
-                 },
-                 warmup_step=5,
-                 chunk_size=-1,
-                 use_distributed_muon=False,
-                 small_param_numel_threshold=65536):
-        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.rank = None
-
-        self.comm_stream = torch.cuda.Stream()
-        self.compute_stream = torch.cuda.Stream()
-        self.debug = debug
-        self.clip_config = clip_config
-        self.warmup_step = warmup_step
-        self.chunk_size = chunk_size
-        self.use_distributed_muon = use_distributed_muon
-        self.small_param_numel_threshold = small_param_numel_threshold
-
-    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 set_rank_once(self, rank):
-        if self.rank is None:
-            self.rank = rank
-        else:
-            assert self.rank == rank
-
-    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)
-
-        # set rank with the local rank in the shard process group
-        self.set_rank_once(dist.get_rank(group=shard_pg))
-
-        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:
-            print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs",
-                  flush=True)
-
-        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 = self.get_qk_clip_info(n, qk_logits)
-
-            param_to_state[id(p)] = _muon_state(
-                worker_rank=worker_rank,
-                process_group=shard_pg,
-                shard_mesh=shard_mesh,
-                shard_placements=shard_placements,
-                name=n,
-                qk_clip_state=qk_clip_state,
-            )
-
-        return param_to_state, ordered_params
-
-    def base(self, names, params, group, lr, weight_decay, momentum,
-             qk_logits):
-        # generate weight updates in distributed fashion
-        for n, p in zip(names, 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
-
-            g = self._update_g(p, g, group, momentum)
-
-            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
-                                             steps=group["ns_steps"])
-
-            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)
-
-            scales_full = self._compute_scales(
-                p, qk_clip_state) if qk_clip_state is not None else None
-            if scales_full is not None:
-                Muon._qk_clip(p, scales_full, qk_clip_state.head_dim)
-
-    def distributed_muon(
-        self,
-        names: list[str],
-        params: list[torch.nn.Parameter],
-        group: dict[str, Any],
-        lr: float,
-        weight_decay: float,
-        momentum: float,
-        qk_logits: list[torch.Tensor | DTensor] | None,
-    ):
-        """ Implementation of Distributed Muon by Liu et al. """
-
-        for n, p in zip(names, 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
-
-            g = self._update_g(p, g, group, momentum)
-
-            # Gather G
-            if isinstance(p.data, DTensor):
-                g_full = g.full_tensor()
-                p_full = p.data.full_tensor()
-            else:
-                g_full = g
-                p_full = p
-
-            u_full = _zeropower_via_newtonschulz5(g_full.to(COMM_DTYPE),
-                                                  steps=group["ns_steps"])
-
-            adjusted_lr = self.adjust_lr_for_muon(lr, p_full.shape)
-            Muon._update_p(p_full, u_full, lr, adjusted_lr, weight_decay)
-
-            qk_clip_state = self.get_qk_clip_info(n, qk_logits)
-
-            scales_full = self._compute_scales(
-                p_full, qk_clip_state) if qk_clip_state is not None else None
-
-            if scales_full is not None:
-                Muon._qk_clip(p_full, scales_full, qk_clip_state.head_dim)
-
-            if isinstance(p.data, DTensor):
-                ndims = len(p.device_mesh.mesh.shape)
-                p_replicate = DTensor.from_local(
-                    p_full,
-                    device_mesh=p.device_mesh,
-                    placements=[Replicate() for _ in range(ndims)],
-                )
-
-                p_sharded = p_replicate.redistribute(
-                    device_mesh=p.device_mesh,
-                    placements=p.placements,
-                )
-
-                p.copy_(p_sharded)
-
-    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 group["nesterov"]:
-            g.add_(buf, alpha=momentum)
-            return g
-        return buf
-
-    @staticmethod
-    def _update_p(p, u, lr, adjusted_lr, weight_decay):
-        if isinstance(p, torch.nn.Parameter):
-            # apply weight decay
-            p.data.mul_(1 - lr * weight_decay)
-            # apply update
-            p.data.add_(u, alpha=-adjusted_lr)
-        else:
-            p.mul_(1 - lr * weight_decay)
-            p.add_(u, alpha=-adjusted_lr)
-
-    def get_qk_clip_info(self, n, qk_logits):
-        if self.clip_config is None:
-            return None
-
-        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 []
-
-            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()
-
-        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):
-        if isinstance(p, torch.nn.Parameter):
-            W = p.data.view(-1, head_dim, p.data.shape[1])
-            W.mul_(scales.view(-1, 1, 1))
-        else:
-            W = p.view(-1, head_dim, p.shape[1])
-            W.mul_(scales.view(-1, 1, 1))
-
-    def parallel(self, names, params, group, lr, weight_decay, momentum,
-                 qk_logits):
-        """
-        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(
-            names, params, group, qk_logits)
-
-        assert self.rank is not None
-
-        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_computes(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]:
-                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)
-
-        if self.chunk_size == -1:
-            shard_ranks = dist.get_world_size(param_to_state[id(
-                params[0])].process_group)
-            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.")
-
-        # Wait grad update
-        self.comm_stream.wait_stream(torch.cuda.current_stream())
-
-        warmup_step = self.warmup_step
-        for i in range(0, warmup_step):
-            enqueue_all2all_gather(i * chunk_size, chunk_size)
-            enqueue_computes(i * chunk_size, chunk_size)
-
-        for i in range(0, len(params) + chunk_size - 1, chunk_size):
-            enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + warmup_step * chunk_size, chunk_size)
-            enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + warmup_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: 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: DeviceDict | 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 _step_muon(self, group, qk_logits=None):
-        params = group["params"]
-        lr = group["lr"]
-        weight_decay = group["weight_decay"]
-        momentum = group["momentum"]
-        names = group["names"]
-
-        param_dtensors = []
-        name_dtensors = []
-
-        param_tensors = []
-        name_tensors = []
-
-        param_dtensors_small = []
-        name_dtensors_small = []
-
-        if self.use_distributed_muon:
-            self.distributed_muon(names=names,
-                                  params=params,
-                                  group=group,
-                                  lr=lr,
-                                  weight_decay=weight_decay,
-                                  momentum=momentum,
-                                  qk_logits=qk_logits)
-            return
-
-        # For simplicity, we use distributed Muon for small parameters
-        # whose number of elements is below a threshold.
-        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)
-                elif p.data.numel() <= self.small_param_numel_threshold:
-                    param_dtensors_small.append(p)
-                    name_dtensors_small.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)}")
-
-        logger.debug(
-            f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors, "
-            f"{len(param_dtensors_small)} Small DTensors")
-
-        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: ([], []))
-            # type: dict[tuple[Placement, DeviceMesh], tuple[list[str], list[DTensor]]]
-
-            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_small) > 0:
-            if not dist.is_initialized():
-                raise RuntimeError(
-                    "Parallel Muon requires torch.distributed to be initialized."
-                )
-
-            self.distributed_muon(
-                params=param_dtensors_small,
-                names=name_dtensors_small,
-                group=group,
-                lr=lr,
-                weight_decay=weight_decay,
-                momentum=momentum,
-                qk_logits=qk_logits,
-            )
-
-        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)
-            for _, (names, params) in dtensor_group.items():
-                self.parallel(
-                    names,
-                    params,
-                    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,
-            )
-
-    def _step_adamw_params(self, params, group):
-        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,
-            lr=lr,
-            weight_decay=weight_decay,
-            eps=eps,
-            maximize=False,
-        )
-
-    def _step_adamw(self, group):
-        params = group["params"]
-
-        # group params with it's type and placement
-        placement_to_params: dict[tuple[Placement | type,
-                                        DeviceMesh | None]] = defaultdict(list)
-        for p in params:
-            match p:
-                case DTensor():
-                    placement_to_params[tuple([p.placements,
-                                               p.device_mesh])].append(p)
-                case torch.Tensor():
-                    placement_to_params[tuple([torch.Tensor, None])].append(p)
-
-        for params in placement_to_params.values():
-            self._step_adamw_params(params, group)
-
-    @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()
-
-        for group in self.param_groups:
-            if group["use_muon"]:
-                self._step_muon(group, qk_logits=qk_logits)
-            else:
-                self._step_adamw(group)
-
-        return loss

build/torch210-cxx11-cu126-x86_64-linux/optimizer/__init__.py

@@ -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")))

build/torch210-cxx11-cu128-x86_64-linux/distributed/utils.py

@@ -1,175 +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 get_slices_of_dtensor(
-    target: DTensor | torch.Tensor,
-    local_rank: int,
-    shard_mesh: DeviceMesh,
-    shard_placements: tuple[Placement],
-) -> tuple[slice]:
-    """
-    Get the slice of local tensor for a given rank from a tensor.
-    Args:
-        target (DTensor | torch.Tensor): The target tensor.
-        rank (int): The local rank of the shard group.
-        shard_mesh (DeviceMesh): The shard mesh. It consists of global ranks.
-        shard_placements (tuple[Placement]): The shard placements.
-    """
-
-    slices: list[slice] = [slice(0, dim_size) for dim_size in target.size()]
-
-    # 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)
-
-    # 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 i, (rank_coord,
-            placement) in enumerate(zip(rank_coords, shard_placements)):
-        assert isinstance(placement, Shard)
-
-        num_ranks = shard_mesh.mesh.shape[i]
-
-        dim = placement.dim
-        dim_size = (slices[dim].stop - slices[dim].start)
-
-        if dim_size % num_ranks != 0:
-            raise NotImplementedError(
-                f"Dimension size {dim_size} is not divisible "
-                f"by number of ranks {num_ranks} for shard "
-                f"placement on dim {dim}. (shape: {target.shape})")
-
-        shard_size = dim_size // num_ranks
-
-        start = slices[dim].start + rank_coord * shard_size
-        end = start + shard_size
-
-        assert start < end <= slices[dim].stop
-
-        slices[dim] = slice(start, end)
-
-    return tuple(slices)
-
-
-_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh,
-                                                  ProcessGroup]] = dict()
-
-
-def construct_shard_mesh(
-    placements: tuple[Placement],
-    mesh: DeviceMesh,
-) -> (DeviceMesh, ProcessGroup, tuple[Placement]):
-    """
-    Construct Shard Mesh and Placements for unsharding.
-    It removes Replicate placements and constructs a new Mesh and ProcessGroup.
-    """
-    my_rank = dist.get_rank()
-
-    assert mesh.mesh.device.type == 'cpu'
-
-    # Copy mesh to avoid modifying the original mesh
-    mesh = mesh.mesh.clone()
-
-    # 1. Sort placements. Replicate first, then Shard by dim ascending.
-
-    # For Shard, strided shard comes after regular shard on the same dim
-    # to preserve left-to-right order of replicate-to-shard.
-    # This is because that strided shard is using stride to represent
-    # more fine-grained sharding on the same dim.
-    # Please check the URL below for _StridedShard.
-    # https://github.com/pytorch/pytorch/blob/v2.8.0/torch/distributed/tensor/placement_types.py#L366
-
-    def placement_sort_key(
-        placement_with_index: tuple[float, Placement]
-    ) -> tuple[int, float, int]:  # (dim, split factor, original index)
-        index, placement = placement_with_index
-        is_replicate = placement.is_replicate()
-        is_shard = placement.is_shard()
-        is_partial = placement.is_partial()
-
-        assert is_replicate or is_shard, f"Unsupported placement type: {type(placement)}"
-        assert not is_partial, "Partial placement is not supported."
-
-        if is_replicate:
-            return (-1.0, 0, index)
-        elif is_shard:
-            if isinstance(placement, _StridedShard):
-                return (placement.dim, 1 / placement.split_factor, index)
-            return (placement.dim, 0, index)
-        else:
-            raise TypeError(f"Unknown placement type: {type(placement)}")
-
-    placements_with_index: list[tuple[int,
-                                      Placement]] = list(enumerate(placements))
-    placements_with_index = sorted(placements_with_index,
-                                   key=placement_sort_key)
-
-    sorted_indices, sorted_placements = zip(*placements_with_index)
-
-    # 2. Permute mesh according to sorted placements.
-    sorted_mesh = mesh.permute(sorted_indices)
-
-    # 3. Collect list of shard meshes by removing replicate dims
-    # For example, (2, 3, 4, 4) with placements [R, R, S(0), S(1)]
-    # shard_meshes should be list with 2 * 3 = 6 shard meshes of shape (4, 4)
-    num_replicates = sum(1 for p in sorted_placements if p.is_replicate())
-
-    # merge replicate dims
-    # shard_meshes became a list of shard meshes with a length of replicate degree
-    if num_replicates > 0:
-        sorted_mesh = sorted_mesh.flatten(
-            0, num_replicates - 1) if num_replicates > 1 else sorted_mesh
-        shard_meshes = list(torch.unbind(sorted_mesh, dim=0))
-    else:
-        shard_meshes = [sorted_mesh]
-    shard_placements = sorted_placements[num_replicates:]
-
-    # assume all shard placements are different
-    assert len(shard_placements) == len(set(shard_placements))
-
-    # 4. Construct ProcessGroups
-    # Caution: all groups should be created in the same order in all processes,
-    # even though each process only needs its own group.
-
-    # To use tensor as dict key, convert it to tuple
-    def tensor_to_tuple(t):
-        if isinstance(t, torch.Tensor):
-            t = t.tolist()
-        if isinstance(t, list):
-            return tuple(tensor_to_tuple(x) for x in t)
-        return t
-
-    my_shard_mesh_as_tuple = None
-    for shard_mesh in shard_meshes:
-        assert isinstance(shard_mesh, torch.Tensor)
-        shard_mesh_as_tuple = tensor_to_tuple(shard_mesh)
-
-        if (my_rank == shard_mesh).any().item():
-            assert my_shard_mesh_as_tuple is None
-            my_shard_mesh_as_tuple = shard_mesh_as_tuple
-
-        # update global cache
-        if shard_mesh_as_tuple not in _ranks_to_dist_cache:
-            shard_process_group = dist.new_group(shard_mesh.flatten().tolist())
-            _ranks_to_dist_cache[shard_mesh_as_tuple] = (
-                DeviceMesh(device_type="cuda", mesh=shard_mesh),
-                shard_process_group,
-            )
-
-    my_shard_mesh, my_shard_process_group = _ranks_to_dist_cache[
-        my_shard_mesh_as_tuple]
-
-    return my_shard_mesh, my_shard_process_group, shard_placements

build/torch210-cxx11-cu128-x86_64-linux/matmul_transpose_triton.py

@@ -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

build/torch210-cxx11-cu128-x86_64-linux/metadata.json

@@ -1 +0,0 @@
-{"python-depends":[]}

build/torch210-cxx11-cu128-x86_64-linux/muon.py

@@ -1,1268 +0,0 @@
-import logging
-import math
-import types
-from collections import defaultdict
-from dataclasses import dataclass
-from typing import Any, cast
-
-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, Replicate
-from torch.distributed.tensor.placement_types import Placement
-
-from .distributed.utils import construct_shard_mesh, get_slices_of_dtensor
-from .matmul_transpose_triton import matmul_transpose_assign
-
-logger = logging.getLogger(__name__)
-
-COMM_DTYPE = torch.bfloat16
-DEFAULT_CHUNK_SIZE_RATIO = 4
-
-
-# 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
-    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
-    assert G.dtype == COMM_DTYPE
-    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)
-    # 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)
-
-    if G.size(0) > G.size(1):
-        X = X.T
-    return X
-
-
-@dataclass
-class _muon_state:
-    # TODO: use Optional
-    worker_rank: int
-    process_group: ProcessGroup
-    shard_mesh: DeviceMesh
-    shard_placements: tuple[Placement, ...]
-    name: str
-    qk_clip_state: torch.Tensor | 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
-
-
-def numel_for_rank(
-    param: DTensor,
-    local_rank: int,
-    state: _muon_state,
-) -> int:
-    slices = get_slices_of_dtensor(
-        param,
-        local_rank,
-        state.shard_mesh,
-        state.shard_placements,
-    )
-
-    numel = 1
-    for s, dim in zip(slices, param.shape):
-        start, stop, step = s.indices(dim)
-        length = max(0, (stop - start + (step - 1)) // step)
-        numel *= length
-
-    return numel
-
-
-@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:
-                state.gathered_grad = torch.empty(p.shape,
-                                                  dtype=COMM_DTYPE,
-                                                  device="cuda")
-            else:
-                state.gathered_grad = None
-
-        alloc_event = torch.cuda.Event()
-        alloc_event.record(compute_stream)
-        return alloc_event
-
-
-@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 = numel_for_rank(p, rank, state)
-            g = p.grad
-            g = g.to_local().to(COMM_DTYPE).contiguous()
-            assert g.numel() == shard_elems
-            per_dst[dst].append(g.view(-1))
-            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 += numel_for_rank(p, src, state)
-            recv_counts[src] = total
-
-        recv_total = sum(recv_counts)
-        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
-
-        #All2All
-        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)}")
-        dist.all_to_all_single(
-            recv_buf,
-            send_buf,
-            output_split_sizes=recv_counts,
-            input_split_sizes=send_counts,
-            group=process_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
-        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
-
-                # get the slice of the full dtensor corresponding to rank src.
-                slices = get_slices_of_dtensor(state.gathered_grad, src,
-                                               state.shard_mesh,
-                                               state.shard_placements)
-
-                dst = state.gathered_grad[slices]
-                assert dst._base is state.gathered_grad
-
-                n = dst.numel()
-                assert n > 0
-
-                sg = recv_buf.narrow(0, off + inner_off, n)
-                sg = sg.reshape_as(dst)
-                dst.copy_(sg)
-
-                inner_off += n
-            off += block
-
-        for p in params:
-            state = param_to_state[id(p)]
-            if state.worker_rank == rank:
-                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
-
-
-@torch.no_grad()
-def _compute_u(p, state, steps, rank, compute_stream):
-    """
-    On worker_rank, compute the orthogonalized update using Newton-Schulz iteration.
-    """
-    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()
-        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 _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()
-
-                offset = 0
-                for dst in range(num_ranks):
-                    # get the slice of the full tensor corresponding to rank dst.
-                    slices = get_slices_of_dtensor(u_full, dst,
-                                                   state.shard_mesh,
-                                                   state.shard_placements)
-                    su = u_full[slices].flatten()
-
-                    n = su.numel()
-                    assert n > 0
-
-                    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)
-        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 += numel_for_rank(p, rank, state)
-            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,
-        )
-
-        # 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 = numel_for_rank(p, rank, state)
-                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
-            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 state.qk_clip_state is not None else None
-        if scales_full is not None:
-            # Have to slice scales_full among dim 0
-            weight_slices = get_slices_of_dtensor(p, rank, state.shard_mesh,
-                                                  state.shard_placements)
-            ratio = p.shape[0] // scales_full.shape[0]
-            scales_slice = slice(
-                None if weight_slices[0].start is None else
-                weight_slices[0].start // ratio,
-                None if weight_slices[0].stop is None else
-                weight_slices[0].stop // ratio,
-                None,
-            )
-
-            scales_local = scales_full[scales_slice]
-            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: str | None  # '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: torch.Tensor | None
-
-
-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.
-        {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.
-        small_param_numel_threshold: Threshold for classifying parameters as small and falling back to distributed Muon
-    """
-
-    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
-                 },
-                 warmup_step=5,
-                 chunk_size=-1,
-                 use_distributed_muon=False,
-                 small_param_numel_threshold=65536):
-        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.rank = None
-
-        self.comm_stream = torch.cuda.Stream()
-        self.compute_stream = torch.cuda.Stream()
-        self.debug = debug
-        self.clip_config = clip_config
-        self.warmup_step = warmup_step
-        self.chunk_size = chunk_size
-        self.use_distributed_muon = use_distributed_muon
-        self.small_param_numel_threshold = small_param_numel_threshold
-
-    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 set_rank_once(self, rank):
-        if self.rank is None:
-            self.rank = rank
-        else:
-            assert self.rank == rank
-
-    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)
-
-        # set rank with the local rank in the shard process group
-        self.set_rank_once(dist.get_rank(group=shard_pg))
-
-        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:
-            print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs",
-                  flush=True)
-
-        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 = self.get_qk_clip_info(n, qk_logits)
-
-            param_to_state[id(p)] = _muon_state(
-                worker_rank=worker_rank,
-                process_group=shard_pg,
-                shard_mesh=shard_mesh,
-                shard_placements=shard_placements,
-                name=n,
-                qk_clip_state=qk_clip_state,
-            )
-
-        return param_to_state, ordered_params
-
-    def base(self, names, params, group, lr, weight_decay, momentum,
-             qk_logits):
-        # generate weight updates in distributed fashion
-        for n, p in zip(names, 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
-
-            g = self._update_g(p, g, group, momentum)
-
-            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
-                                             steps=group["ns_steps"])
-
-            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)
-
-            scales_full = self._compute_scales(
-                p, qk_clip_state) if qk_clip_state is not None else None
-            if scales_full is not None:
-                Muon._qk_clip(p, scales_full, qk_clip_state.head_dim)
-
-    def distributed_muon(
-        self,
-        names: list[str],
-        params: list[torch.nn.Parameter],
-        group: dict[str, Any],
-        lr: float,
-        weight_decay: float,
-        momentum: float,
-        qk_logits: list[torch.Tensor | DTensor] | None,
-    ):
-        """ Implementation of Distributed Muon by Liu et al. """
-
-        for n, p in zip(names, 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
-
-            g = self._update_g(p, g, group, momentum)
-
-            # Gather G
-            if isinstance(p.data, DTensor):
-                g_full = g.full_tensor()
-                p_full = p.data.full_tensor()
-            else:
-                g_full = g
-                p_full = p
-
-            u_full = _zeropower_via_newtonschulz5(g_full.to(COMM_DTYPE),
-                                                  steps=group["ns_steps"])
-
-            adjusted_lr = self.adjust_lr_for_muon(lr, p_full.shape)
-            Muon._update_p(p_full, u_full, lr, adjusted_lr, weight_decay)
-
-            qk_clip_state = self.get_qk_clip_info(n, qk_logits)
-
-            scales_full = self._compute_scales(
-                p_full, qk_clip_state) if qk_clip_state is not None else None
-
-            if scales_full is not None:
-                Muon._qk_clip(p_full, scales_full, qk_clip_state.head_dim)
-
-            if isinstance(p.data, DTensor):
-                ndims = len(p.device_mesh.mesh.shape)
-                p_replicate = DTensor.from_local(
-                    p_full,
-                    device_mesh=p.device_mesh,
-                    placements=[Replicate() for _ in range(ndims)],
-                )
-
-                p_sharded = p_replicate.redistribute(
-                    device_mesh=p.device_mesh,
-                    placements=p.placements,
-                )
-
-                p.copy_(p_sharded)
-
-    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 group["nesterov"]:
-            g.add_(buf, alpha=momentum)
-            return g
-        return buf
-
-    @staticmethod
-    def _update_p(p, u, lr, adjusted_lr, weight_decay):
-        if isinstance(p, torch.nn.Parameter):
-            # apply weight decay
-            p.data.mul_(1 - lr * weight_decay)
-            # apply update
-            p.data.add_(u, alpha=-adjusted_lr)
-        else:
-            p.mul_(1 - lr * weight_decay)
-            p.add_(u, alpha=-adjusted_lr)
-
-    def get_qk_clip_info(self, n, qk_logits):
-        if self.clip_config is None:
-            return None
-
-        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 []
-
-            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()
-
-        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):
-        if isinstance(p, torch.nn.Parameter):
-            W = p.data.view(-1, head_dim, p.data.shape[1])
-            W.mul_(scales.view(-1, 1, 1))
-        else:
-            W = p.view(-1, head_dim, p.shape[1])
-            W.mul_(scales.view(-1, 1, 1))
-
-    def parallel(self, names, params, group, lr, weight_decay, momentum,
-                 qk_logits):
-        """
-        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(
-            names, params, group, qk_logits)
-
-        assert self.rank is not None
-
-        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_computes(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]:
-                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)
-
-        if self.chunk_size == -1:
-            shard_ranks = dist.get_world_size(param_to_state[id(
-                params[0])].process_group)
-            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.")
-
-        # Wait grad update
-        self.comm_stream.wait_stream(torch.cuda.current_stream())
-
-        warmup_step = self.warmup_step
-        for i in range(0, warmup_step):
-            enqueue_all2all_gather(i * chunk_size, chunk_size)
-            enqueue_computes(i * chunk_size, chunk_size)
-
-        for i in range(0, len(params) + chunk_size - 1, chunk_size):
-            enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + warmup_step * chunk_size, chunk_size)
-            enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + warmup_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: 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: DeviceDict | 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 _step_muon(self, group, qk_logits=None):
-        params = group["params"]
-        lr = group["lr"]
-        weight_decay = group["weight_decay"]
-        momentum = group["momentum"]
-        names = group["names"]
-
-        param_dtensors = []
-        name_dtensors = []
-
-        param_tensors = []
-        name_tensors = []
-
-        param_dtensors_small = []
-        name_dtensors_small = []
-
-        if self.use_distributed_muon:
-            self.distributed_muon(names=names,
-                                  params=params,
-                                  group=group,
-                                  lr=lr,
-                                  weight_decay=weight_decay,
-                                  momentum=momentum,
-                                  qk_logits=qk_logits)
-            return
-
-        # For simplicity, we use distributed Muon for small parameters
-        # whose number of elements is below a threshold.
-        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)
-                elif p.data.numel() <= self.small_param_numel_threshold:
-                    param_dtensors_small.append(p)
-                    name_dtensors_small.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)}")
-
-        logger.debug(
-            f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors, "
-            f"{len(param_dtensors_small)} Small DTensors")
-
-        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: ([], []))
-            # type: dict[tuple[Placement, DeviceMesh], tuple[list[str], list[DTensor]]]
-
-            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_small) > 0:
-            if not dist.is_initialized():
-                raise RuntimeError(
-                    "Parallel Muon requires torch.distributed to be initialized."
-                )
-
-            self.distributed_muon(
-                params=param_dtensors_small,
-                names=name_dtensors_small,
-                group=group,
-                lr=lr,
-                weight_decay=weight_decay,
-                momentum=momentum,
-                qk_logits=qk_logits,
-            )
-
-        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)
-            for _, (names, params) in dtensor_group.items():
-                self.parallel(
-                    names,
-                    params,
-                    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,
-            )
-
-    def _step_adamw_params(self, params, group):
-        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,
-            lr=lr,
-            weight_decay=weight_decay,
-            eps=eps,
-            maximize=False,
-        )
-
-    def _step_adamw(self, group):
-        params = group["params"]
-
-        # group params with it's type and placement
-        placement_to_params: dict[tuple[Placement | type,
-                                        DeviceMesh | None]] = defaultdict(list)
-        for p in params:
-            match p:
-                case DTensor():
-                    placement_to_params[tuple([p.placements,
-                                               p.device_mesh])].append(p)
-                case torch.Tensor():
-                    placement_to_params[tuple([torch.Tensor, None])].append(p)
-
-        for params in placement_to_params.values():
-            self._step_adamw_params(params, group)
-
-    @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()
-
-        for group in self.param_groups:
-            if group["use_muon"]:
-                self._step_muon(group, qk_logits=qk_logits)
-            else:
-                self._step_adamw(group)
-
-        return loss

build/torch210-cxx11-cu128-x86_64-linux/optimizer/__init__.py

@@ -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")))

build/torch210-cxx11-cu130-x86_64-linux/distributed/utils.py

@@ -1,175 +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 get_slices_of_dtensor(
-    target: DTensor | torch.Tensor,
-    local_rank: int,
-    shard_mesh: DeviceMesh,
-    shard_placements: tuple[Placement],
-) -> tuple[slice]:
-    """
-    Get the slice of local tensor for a given rank from a tensor.
-    Args:
-        target (DTensor | torch.Tensor): The target tensor.
-        rank (int): The local rank of the shard group.
-        shard_mesh (DeviceMesh): The shard mesh. It consists of global ranks.
-        shard_placements (tuple[Placement]): The shard placements.
-    """
-
-    slices: list[slice] = [slice(0, dim_size) for dim_size in target.size()]
-
-    # 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)
-
-    # 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 i, (rank_coord,
-            placement) in enumerate(zip(rank_coords, shard_placements)):
-        assert isinstance(placement, Shard)
-
-        num_ranks = shard_mesh.mesh.shape[i]
-
-        dim = placement.dim
-        dim_size = (slices[dim].stop - slices[dim].start)
-
-        if dim_size % num_ranks != 0:
-            raise NotImplementedError(
-                f"Dimension size {dim_size} is not divisible "
-                f"by number of ranks {num_ranks} for shard "
-                f"placement on dim {dim}. (shape: {target.shape})")
-
-        shard_size = dim_size // num_ranks
-
-        start = slices[dim].start + rank_coord * shard_size
-        end = start + shard_size
-
-        assert start < end <= slices[dim].stop
-
-        slices[dim] = slice(start, end)
-
-    return tuple(slices)
-
-
-_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh,
-                                                  ProcessGroup]] = dict()
-
-
-def construct_shard_mesh(
-    placements: tuple[Placement],
-    mesh: DeviceMesh,
-) -> (DeviceMesh, ProcessGroup, tuple[Placement]):
-    """
-    Construct Shard Mesh and Placements for unsharding.
-    It removes Replicate placements and constructs a new Mesh and ProcessGroup.
-    """
-    my_rank = dist.get_rank()
-
-    assert mesh.mesh.device.type == 'cpu'
-
-    # Copy mesh to avoid modifying the original mesh
-    mesh = mesh.mesh.clone()
-
-    # 1. Sort placements. Replicate first, then Shard by dim ascending.
-
-    # For Shard, strided shard comes after regular shard on the same dim
-    # to preserve left-to-right order of replicate-to-shard.
-    # This is because that strided shard is using stride to represent
-    # more fine-grained sharding on the same dim.
-    # Please check the URL below for _StridedShard.
-    # https://github.com/pytorch/pytorch/blob/v2.8.0/torch/distributed/tensor/placement_types.py#L366
-
-    def placement_sort_key(
-        placement_with_index: tuple[float, Placement]
-    ) -> tuple[int, float, int]:  # (dim, split factor, original index)
-        index, placement = placement_with_index
-        is_replicate = placement.is_replicate()
-        is_shard = placement.is_shard()
-        is_partial = placement.is_partial()
-
-        assert is_replicate or is_shard, f"Unsupported placement type: {type(placement)}"
-        assert not is_partial, "Partial placement is not supported."
-
-        if is_replicate:
-            return (-1.0, 0, index)
-        elif is_shard:
-            if isinstance(placement, _StridedShard):
-                return (placement.dim, 1 / placement.split_factor, index)
-            return (placement.dim, 0, index)
-        else:
-            raise TypeError(f"Unknown placement type: {type(placement)}")
-
-    placements_with_index: list[tuple[int,
-                                      Placement]] = list(enumerate(placements))
-    placements_with_index = sorted(placements_with_index,
-                                   key=placement_sort_key)
-
-    sorted_indices, sorted_placements = zip(*placements_with_index)
-
-    # 2. Permute mesh according to sorted placements.
-    sorted_mesh = mesh.permute(sorted_indices)
-
-    # 3. Collect list of shard meshes by removing replicate dims
-    # For example, (2, 3, 4, 4) with placements [R, R, S(0), S(1)]
-    # shard_meshes should be list with 2 * 3 = 6 shard meshes of shape (4, 4)
-    num_replicates = sum(1 for p in sorted_placements if p.is_replicate())
-
-    # merge replicate dims
-    # shard_meshes became a list of shard meshes with a length of replicate degree
-    if num_replicates > 0:
-        sorted_mesh = sorted_mesh.flatten(
-            0, num_replicates - 1) if num_replicates > 1 else sorted_mesh
-        shard_meshes = list(torch.unbind(sorted_mesh, dim=0))
-    else:
-        shard_meshes = [sorted_mesh]
-    shard_placements = sorted_placements[num_replicates:]
-
-    # assume all shard placements are different
-    assert len(shard_placements) == len(set(shard_placements))
-
-    # 4. Construct ProcessGroups
-    # Caution: all groups should be created in the same order in all processes,
-    # even though each process only needs its own group.
-
-    # To use tensor as dict key, convert it to tuple
-    def tensor_to_tuple(t):
-        if isinstance(t, torch.Tensor):
-            t = t.tolist()
-        if isinstance(t, list):
-            return tuple(tensor_to_tuple(x) for x in t)
-        return t
-
-    my_shard_mesh_as_tuple = None
-    for shard_mesh in shard_meshes:
-        assert isinstance(shard_mesh, torch.Tensor)
-        shard_mesh_as_tuple = tensor_to_tuple(shard_mesh)
-
-        if (my_rank == shard_mesh).any().item():
-            assert my_shard_mesh_as_tuple is None
-            my_shard_mesh_as_tuple = shard_mesh_as_tuple
-
-        # update global cache
-        if shard_mesh_as_tuple not in _ranks_to_dist_cache:
-            shard_process_group = dist.new_group(shard_mesh.flatten().tolist())
-            _ranks_to_dist_cache[shard_mesh_as_tuple] = (
-                DeviceMesh(device_type="cuda", mesh=shard_mesh),
-                shard_process_group,
-            )
-
-    my_shard_mesh, my_shard_process_group = _ranks_to_dist_cache[
-        my_shard_mesh_as_tuple]
-
-    return my_shard_mesh, my_shard_process_group, shard_placements

build/torch210-cxx11-cu130-x86_64-linux/matmul_transpose_triton.py

@@ -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

build/torch210-cxx11-cu130-x86_64-linux/metadata.json

@@ -1 +0,0 @@
-{"python-depends":[]}

build/torch210-cxx11-cu130-x86_64-linux/muon.py

@@ -1,1268 +0,0 @@
-import logging
-import math
-import types
-from collections import defaultdict
-from dataclasses import dataclass
-from typing import Any, cast
-
-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, Replicate
-from torch.distributed.tensor.placement_types import Placement
-
-from .distributed.utils import construct_shard_mesh, get_slices_of_dtensor
-from .matmul_transpose_triton import matmul_transpose_assign
-
-logger = logging.getLogger(__name__)
-
-COMM_DTYPE = torch.bfloat16
-DEFAULT_CHUNK_SIZE_RATIO = 4
-
-
-# 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
-    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
-    assert G.dtype == COMM_DTYPE
-    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)
-    # 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)
-
-    if G.size(0) > G.size(1):
-        X = X.T
-    return X
-
-
-@dataclass
-class _muon_state:
-    # TODO: use Optional
-    worker_rank: int
-    process_group: ProcessGroup
-    shard_mesh: DeviceMesh
-    shard_placements: tuple[Placement, ...]
-    name: str
-    qk_clip_state: torch.Tensor | 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
-
-
-def numel_for_rank(
-    param: DTensor,
-    local_rank: int,
-    state: _muon_state,
-) -> int:
-    slices = get_slices_of_dtensor(
-        param,
-        local_rank,
-        state.shard_mesh,
-        state.shard_placements,
-    )
-
-    numel = 1
-    for s, dim in zip(slices, param.shape):
-        start, stop, step = s.indices(dim)
-        length = max(0, (stop - start + (step - 1)) // step)
-        numel *= length
-
-    return numel
-
-
-@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:
-                state.gathered_grad = torch.empty(p.shape,
-                                                  dtype=COMM_DTYPE,
-                                                  device="cuda")
-            else:
-                state.gathered_grad = None
-
-        alloc_event = torch.cuda.Event()
-        alloc_event.record(compute_stream)
-        return alloc_event
-
-
-@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 = numel_for_rank(p, rank, state)
-            g = p.grad
-            g = g.to_local().to(COMM_DTYPE).contiguous()
-            assert g.numel() == shard_elems
-            per_dst[dst].append(g.view(-1))
-            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 += numel_for_rank(p, src, state)
-            recv_counts[src] = total
-
-        recv_total = sum(recv_counts)
-        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
-
-        #All2All
-        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)}")
-        dist.all_to_all_single(
-            recv_buf,
-            send_buf,
-            output_split_sizes=recv_counts,
-            input_split_sizes=send_counts,
-            group=process_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
-        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
-
-                # get the slice of the full dtensor corresponding to rank src.
-                slices = get_slices_of_dtensor(state.gathered_grad, src,
-                                               state.shard_mesh,
-                                               state.shard_placements)
-
-                dst = state.gathered_grad[slices]
-                assert dst._base is state.gathered_grad
-
-                n = dst.numel()
-                assert n > 0
-
-                sg = recv_buf.narrow(0, off + inner_off, n)
-                sg = sg.reshape_as(dst)
-                dst.copy_(sg)
-
-                inner_off += n
-            off += block
-
-        for p in params:
-            state = param_to_state[id(p)]
-            if state.worker_rank == rank:
-                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
-
-
-@torch.no_grad()
-def _compute_u(p, state, steps, rank, compute_stream):
-    """
-    On worker_rank, compute the orthogonalized update using Newton-Schulz iteration.
-    """
-    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()
-        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 _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()
-
-                offset = 0
-                for dst in range(num_ranks):
-                    # get the slice of the full tensor corresponding to rank dst.
-                    slices = get_slices_of_dtensor(u_full, dst,
-                                                   state.shard_mesh,
-                                                   state.shard_placements)
-                    su = u_full[slices].flatten()
-
-                    n = su.numel()
-                    assert n > 0
-
-                    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)
-        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 += numel_for_rank(p, rank, state)
-            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,
-        )
-
-        # 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 = numel_for_rank(p, rank, state)
-                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
-            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 state.qk_clip_state is not None else None
-        if scales_full is not None:
-            # Have to slice scales_full among dim 0
-            weight_slices = get_slices_of_dtensor(p, rank, state.shard_mesh,
-                                                  state.shard_placements)
-            ratio = p.shape[0] // scales_full.shape[0]
-            scales_slice = slice(
-                None if weight_slices[0].start is None else
-                weight_slices[0].start // ratio,
-                None if weight_slices[0].stop is None else
-                weight_slices[0].stop // ratio,
-                None,
-            )
-
-            scales_local = scales_full[scales_slice]
-            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: str | None  # '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: torch.Tensor | None
-
-
-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.
-        {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.
-        small_param_numel_threshold: Threshold for classifying parameters as small and falling back to distributed Muon
-    """
-
-    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
-                 },
-                 warmup_step=5,
-                 chunk_size=-1,
-                 use_distributed_muon=False,
-                 small_param_numel_threshold=65536):
-        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.rank = None
-
-        self.comm_stream = torch.cuda.Stream()
-        self.compute_stream = torch.cuda.Stream()
-        self.debug = debug
-        self.clip_config = clip_config
-        self.warmup_step = warmup_step
-        self.chunk_size = chunk_size
-        self.use_distributed_muon = use_distributed_muon
-        self.small_param_numel_threshold = small_param_numel_threshold
-
-    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 set_rank_once(self, rank):
-        if self.rank is None:
-            self.rank = rank
-        else:
-            assert self.rank == rank
-
-    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)
-
-        # set rank with the local rank in the shard process group
-        self.set_rank_once(dist.get_rank(group=shard_pg))
-
-        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:
-            print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs",
-                  flush=True)
-
-        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 = self.get_qk_clip_info(n, qk_logits)
-
-            param_to_state[id(p)] = _muon_state(
-                worker_rank=worker_rank,
-                process_group=shard_pg,
-                shard_mesh=shard_mesh,
-                shard_placements=shard_placements,
-                name=n,
-                qk_clip_state=qk_clip_state,
-            )
-
-        return param_to_state, ordered_params
-
-    def base(self, names, params, group, lr, weight_decay, momentum,
-             qk_logits):
-        # generate weight updates in distributed fashion
-        for n, p in zip(names, 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
-
-            g = self._update_g(p, g, group, momentum)
-
-            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
-                                             steps=group["ns_steps"])
-
-            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)
-
-            scales_full = self._compute_scales(
-                p, qk_clip_state) if qk_clip_state is not None else None
-            if scales_full is not None:
-                Muon._qk_clip(p, scales_full, qk_clip_state.head_dim)
-
-    def distributed_muon(
-        self,
-        names: list[str],
-        params: list[torch.nn.Parameter],
-        group: dict[str, Any],
-        lr: float,
-        weight_decay: float,
-        momentum: float,
-        qk_logits: list[torch.Tensor | DTensor] | None,
-    ):
-        """ Implementation of Distributed Muon by Liu et al. """
-
-        for n, p in zip(names, 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
-
-            g = self._update_g(p, g, group, momentum)
-
-            # Gather G
-            if isinstance(p.data, DTensor):
-                g_full = g.full_tensor()
-                p_full = p.data.full_tensor()
-            else:
-                g_full = g
-                p_full = p
-
-            u_full = _zeropower_via_newtonschulz5(g_full.to(COMM_DTYPE),
-                                                  steps=group["ns_steps"])
-
-            adjusted_lr = self.adjust_lr_for_muon(lr, p_full.shape)
-            Muon._update_p(p_full, u_full, lr, adjusted_lr, weight_decay)
-
-            qk_clip_state = self.get_qk_clip_info(n, qk_logits)
-
-            scales_full = self._compute_scales(
-                p_full, qk_clip_state) if qk_clip_state is not None else None
-
-            if scales_full is not None:
-                Muon._qk_clip(p_full, scales_full, qk_clip_state.head_dim)
-
-            if isinstance(p.data, DTensor):
-                ndims = len(p.device_mesh.mesh.shape)
-                p_replicate = DTensor.from_local(
-                    p_full,
-                    device_mesh=p.device_mesh,
-                    placements=[Replicate() for _ in range(ndims)],
-                )
-
-                p_sharded = p_replicate.redistribute(
-                    device_mesh=p.device_mesh,
-                    placements=p.placements,
-                )
-
-                p.copy_(p_sharded)
-
-    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 group["nesterov"]:
-            g.add_(buf, alpha=momentum)
-            return g
-        return buf
-
-    @staticmethod
-    def _update_p(p, u, lr, adjusted_lr, weight_decay):
-        if isinstance(p, torch.nn.Parameter):
-            # apply weight decay
-            p.data.mul_(1 - lr * weight_decay)
-            # apply update
-            p.data.add_(u, alpha=-adjusted_lr)
-        else:
-            p.mul_(1 - lr * weight_decay)
-            p.add_(u, alpha=-adjusted_lr)
-
-    def get_qk_clip_info(self, n, qk_logits):
-        if self.clip_config is None:
-            return None
-
-        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 []
-
-            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()
-
-        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):
-        if isinstance(p, torch.nn.Parameter):
-            W = p.data.view(-1, head_dim, p.data.shape[1])
-            W.mul_(scales.view(-1, 1, 1))
-        else:
-            W = p.view(-1, head_dim, p.shape[1])
-            W.mul_(scales.view(-1, 1, 1))
-
-    def parallel(self, names, params, group, lr, weight_decay, momentum,
-                 qk_logits):
-        """
-        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(
-            names, params, group, qk_logits)
-
-        assert self.rank is not None
-
-        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_computes(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]:
-                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)
-
-        if self.chunk_size == -1:
-            shard_ranks = dist.get_world_size(param_to_state[id(
-                params[0])].process_group)
-            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.")
-
-        # Wait grad update
-        self.comm_stream.wait_stream(torch.cuda.current_stream())
-
-        warmup_step = self.warmup_step
-        for i in range(0, warmup_step):
-            enqueue_all2all_gather(i * chunk_size, chunk_size)
-            enqueue_computes(i * chunk_size, chunk_size)
-
-        for i in range(0, len(params) + chunk_size - 1, chunk_size):
-            enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + warmup_step * chunk_size, chunk_size)
-            enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + warmup_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: 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: DeviceDict | 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 _step_muon(self, group, qk_logits=None):
-        params = group["params"]
-        lr = group["lr"]
-        weight_decay = group["weight_decay"]
-        momentum = group["momentum"]
-        names = group["names"]
-
-        param_dtensors = []
-        name_dtensors = []
-
-        param_tensors = []
-        name_tensors = []
-
-        param_dtensors_small = []
-        name_dtensors_small = []
-
-        if self.use_distributed_muon:
-            self.distributed_muon(names=names,
-                                  params=params,
-                                  group=group,
-                                  lr=lr,
-                                  weight_decay=weight_decay,
-                                  momentum=momentum,
-                                  qk_logits=qk_logits)
-            return
-
-        # For simplicity, we use distributed Muon for small parameters
-        # whose number of elements is below a threshold.
-        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)
-                elif p.data.numel() <= self.small_param_numel_threshold:
-                    param_dtensors_small.append(p)
-                    name_dtensors_small.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)}")
-
-        logger.debug(
-            f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors, "
-            f"{len(param_dtensors_small)} Small DTensors")
-
-        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: ([], []))
-            # type: dict[tuple[Placement, DeviceMesh], tuple[list[str], list[DTensor]]]
-
-            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_small) > 0:
-            if not dist.is_initialized():
-                raise RuntimeError(
-                    "Parallel Muon requires torch.distributed to be initialized."
-                )
-
-            self.distributed_muon(
-                params=param_dtensors_small,
-                names=name_dtensors_small,
-                group=group,
-                lr=lr,
-                weight_decay=weight_decay,
-                momentum=momentum,
-                qk_logits=qk_logits,
-            )
-
-        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)
-            for _, (names, params) in dtensor_group.items():
-                self.parallel(
-                    names,
-                    params,
-                    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,
-            )
-
-    def _step_adamw_params(self, params, group):
-        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,
-            lr=lr,
-            weight_decay=weight_decay,
-            eps=eps,
-            maximize=False,
-        )
-
-    def _step_adamw(self, group):
-        params = group["params"]
-
-        # group params with it's type and placement
-        placement_to_params: dict[tuple[Placement | type,
-                                        DeviceMesh | None]] = defaultdict(list)
-        for p in params:
-            match p:
-                case DTensor():
-                    placement_to_params[tuple([p.placements,
-                                               p.device_mesh])].append(p)
-                case torch.Tensor():
-                    placement_to_params[tuple([torch.Tensor, None])].append(p)
-
-        for params in placement_to_params.values():
-            self._step_adamw_params(params, group)
-
-    @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()
-
-        for group in self.param_groups:
-            if group["use_muon"]:
-                self._step_muon(group, qk_logits=qk_logits)
-            else:
-                self._step_adamw(group)
-
-        return loss

build/torch210-cxx11-cu130-x86_64-linux/optimizer/__init__.py

@@ -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")))

build/torch210-cxx11-rocm70-x86_64-linux/distributed/utils.py

@@ -1,175 +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 get_slices_of_dtensor(
-    target: DTensor | torch.Tensor,
-    local_rank: int,
-    shard_mesh: DeviceMesh,
-    shard_placements: tuple[Placement],
-) -> tuple[slice]:
-    """
-    Get the slice of local tensor for a given rank from a tensor.
-    Args:
-        target (DTensor | torch.Tensor): The target tensor.
-        rank (int): The local rank of the shard group.
-        shard_mesh (DeviceMesh): The shard mesh. It consists of global ranks.
-        shard_placements (tuple[Placement]): The shard placements.
-    """
-
-    slices: list[slice] = [slice(0, dim_size) for dim_size in target.size()]
-
-    # 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)
-
-    # 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 i, (rank_coord,
-            placement) in enumerate(zip(rank_coords, shard_placements)):
-        assert isinstance(placement, Shard)
-
-        num_ranks = shard_mesh.mesh.shape[i]
-
-        dim = placement.dim
-        dim_size = (slices[dim].stop - slices[dim].start)
-
-        if dim_size % num_ranks != 0:
-            raise NotImplementedError(
-                f"Dimension size {dim_size} is not divisible "
-                f"by number of ranks {num_ranks} for shard "
-                f"placement on dim {dim}. (shape: {target.shape})")
-
-        shard_size = dim_size // num_ranks
-
-        start = slices[dim].start + rank_coord * shard_size
-        end = start + shard_size
-
-        assert start < end <= slices[dim].stop
-
-        slices[dim] = slice(start, end)
-
-    return tuple(slices)
-
-
-_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh,
-                                                  ProcessGroup]] = dict()
-
-
-def construct_shard_mesh(
-    placements: tuple[Placement],
-    mesh: DeviceMesh,
-) -> (DeviceMesh, ProcessGroup, tuple[Placement]):
-    """
-    Construct Shard Mesh and Placements for unsharding.
-    It removes Replicate placements and constructs a new Mesh and ProcessGroup.
-    """
-    my_rank = dist.get_rank()
-
-    assert mesh.mesh.device.type == 'cpu'
-
-    # Copy mesh to avoid modifying the original mesh
-    mesh = mesh.mesh.clone()
-
-    # 1. Sort placements. Replicate first, then Shard by dim ascending.
-
-    # For Shard, strided shard comes after regular shard on the same dim
-    # to preserve left-to-right order of replicate-to-shard.
-    # This is because that strided shard is using stride to represent
-    # more fine-grained sharding on the same dim.
-    # Please check the URL below for _StridedShard.
-    # https://github.com/pytorch/pytorch/blob/v2.8.0/torch/distributed/tensor/placement_types.py#L366
-
-    def placement_sort_key(
-        placement_with_index: tuple[float, Placement]
-    ) -> tuple[int, float, int]:  # (dim, split factor, original index)
-        index, placement = placement_with_index
-        is_replicate = placement.is_replicate()
-        is_shard = placement.is_shard()
-        is_partial = placement.is_partial()
-
-        assert is_replicate or is_shard, f"Unsupported placement type: {type(placement)}"
-        assert not is_partial, "Partial placement is not supported."
-
-        if is_replicate:
-            return (-1.0, 0, index)
-        elif is_shard:
-            if isinstance(placement, _StridedShard):
-                return (placement.dim, 1 / placement.split_factor, index)
-            return (placement.dim, 0, index)
-        else:
-            raise TypeError(f"Unknown placement type: {type(placement)}")
-
-    placements_with_index: list[tuple[int,
-                                      Placement]] = list(enumerate(placements))
-    placements_with_index = sorted(placements_with_index,
-                                   key=placement_sort_key)
-
-    sorted_indices, sorted_placements = zip(*placements_with_index)
-
-    # 2. Permute mesh according to sorted placements.
-    sorted_mesh = mesh.permute(sorted_indices)
-
-    # 3. Collect list of shard meshes by removing replicate dims
-    # For example, (2, 3, 4, 4) with placements [R, R, S(0), S(1)]
-    # shard_meshes should be list with 2 * 3 = 6 shard meshes of shape (4, 4)
-    num_replicates = sum(1 for p in sorted_placements if p.is_replicate())
-
-    # merge replicate dims
-    # shard_meshes became a list of shard meshes with a length of replicate degree
-    if num_replicates > 0:
-        sorted_mesh = sorted_mesh.flatten(
-            0, num_replicates - 1) if num_replicates > 1 else sorted_mesh
-        shard_meshes = list(torch.unbind(sorted_mesh, dim=0))
-    else:
-        shard_meshes = [sorted_mesh]
-    shard_placements = sorted_placements[num_replicates:]
-
-    # assume all shard placements are different
-    assert len(shard_placements) == len(set(shard_placements))
-
-    # 4. Construct ProcessGroups
-    # Caution: all groups should be created in the same order in all processes,
-    # even though each process only needs its own group.
-
-    # To use tensor as dict key, convert it to tuple
-    def tensor_to_tuple(t):
-        if isinstance(t, torch.Tensor):
-            t = t.tolist()
-        if isinstance(t, list):
-            return tuple(tensor_to_tuple(x) for x in t)
-        return t
-
-    my_shard_mesh_as_tuple = None
-    for shard_mesh in shard_meshes:
-        assert isinstance(shard_mesh, torch.Tensor)
-        shard_mesh_as_tuple = tensor_to_tuple(shard_mesh)
-
-        if (my_rank == shard_mesh).any().item():
-            assert my_shard_mesh_as_tuple is None
-            my_shard_mesh_as_tuple = shard_mesh_as_tuple
-
-        # update global cache
-        if shard_mesh_as_tuple not in _ranks_to_dist_cache:
-            shard_process_group = dist.new_group(shard_mesh.flatten().tolist())
-            _ranks_to_dist_cache[shard_mesh_as_tuple] = (
-                DeviceMesh(device_type="cuda", mesh=shard_mesh),
-                shard_process_group,
-            )
-
-    my_shard_mesh, my_shard_process_group = _ranks_to_dist_cache[
-        my_shard_mesh_as_tuple]
-
-    return my_shard_mesh, my_shard_process_group, shard_placements

build/torch210-cxx11-rocm70-x86_64-linux/matmul_transpose_triton.py

@@ -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

build/torch210-cxx11-rocm70-x86_64-linux/metadata.json

@@ -1 +0,0 @@
-{"python-depends":[]}

build/torch210-cxx11-rocm70-x86_64-linux/muon.py

@@ -1,1268 +0,0 @@
-import logging
-import math
-import types
-from collections import defaultdict
-from dataclasses import dataclass
-from typing import Any, cast
-
-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, Replicate
-from torch.distributed.tensor.placement_types import Placement
-
-from .distributed.utils import construct_shard_mesh, get_slices_of_dtensor
-from .matmul_transpose_triton import matmul_transpose_assign
-
-logger = logging.getLogger(__name__)
-
-COMM_DTYPE = torch.bfloat16
-DEFAULT_CHUNK_SIZE_RATIO = 4
-
-
-# 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
-    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
-    assert G.dtype == COMM_DTYPE
-    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)
-    # 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)
-
-    if G.size(0) > G.size(1):
-        X = X.T
-    return X
-
-
-@dataclass
-class _muon_state:
-    # TODO: use Optional
-    worker_rank: int
-    process_group: ProcessGroup
-    shard_mesh: DeviceMesh
-    shard_placements: tuple[Placement, ...]
-    name: str
-    qk_clip_state: torch.Tensor | 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
-
-
-def numel_for_rank(
-    param: DTensor,
-    local_rank: int,
-    state: _muon_state,
-) -> int:
-    slices = get_slices_of_dtensor(
-        param,
-        local_rank,
-        state.shard_mesh,
-        state.shard_placements,
-    )
-
-    numel = 1
-    for s, dim in zip(slices, param.shape):
-        start, stop, step = s.indices(dim)
-        length = max(0, (stop - start + (step - 1)) // step)
-        numel *= length
-
-    return numel
-
-
-@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:
-                state.gathered_grad = torch.empty(p.shape,
-                                                  dtype=COMM_DTYPE,
-                                                  device="cuda")
-            else:
-                state.gathered_grad = None
-
-        alloc_event = torch.cuda.Event()
-        alloc_event.record(compute_stream)
-        return alloc_event
-
-
-@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 = numel_for_rank(p, rank, state)
-            g = p.grad
-            g = g.to_local().to(COMM_DTYPE).contiguous()
-            assert g.numel() == shard_elems
-            per_dst[dst].append(g.view(-1))
-            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 += numel_for_rank(p, src, state)
-            recv_counts[src] = total
-
-        recv_total = sum(recv_counts)
-        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
-
-        #All2All
-        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)}")
-        dist.all_to_all_single(
-            recv_buf,
-            send_buf,
-            output_split_sizes=recv_counts,
-            input_split_sizes=send_counts,
-            group=process_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
-        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
-
-                # get the slice of the full dtensor corresponding to rank src.
-                slices = get_slices_of_dtensor(state.gathered_grad, src,
-                                               state.shard_mesh,
-                                               state.shard_placements)
-
-                dst = state.gathered_grad[slices]
-                assert dst._base is state.gathered_grad
-
-                n = dst.numel()
-                assert n > 0
-
-                sg = recv_buf.narrow(0, off + inner_off, n)
-                sg = sg.reshape_as(dst)
-                dst.copy_(sg)
-
-                inner_off += n
-            off += block
-
-        for p in params:
-            state = param_to_state[id(p)]
-            if state.worker_rank == rank:
-                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
-
-
-@torch.no_grad()
-def _compute_u(p, state, steps, rank, compute_stream):
-    """
-    On worker_rank, compute the orthogonalized update using Newton-Schulz iteration.
-    """
-    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()
-        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 _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()
-
-                offset = 0
-                for dst in range(num_ranks):
-                    # get the slice of the full tensor corresponding to rank dst.
-                    slices = get_slices_of_dtensor(u_full, dst,
-                                                   state.shard_mesh,
-                                                   state.shard_placements)
-                    su = u_full[slices].flatten()
-
-                    n = su.numel()
-                    assert n > 0
-
-                    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)
-        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 += numel_for_rank(p, rank, state)
-            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,
-        )
-
-        # 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 = numel_for_rank(p, rank, state)
-                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
-            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 state.qk_clip_state is not None else None
-        if scales_full is not None:
-            # Have to slice scales_full among dim 0
-            weight_slices = get_slices_of_dtensor(p, rank, state.shard_mesh,
-                                                  state.shard_placements)
-            ratio = p.shape[0] // scales_full.shape[0]
-            scales_slice = slice(
-                None if weight_slices[0].start is None else
-                weight_slices[0].start // ratio,
-                None if weight_slices[0].stop is None else
-                weight_slices[0].stop // ratio,
-                None,
-            )
-
-            scales_local = scales_full[scales_slice]
-            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: str | None  # '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: torch.Tensor | None
-
-
-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.
-        {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.
-        small_param_numel_threshold: Threshold for classifying parameters as small and falling back to distributed Muon
-    """
-
-    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
-                 },
-                 warmup_step=5,
-                 chunk_size=-1,
-                 use_distributed_muon=False,
-                 small_param_numel_threshold=65536):
-        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.rank = None
-
-        self.comm_stream = torch.cuda.Stream()
-        self.compute_stream = torch.cuda.Stream()
-        self.debug = debug
-        self.clip_config = clip_config
-        self.warmup_step = warmup_step
-        self.chunk_size = chunk_size
-        self.use_distributed_muon = use_distributed_muon
-        self.small_param_numel_threshold = small_param_numel_threshold
-
-    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 set_rank_once(self, rank):
-        if self.rank is None:
-            self.rank = rank
-        else:
-            assert self.rank == rank
-
-    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)
-
-        # set rank with the local rank in the shard process group
-        self.set_rank_once(dist.get_rank(group=shard_pg))
-
-        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:
-            print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs",
-                  flush=True)
-
-        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 = self.get_qk_clip_info(n, qk_logits)
-
-            param_to_state[id(p)] = _muon_state(
-                worker_rank=worker_rank,
-                process_group=shard_pg,
-                shard_mesh=shard_mesh,
-                shard_placements=shard_placements,
-                name=n,
-                qk_clip_state=qk_clip_state,
-            )
-
-        return param_to_state, ordered_params
-
-    def base(self, names, params, group, lr, weight_decay, momentum,
-             qk_logits):
-        # generate weight updates in distributed fashion
-        for n, p in zip(names, 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
-
-            g = self._update_g(p, g, group, momentum)
-
-            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
-                                             steps=group["ns_steps"])
-
-            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)
-
-            scales_full = self._compute_scales(
-                p, qk_clip_state) if qk_clip_state is not None else None
-            if scales_full is not None:
-                Muon._qk_clip(p, scales_full, qk_clip_state.head_dim)
-
-    def distributed_muon(
-        self,
-        names: list[str],
-        params: list[torch.nn.Parameter],
-        group: dict[str, Any],
-        lr: float,
-        weight_decay: float,
-        momentum: float,
-        qk_logits: list[torch.Tensor | DTensor] | None,
-    ):
-        """ Implementation of Distributed Muon by Liu et al. """
-
-        for n, p in zip(names, 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
-
-            g = self._update_g(p, g, group, momentum)
-
-            # Gather G
-            if isinstance(p.data, DTensor):
-                g_full = g.full_tensor()
-                p_full = p.data.full_tensor()
-            else:
-                g_full = g
-                p_full = p
-
-            u_full = _zeropower_via_newtonschulz5(g_full.to(COMM_DTYPE),
-                                                  steps=group["ns_steps"])
-
-            adjusted_lr = self.adjust_lr_for_muon(lr, p_full.shape)
-            Muon._update_p(p_full, u_full, lr, adjusted_lr, weight_decay)
-
-            qk_clip_state = self.get_qk_clip_info(n, qk_logits)
-
-            scales_full = self._compute_scales(
-                p_full, qk_clip_state) if qk_clip_state is not None else None
-
-            if scales_full is not None:
-                Muon._qk_clip(p_full, scales_full, qk_clip_state.head_dim)
-
-            if isinstance(p.data, DTensor):
-                ndims = len(p.device_mesh.mesh.shape)
-                p_replicate = DTensor.from_local(
-                    p_full,
-                    device_mesh=p.device_mesh,
-                    placements=[Replicate() for _ in range(ndims)],
-                )
-
-                p_sharded = p_replicate.redistribute(
-                    device_mesh=p.device_mesh,
-                    placements=p.placements,
-                )
-
-                p.copy_(p_sharded)
-
-    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 group["nesterov"]:
-            g.add_(buf, alpha=momentum)
-            return g
-        return buf
-
-    @staticmethod
-    def _update_p(p, u, lr, adjusted_lr, weight_decay):
-        if isinstance(p, torch.nn.Parameter):
-            # apply weight decay
-            p.data.mul_(1 - lr * weight_decay)
-            # apply update
-            p.data.add_(u, alpha=-adjusted_lr)
-        else:
-            p.mul_(1 - lr * weight_decay)
-            p.add_(u, alpha=-adjusted_lr)
-
-    def get_qk_clip_info(self, n, qk_logits):
-        if self.clip_config is None:
-            return None
-
-        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 []
-
-            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()
-
-        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):
-        if isinstance(p, torch.nn.Parameter):
-            W = p.data.view(-1, head_dim, p.data.shape[1])
-            W.mul_(scales.view(-1, 1, 1))
-        else:
-            W = p.view(-1, head_dim, p.shape[1])
-            W.mul_(scales.view(-1, 1, 1))
-
-    def parallel(self, names, params, group, lr, weight_decay, momentum,
-                 qk_logits):
-        """
-        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(
-            names, params, group, qk_logits)
-
-        assert self.rank is not None
-
-        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_computes(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]:
-                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)
-
-        if self.chunk_size == -1:
-            shard_ranks = dist.get_world_size(param_to_state[id(
-                params[0])].process_group)
-            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.")
-
-        # Wait grad update
-        self.comm_stream.wait_stream(torch.cuda.current_stream())
-
-        warmup_step = self.warmup_step
-        for i in range(0, warmup_step):
-            enqueue_all2all_gather(i * chunk_size, chunk_size)
-            enqueue_computes(i * chunk_size, chunk_size)
-
-        for i in range(0, len(params) + chunk_size - 1, chunk_size):
-            enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + warmup_step * chunk_size, chunk_size)
-            enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + warmup_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: 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: DeviceDict | 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 _step_muon(self, group, qk_logits=None):
-        params = group["params"]
-        lr = group["lr"]
-        weight_decay = group["weight_decay"]
-        momentum = group["momentum"]
-        names = group["names"]
-
-        param_dtensors = []
-        name_dtensors = []
-
-        param_tensors = []
-        name_tensors = []
-
-        param_dtensors_small = []
-        name_dtensors_small = []
-
-        if self.use_distributed_muon:
-            self.distributed_muon(names=names,
-                                  params=params,
-                                  group=group,
-                                  lr=lr,
-                                  weight_decay=weight_decay,
-                                  momentum=momentum,
-                                  qk_logits=qk_logits)
-            return
-
-        # For simplicity, we use distributed Muon for small parameters
-        # whose number of elements is below a threshold.
-        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)
-                elif p.data.numel() <= self.small_param_numel_threshold:
-                    param_dtensors_small.append(p)
-                    name_dtensors_small.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)}")
-
-        logger.debug(
-            f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors, "
-            f"{len(param_dtensors_small)} Small DTensors")
-
-        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: ([], []))
-            # type: dict[tuple[Placement, DeviceMesh], tuple[list[str], list[DTensor]]]
-
-            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_small) > 0:
-            if not dist.is_initialized():
-                raise RuntimeError(
-                    "Parallel Muon requires torch.distributed to be initialized."
-                )
-
-            self.distributed_muon(
-                params=param_dtensors_small,
-                names=name_dtensors_small,
-                group=group,
-                lr=lr,
-                weight_decay=weight_decay,
-                momentum=momentum,
-                qk_logits=qk_logits,
-            )
-
-        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)
-            for _, (names, params) in dtensor_group.items():
-                self.parallel(
-                    names,
-                    params,
-                    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,
-            )
-
-    def _step_adamw_params(self, params, group):
-        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,
-            lr=lr,
-            weight_decay=weight_decay,
-            eps=eps,
-            maximize=False,
-        )
-
-    def _step_adamw(self, group):
-        params = group["params"]
-
-        # group params with it's type and placement
-        placement_to_params: dict[tuple[Placement | type,
-                                        DeviceMesh | None]] = defaultdict(list)
-        for p in params:
-            match p:
-                case DTensor():
-                    placement_to_params[tuple([p.placements,
-                                               p.device_mesh])].append(p)
-                case torch.Tensor():
-                    placement_to_params[tuple([torch.Tensor, None])].append(p)
-
-        for params in placement_to_params.values():
-            self._step_adamw_params(params, group)
-
-    @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()
-
-        for group in self.param_groups:
-            if group["use_muon"]:
-                self._step_muon(group, qk_logits=qk_logits)
-            else:
-                self._step_adamw(group)
-
-        return loss

build/torch210-cxx11-rocm70-x86_64-linux/optimizer/__init__.py

@@ -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")))

build/torch210-cxx11-rocm71-x86_64-linux/_ops.py

@@ -1,9 +0,0 @@
-import torch
-from . import _optimizer_06a260a_dirty
-ops = torch.ops._optimizer_06a260a_dirty
-
-def add_op_namespace_prefix(op_name: str):
-    """
-    Prefix op by namespace.
-    """
-    return f"_optimizer_06a260a_dirty::{op_name}"

build/torch210-cxx11-rocm71-x86_64-linux/_optimizer_06a260a_dirty.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:d804ba4d3ed9716c80e9819ba16a2bef300fb23fa4c456c550f4a96167a2eb00
-size 1866112

build/torch210-cxx11-rocm71-x86_64-linux/distributed/utils.py

@@ -1,175 +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 get_slices_of_dtensor(
-    target: DTensor | torch.Tensor,
-    local_rank: int,
-    shard_mesh: DeviceMesh,
-    shard_placements: tuple[Placement],
-) -> tuple[slice]:
-    """
-    Get the slice of local tensor for a given rank from a tensor.
-    Args:
-        target (DTensor | torch.Tensor): The target tensor.
-        rank (int): The local rank of the shard group.
-        shard_mesh (DeviceMesh): The shard mesh. It consists of global ranks.
-        shard_placements (tuple[Placement]): The shard placements.
-    """
-
-    slices: list[slice] = [slice(0, dim_size) for dim_size in target.size()]
-
-    # 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)
-
-    # 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 i, (rank_coord,
-            placement) in enumerate(zip(rank_coords, shard_placements)):
-        assert isinstance(placement, Shard)
-
-        num_ranks = shard_mesh.mesh.shape[i]
-
-        dim = placement.dim
-        dim_size = (slices[dim].stop - slices[dim].start)
-
-        if dim_size % num_ranks != 0:
-            raise NotImplementedError(
-                f"Dimension size {dim_size} is not divisible "
-                f"by number of ranks {num_ranks} for shard "
-                f"placement on dim {dim}. (shape: {target.shape})")
-
-        shard_size = dim_size // num_ranks
-
-        start = slices[dim].start + rank_coord * shard_size
-        end = start + shard_size
-
-        assert start < end <= slices[dim].stop
-
-        slices[dim] = slice(start, end)
-
-    return tuple(slices)
-
-
-_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh,
-                                                  ProcessGroup]] = dict()
-
-
-def construct_shard_mesh(
-    placements: tuple[Placement],
-    mesh: DeviceMesh,
-) -> (DeviceMesh, ProcessGroup, tuple[Placement]):
-    """
-    Construct Shard Mesh and Placements for unsharding.
-    It removes Replicate placements and constructs a new Mesh and ProcessGroup.
-    """
-    my_rank = dist.get_rank()
-
-    assert mesh.mesh.device.type == 'cpu'
-
-    # Copy mesh to avoid modifying the original mesh
-    mesh = mesh.mesh.clone()
-
-    # 1. Sort placements. Replicate first, then Shard by dim ascending.
-
-    # For Shard, strided shard comes after regular shard on the same dim
-    # to preserve left-to-right order of replicate-to-shard.
-    # This is because that strided shard is using stride to represent
-    # more fine-grained sharding on the same dim.
-    # Please check the URL below for _StridedShard.
-    # https://github.com/pytorch/pytorch/blob/v2.8.0/torch/distributed/tensor/placement_types.py#L366
-
-    def placement_sort_key(
-        placement_with_index: tuple[float, Placement]
-    ) -> tuple[int, float, int]:  # (dim, split factor, original index)
-        index, placement = placement_with_index
-        is_replicate = placement.is_replicate()
-        is_shard = placement.is_shard()
-        is_partial = placement.is_partial()
-
-        assert is_replicate or is_shard, f"Unsupported placement type: {type(placement)}"
-        assert not is_partial, "Partial placement is not supported."
-
-        if is_replicate:
-            return (-1.0, 0, index)
-        elif is_shard:
-            if isinstance(placement, _StridedShard):
-                return (placement.dim, 1 / placement.split_factor, index)
-            return (placement.dim, 0, index)
-        else:
-            raise TypeError(f"Unknown placement type: {type(placement)}")
-
-    placements_with_index: list[tuple[int,
-                                      Placement]] = list(enumerate(placements))
-    placements_with_index = sorted(placements_with_index,
-                                   key=placement_sort_key)
-
-    sorted_indices, sorted_placements = zip(*placements_with_index)
-
-    # 2. Permute mesh according to sorted placements.
-    sorted_mesh = mesh.permute(sorted_indices)
-
-    # 3. Collect list of shard meshes by removing replicate dims
-    # For example, (2, 3, 4, 4) with placements [R, R, S(0), S(1)]
-    # shard_meshes should be list with 2 * 3 = 6 shard meshes of shape (4, 4)
-    num_replicates = sum(1 for p in sorted_placements if p.is_replicate())
-
-    # merge replicate dims
-    # shard_meshes became a list of shard meshes with a length of replicate degree
-    if num_replicates > 0:
-        sorted_mesh = sorted_mesh.flatten(
-            0, num_replicates - 1) if num_replicates > 1 else sorted_mesh
-        shard_meshes = list(torch.unbind(sorted_mesh, dim=0))
-    else:
-        shard_meshes = [sorted_mesh]
-    shard_placements = sorted_placements[num_replicates:]
-
-    # assume all shard placements are different
-    assert len(shard_placements) == len(set(shard_placements))
-
-    # 4. Construct ProcessGroups
-    # Caution: all groups should be created in the same order in all processes,
-    # even though each process only needs its own group.
-
-    # To use tensor as dict key, convert it to tuple
-    def tensor_to_tuple(t):
-        if isinstance(t, torch.Tensor):
-            t = t.tolist()
-        if isinstance(t, list):
-            return tuple(tensor_to_tuple(x) for x in t)
-        return t
-
-    my_shard_mesh_as_tuple = None
-    for shard_mesh in shard_meshes:
-        assert isinstance(shard_mesh, torch.Tensor)
-        shard_mesh_as_tuple = tensor_to_tuple(shard_mesh)
-
-        if (my_rank == shard_mesh).any().item():
-            assert my_shard_mesh_as_tuple is None
-            my_shard_mesh_as_tuple = shard_mesh_as_tuple
-
-        # update global cache
-        if shard_mesh_as_tuple not in _ranks_to_dist_cache:
-            shard_process_group = dist.new_group(shard_mesh.flatten().tolist())
-            _ranks_to_dist_cache[shard_mesh_as_tuple] = (
-                DeviceMesh(device_type="cuda", mesh=shard_mesh),
-                shard_process_group,
-            )
-
-    my_shard_mesh, my_shard_process_group = _ranks_to_dist_cache[
-        my_shard_mesh_as_tuple]
-
-    return my_shard_mesh, my_shard_process_group, shard_placements

build/torch210-cxx11-rocm71-x86_64-linux/matmul_transpose_triton.py

@@ -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

build/torch210-cxx11-rocm71-x86_64-linux/metadata.json

@@ -1 +0,0 @@
-{"python-depends":[]}

build/torch210-cxx11-rocm71-x86_64-linux/muon.py

@@ -1,1268 +0,0 @@
-import logging
-import math
-import types
-from collections import defaultdict
-from dataclasses import dataclass
-from typing import Any, cast
-
-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, Replicate
-from torch.distributed.tensor.placement_types import Placement
-
-from .distributed.utils import construct_shard_mesh, get_slices_of_dtensor
-from .matmul_transpose_triton import matmul_transpose_assign
-
-logger = logging.getLogger(__name__)
-
-COMM_DTYPE = torch.bfloat16
-DEFAULT_CHUNK_SIZE_RATIO = 4
-
-
-# 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
-    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
-    assert G.dtype == COMM_DTYPE
-    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)
-    # 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)
-
-    if G.size(0) > G.size(1):
-        X = X.T
-    return X
-
-
-@dataclass
-class _muon_state:
-    # TODO: use Optional
-    worker_rank: int
-    process_group: ProcessGroup
-    shard_mesh: DeviceMesh
-    shard_placements: tuple[Placement, ...]
-    name: str
-    qk_clip_state: torch.Tensor | 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
-
-
-def numel_for_rank(
-    param: DTensor,
-    local_rank: int,
-    state: _muon_state,
-) -> int:
-    slices = get_slices_of_dtensor(
-        param,
-        local_rank,
-        state.shard_mesh,
-        state.shard_placements,
-    )
-
-    numel = 1
-    for s, dim in zip(slices, param.shape):
-        start, stop, step = s.indices(dim)
-        length = max(0, (stop - start + (step - 1)) // step)
-        numel *= length
-
-    return numel
-
-
-@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:
-                state.gathered_grad = torch.empty(p.shape,
-                                                  dtype=COMM_DTYPE,
-                                                  device="cuda")
-            else:
-                state.gathered_grad = None
-
-        alloc_event = torch.cuda.Event()
-        alloc_event.record(compute_stream)
-        return alloc_event
-
-
-@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 = numel_for_rank(p, rank, state)
-            g = p.grad
-            g = g.to_local().to(COMM_DTYPE).contiguous()
-            assert g.numel() == shard_elems
-            per_dst[dst].append(g.view(-1))
-            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 += numel_for_rank(p, src, state)
-            recv_counts[src] = total
-
-        recv_total = sum(recv_counts)
-        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
-
-        #All2All
-        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)}")
-        dist.all_to_all_single(
-            recv_buf,
-            send_buf,
-            output_split_sizes=recv_counts,
-            input_split_sizes=send_counts,
-            group=process_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
-        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
-
-                # get the slice of the full dtensor corresponding to rank src.
-                slices = get_slices_of_dtensor(state.gathered_grad, src,
-                                               state.shard_mesh,
-                                               state.shard_placements)
-
-                dst = state.gathered_grad[slices]
-                assert dst._base is state.gathered_grad
-
-                n = dst.numel()
-                assert n > 0
-
-                sg = recv_buf.narrow(0, off + inner_off, n)
-                sg = sg.reshape_as(dst)
-                dst.copy_(sg)
-
-                inner_off += n
-            off += block
-
-        for p in params:
-            state = param_to_state[id(p)]
-            if state.worker_rank == rank:
-                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
-
-
-@torch.no_grad()
-def _compute_u(p, state, steps, rank, compute_stream):
-    """
-    On worker_rank, compute the orthogonalized update using Newton-Schulz iteration.
-    """
-    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()
-        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 _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()
-
-                offset = 0
-                for dst in range(num_ranks):
-                    # get the slice of the full tensor corresponding to rank dst.
-                    slices = get_slices_of_dtensor(u_full, dst,
-                                                   state.shard_mesh,
-                                                   state.shard_placements)
-                    su = u_full[slices].flatten()
-
-                    n = su.numel()
-                    assert n > 0
-
-                    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)
-        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 += numel_for_rank(p, rank, state)
-            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,
-        )
-
-        # 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 = numel_for_rank(p, rank, state)
-                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
-            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 state.qk_clip_state is not None else None
-        if scales_full is not None:
-            # Have to slice scales_full among dim 0
-            weight_slices = get_slices_of_dtensor(p, rank, state.shard_mesh,
-                                                  state.shard_placements)
-            ratio = p.shape[0] // scales_full.shape[0]
-            scales_slice = slice(
-                None if weight_slices[0].start is None else
-                weight_slices[0].start // ratio,
-                None if weight_slices[0].stop is None else
-                weight_slices[0].stop // ratio,
-                None,
-            )
-
-            scales_local = scales_full[scales_slice]
-            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: str | None  # '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: torch.Tensor | None
-
-
-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.
-        {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.
-        small_param_numel_threshold: Threshold for classifying parameters as small and falling back to distributed Muon
-    """
-
-    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
-                 },
-                 warmup_step=5,
-                 chunk_size=-1,
-                 use_distributed_muon=False,
-                 small_param_numel_threshold=65536):
-        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.rank = None
-
-        self.comm_stream = torch.cuda.Stream()
-        self.compute_stream = torch.cuda.Stream()
-        self.debug = debug
-        self.clip_config = clip_config
-        self.warmup_step = warmup_step
-        self.chunk_size = chunk_size
-        self.use_distributed_muon = use_distributed_muon
-        self.small_param_numel_threshold = small_param_numel_threshold
-
-    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 set_rank_once(self, rank):
-        if self.rank is None:
-            self.rank = rank
-        else:
-            assert self.rank == rank
-
-    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)
-
-        # set rank with the local rank in the shard process group
-        self.set_rank_once(dist.get_rank(group=shard_pg))
-
-        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:
-            print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs",
-                  flush=True)
-
-        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 = self.get_qk_clip_info(n, qk_logits)
-
-            param_to_state[id(p)] = _muon_state(
-                worker_rank=worker_rank,
-                process_group=shard_pg,
-                shard_mesh=shard_mesh,
-                shard_placements=shard_placements,
-                name=n,
-                qk_clip_state=qk_clip_state,
-            )
-
-        return param_to_state, ordered_params
-
-    def base(self, names, params, group, lr, weight_decay, momentum,
-             qk_logits):
-        # generate weight updates in distributed fashion
-        for n, p in zip(names, 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
-
-            g = self._update_g(p, g, group, momentum)
-
-            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
-                                             steps=group["ns_steps"])
-
-            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)
-
-            scales_full = self._compute_scales(
-                p, qk_clip_state) if qk_clip_state is not None else None
-            if scales_full is not None:
-                Muon._qk_clip(p, scales_full, qk_clip_state.head_dim)
-
-    def distributed_muon(
-        self,
-        names: list[str],
-        params: list[torch.nn.Parameter],
-        group: dict[str, Any],
-        lr: float,
-        weight_decay: float,
-        momentum: float,
-        qk_logits: list[torch.Tensor | DTensor] | None,
-    ):
-        """ Implementation of Distributed Muon by Liu et al. """
-
-        for n, p in zip(names, 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
-
-            g = self._update_g(p, g, group, momentum)
-
-            # Gather G
-            if isinstance(p.data, DTensor):
-                g_full = g.full_tensor()
-                p_full = p.data.full_tensor()
-            else:
-                g_full = g
-                p_full = p
-
-            u_full = _zeropower_via_newtonschulz5(g_full.to(COMM_DTYPE),
-                                                  steps=group["ns_steps"])
-
-            adjusted_lr = self.adjust_lr_for_muon(lr, p_full.shape)
-            Muon._update_p(p_full, u_full, lr, adjusted_lr, weight_decay)
-
-            qk_clip_state = self.get_qk_clip_info(n, qk_logits)
-
-            scales_full = self._compute_scales(
-                p_full, qk_clip_state) if qk_clip_state is not None else None
-
-            if scales_full is not None:
-                Muon._qk_clip(p_full, scales_full, qk_clip_state.head_dim)
-
-            if isinstance(p.data, DTensor):
-                ndims = len(p.device_mesh.mesh.shape)
-                p_replicate = DTensor.from_local(
-                    p_full,
-                    device_mesh=p.device_mesh,
-                    placements=[Replicate() for _ in range(ndims)],
-                )
-
-                p_sharded = p_replicate.redistribute(
-                    device_mesh=p.device_mesh,
-                    placements=p.placements,
-                )
-
-                p.copy_(p_sharded)
-
-    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 group["nesterov"]:
-            g.add_(buf, alpha=momentum)
-            return g
-        return buf
-
-    @staticmethod
-    def _update_p(p, u, lr, adjusted_lr, weight_decay):
-        if isinstance(p, torch.nn.Parameter):
-            # apply weight decay
-            p.data.mul_(1 - lr * weight_decay)
-            # apply update
-            p.data.add_(u, alpha=-adjusted_lr)
-        else:
-            p.mul_(1 - lr * weight_decay)
-            p.add_(u, alpha=-adjusted_lr)
-
-    def get_qk_clip_info(self, n, qk_logits):
-        if self.clip_config is None:
-            return None
-
-        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 []
-
-            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()
-
-        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):
-        if isinstance(p, torch.nn.Parameter):
-            W = p.data.view(-1, head_dim, p.data.shape[1])
-            W.mul_(scales.view(-1, 1, 1))
-        else:
-            W = p.view(-1, head_dim, p.shape[1])
-            W.mul_(scales.view(-1, 1, 1))
-
-    def parallel(self, names, params, group, lr, weight_decay, momentum,
-                 qk_logits):
-        """
-        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(
-            names, params, group, qk_logits)
-
-        assert self.rank is not None
-
-        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_computes(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]:
-                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)
-
-        if self.chunk_size == -1:
-            shard_ranks = dist.get_world_size(param_to_state[id(
-                params[0])].process_group)
-            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.")
-
-        # Wait grad update
-        self.comm_stream.wait_stream(torch.cuda.current_stream())
-
-        warmup_step = self.warmup_step
-        for i in range(0, warmup_step):
-            enqueue_all2all_gather(i * chunk_size, chunk_size)
-            enqueue_computes(i * chunk_size, chunk_size)
-
-        for i in range(0, len(params) + chunk_size - 1, chunk_size):
-            enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + warmup_step * chunk_size, chunk_size)
-            enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + warmup_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: 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: DeviceDict | 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 _step_muon(self, group, qk_logits=None):
-        params = group["params"]
-        lr = group["lr"]
-        weight_decay = group["weight_decay"]
-        momentum = group["momentum"]
-        names = group["names"]
-
-        param_dtensors = []
-        name_dtensors = []
-
-        param_tensors = []
-        name_tensors = []
-
-        param_dtensors_small = []
-        name_dtensors_small = []
-
-        if self.use_distributed_muon:
-            self.distributed_muon(names=names,
-                                  params=params,
-                                  group=group,
-                                  lr=lr,
-                                  weight_decay=weight_decay,
-                                  momentum=momentum,
-                                  qk_logits=qk_logits)
-            return
-
-        # For simplicity, we use distributed Muon for small parameters
-        # whose number of elements is below a threshold.
-        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)
-                elif p.data.numel() <= self.small_param_numel_threshold:
-                    param_dtensors_small.append(p)
-                    name_dtensors_small.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)}")
-
-        logger.debug(
-            f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors, "
-            f"{len(param_dtensors_small)} Small DTensors")
-
-        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: ([], []))
-            # type: dict[tuple[Placement, DeviceMesh], tuple[list[str], list[DTensor]]]
-
-            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_small) > 0:
-            if not dist.is_initialized():
-                raise RuntimeError(
-                    "Parallel Muon requires torch.distributed to be initialized."
-                )
-
-            self.distributed_muon(
-                params=param_dtensors_small,
-                names=name_dtensors_small,
-                group=group,
-                lr=lr,
-                weight_decay=weight_decay,
-                momentum=momentum,
-                qk_logits=qk_logits,
-            )
-
-        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)
-            for _, (names, params) in dtensor_group.items():
-                self.parallel(
-                    names,
-                    params,
-                    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,
-            )
-
-    def _step_adamw_params(self, params, group):
-        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,
-            lr=lr,
-            weight_decay=weight_decay,
-            eps=eps,
-            maximize=False,
-        )
-
-    def _step_adamw(self, group):
-        params = group["params"]
-
-        # group params with it's type and placement
-        placement_to_params: dict[tuple[Placement | type,
-                                        DeviceMesh | None]] = defaultdict(list)
-        for p in params:
-            match p:
-                case DTensor():
-                    placement_to_params[tuple([p.placements,
-                                               p.device_mesh])].append(p)
-                case torch.Tensor():
-                    placement_to_params[tuple([torch.Tensor, None])].append(p)
-
-        for params in placement_to_params.values():
-            self._step_adamw_params(params, group)
-
-    @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()
-
-        for group in self.param_groups:
-            if group["use_muon"]:
-                self._step_muon(group, qk_logits=qk_logits)
-            else:
-                self._step_adamw(group)
-
-        return loss

build/torch210-cxx11-rocm71-x86_64-linux/optimizer/__init__.py

@@ -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")))

build/{torch210-cxx11-rocm70-x86_64-linux → torch26-cxx11-cu118-x86_64-linux/optimizer}/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_06a260a_dirty
-ops = torch.ops._optimizer_06a260a_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_06a260a_dirty::{op_name}"
+    return f"_optimizer_036642a_dirty::{op_name}"

build/{torch210-cxx11-cu126-x86_64-linux/_optimizer_06a260a_dirty.abi3.so → torch26-cxx11-cu118-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5384da54f22f488e0646e09915b821b3235cb404b163a570aa377967f853e3cf
-size 1940944
+oid sha256:9c77e5647b6056bfaee25050cca7948c40859db0a88fa4fcf40b67a85c947d8c
+size 1787272

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

build/{torch210-cxx11-cu126-x86_64-linux → torch26-cxx11-cu124-x86_64-linux/optimizer}/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_06a260a_dirty
-ops = torch.ops._optimizer_06a260a_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_06a260a_dirty::{op_name}"
+    return f"_optimizer_036642a_dirty::{op_name}"

build/{torch210-cxx11-cu128-x86_64-linux/_optimizer_06a260a_dirty.abi3.so → torch26-cxx11-cu124-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:976df6a1ec3ec4c462dea18477b56dfb75bcff76f504d55b592ce417931597c0
-size 2004144
+oid sha256:94ea66089cc8d9eda72b017733a9e05e4fee5a2f04c50658b690d2c19f0d3068
+size 1824224

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

build/{torch210-cxx11-cu130-x86_64-linux → torch26-cxx11-cu126-x86_64-linux/optimizer}/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_06a260a_dirty
-ops = torch.ops._optimizer_06a260a_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_06a260a_dirty::{op_name}"
+    return f"_optimizer_036642a_dirty::{op_name}"

build/{torch210-cxx11-cu130-x86_64-linux/_optimizer_06a260a_dirty.abi3.so → torch26-cxx11-cu126-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:330aaa6cb247ba3b5df7a13ced6ef7eff3e5d7a72a0b88f674f948aeaed66ee2
-size 2004728
+oid sha256:46e01e1d957ada2d485b30cd60bc3ef7230b8857dffc59f2e7924339761ec577
+size 1824224

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

build/{torch210-cxx11-cu128-x86_64-linux → torch26-cxx11-rocm62-x86_64-linux/optimizer}/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_06a260a_dirty
-ops = torch.ops._optimizer_06a260a_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_06a260a_dirty::{op_name}"
+    return f"_optimizer_036642a_dirty::{op_name}"

build/{torch210-cxx11-rocm70-x86_64-linux/_optimizer_06a260a_dirty.abi3.so → torch26-cxx11-rocm62-x86_64-linux/optimizer/_optimizer_036642a_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3562c68e8ee85fc5b268e079150ffff69d52860092d59e44fb9b3c4526c5d497
-size 1866400
+oid sha256:a825a0cd31d8c1b91aa9db4b24248d7fc0a506615f625a385b40e6002025c7dd
+size 1749744