diff --git a/README.md b/README.md index 4bfd839..a08c702 100644 --- a/README.md +++ b/README.md @@ -15,8 +15,25 @@

+ Docs + PyTorch + HSDP + Tensor Parallel + License Tech Report - Wiki + status +

+ +

+ 📖 Documentation +  ·  + 🚀 Quick Start +  ·  + 🌐 HSDP Guide +  ·  + 🏛️ Architecture +  ·  + 🇨🇳 中文文档

**DMuon** is a high-performance distributed implementation of the Muon optimizer that drops diff --git a/dmuon/__init__.py b/dmuon/__init__.py index 898dae5..de69e96 100644 --- a/dmuon/__init__.py +++ b/dmuon/__init__.py @@ -41,6 +41,13 @@ clip_grad_norm_, register_muon_grad_clip_strategy, ) +from .fast_clip import ( + GradClipBucket, + GradClipBucketStats, + SegmentedGradClipResult, + clip_grad_norm_buckets_, + try_clip_optimizer_grad_norm_buckets_, +) from .diagnostics import summarize_comm_plan, summarize_param_groups from .policy import DMuonParamPolicy from .optim import ( @@ -89,8 +96,13 @@ "DedicatedCommContext", "DMuonParamPolicy", "clip_grad_norm_", + "clip_grad_norm_buckets_", + "try_clip_optimizer_grad_norm_buckets_", "register_muon_grad_clip_strategy", "MuonGradClipStats", + "GradClipBucket", + "GradClipBucketStats", + "SegmentedGradClipResult", "summarize_comm_plan", "summarize_param_groups", "no_sync", diff --git a/dmuon/_backends/fsdp2/param.py b/dmuon/_backends/fsdp2/param.py index 0cd4d22..3a34715 100644 --- a/dmuon/_backends/fsdp2/param.py +++ b/dmuon/_backends/fsdp2/param.py @@ -240,7 +240,7 @@ def __init__( self.is_tp_owner = False # Storage: every rank in the owner's shard column keeps a populated - # ``_owned_data``. + # ``_owned_data`` for owner-broadcast placement. # # Why not just the global owner? In HSDP mode each replicate row is # a full model instance with its own shard_group; the shard-dim @@ -256,8 +256,12 @@ def __init__( # fans out to the other rows' shard-owner ranks. # # Ranks outside the owner's shard column never need this buffer. + # + # Sharded AdamW is different: every shard rank owns its flat shard in + # ``_sharded_adamw_data`` and reconstructs the forward view via + # all-gather. It must not also keep a full fp32 owner copy here. is_in_owner_shard_column = dp_group.rank() == self.owner_shard - if is_in_owner_shard_column: + if is_in_owner_shard_column and self._dmuon_route != "sharded_adamw": master_dtype = self._master_dtype or self._orig_dtype self._owned_data = local_data.detach().to( device=device, diff --git a/dmuon/csrc/fast_clip_kernel.cu b/dmuon/csrc/fast_clip_kernel.cu new file mode 100644 index 0000000..3a25515 --- /dev/null +++ b/dmuon/csrc/fast_clip_kernel.cu @@ -0,0 +1,200 @@ +#include + +#include +#include +#include +#include +#include + +#include +#include + +#include +#include + +namespace { + +constexpr int kFloat32 = 0; +constexpr int kFloat16 = 1; +constexpr int kBFloat16 = 2; + +template +__device__ inline float load_as_float(const int64_t ptr, const int64_t offset) { + const auto* data = reinterpret_cast(static_cast(ptr)); + return static_cast(data[offset]); +} + +template +__device__ inline void scale_value(const int64_t ptr, const int64_t offset, const float coef) { + auto* data = reinterpret_cast(static_cast(ptr)); + const float value = static_cast(data[offset]) * coef; + data[offset] = static_cast(value); +} + +__global__ void segmented_l2_norm_sq_kernel( + const int64_t* __restrict__ ptrs, + const int64_t* __restrict__ numels, + const int32_t* __restrict__ dtypes, + const int32_t* __restrict__ segments, + const int32_t* __restrict__ job_tensor_ids, + const int64_t* __restrict__ job_offsets, + const int64_t chunk_size, + float* __restrict__ local_sq) { + extern __shared__ float smem[]; + + const int job = blockIdx.x; + const int tensor_id = job_tensor_ids[job]; + const int64_t tensor_numel = numels[tensor_id]; + const int64_t begin = job_offsets[job]; + const int64_t end = min(begin + chunk_size, tensor_numel); + const int dtype = dtypes[tensor_id]; + + float thread_sum = 0.0f; + for (int64_t offset = begin + threadIdx.x; offset < end; offset += blockDim.x) { + float value = 0.0f; + if (dtype == kFloat32) { + value = load_as_float(ptrs[tensor_id], offset); + } else if (dtype == kFloat16) { + value = load_as_float(ptrs[tensor_id], offset); + } else { + value = load_as_float(ptrs[tensor_id], offset); + } + thread_sum += value * value; + } + + smem[threadIdx.x] = thread_sum; + __syncthreads(); + + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + smem[threadIdx.x] += smem[threadIdx.x + stride]; + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + atomicAdd(local_sq + segments[tensor_id], smem[0]); + } +} + +__global__ void segmented_scale_kernel( + const int64_t* __restrict__ ptrs, + const int64_t* __restrict__ numels, + const int32_t* __restrict__ dtypes, + const int32_t* __restrict__ segments, + const int32_t* __restrict__ job_tensor_ids, + const int64_t* __restrict__ job_offsets, + const int64_t chunk_size, + const float* __restrict__ coefs) { + const int job = blockIdx.x; + const int tensor_id = job_tensor_ids[job]; + const int64_t tensor_numel = numels[tensor_id]; + const int64_t begin = job_offsets[job]; + const int64_t end = min(begin + chunk_size, tensor_numel); + const int dtype = dtypes[tensor_id]; + const float coef = coefs[segments[tensor_id]]; + + for (int64_t offset = begin + threadIdx.x; offset < end; offset += blockDim.x) { + if (dtype == kFloat32) { + scale_value(ptrs[tensor_id], offset, coef); + } else if (dtype == kFloat16) { + scale_value(ptrs[tensor_id], offset, coef); + } else { + scale_value(ptrs[tensor_id], offset, coef); + } + } +} + +void check_cuda_int_tensor(const torch::Tensor& tensor, const c10::ScalarType dtype) { + TORCH_CHECK(tensor.is_cuda(), "metadata tensor must be CUDA"); + TORCH_CHECK(tensor.is_contiguous(), "metadata tensor must be contiguous"); + TORCH_CHECK(tensor.scalar_type() == dtype, "metadata tensor has unexpected dtype"); +} + +} // namespace + +void segmented_l2_norm_sq( + torch::Tensor ptrs, + torch::Tensor numels, + torch::Tensor dtypes, + torch::Tensor segments, + torch::Tensor job_tensor_ids, + torch::Tensor job_offsets, + torch::Tensor local_sq, + int64_t chunk_size) { + check_cuda_int_tensor(ptrs, c10::kLong); + check_cuda_int_tensor(numels, c10::kLong); + check_cuda_int_tensor(dtypes, c10::kInt); + check_cuda_int_tensor(segments, c10::kInt); + check_cuda_int_tensor(job_tensor_ids, c10::kInt); + check_cuda_int_tensor(job_offsets, c10::kLong); + TORCH_CHECK(local_sq.is_cuda(), "local_sq must be CUDA"); + TORCH_CHECK(local_sq.is_contiguous(), "local_sq must be contiguous"); + TORCH_CHECK(local_sq.scalar_type() == c10::kFloat, "local_sq must be float32"); + TORCH_CHECK(chunk_size > 0, "chunk_size must be positive"); + + const auto jobs = job_tensor_ids.numel(); + if (jobs == 0) { + return; + } + + const c10::cuda::CUDAGuard device_guard(ptrs.device()); + constexpr int threads = 256; + const dim3 blocks(static_cast(jobs)); + const size_t smem = threads * sizeof(float); + segmented_l2_norm_sq_kernel<<>>( + ptrs.data_ptr(), + numels.data_ptr(), + dtypes.data_ptr(), + segments.data_ptr(), + job_tensor_ids.data_ptr(), + job_offsets.data_ptr(), + chunk_size, + local_sq.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void segmented_scale( + torch::Tensor ptrs, + torch::Tensor numels, + torch::Tensor dtypes, + torch::Tensor segments, + torch::Tensor job_tensor_ids, + torch::Tensor job_offsets, + torch::Tensor coefs, + int64_t chunk_size) { + check_cuda_int_tensor(ptrs, c10::kLong); + check_cuda_int_tensor(numels, c10::kLong); + check_cuda_int_tensor(dtypes, c10::kInt); + check_cuda_int_tensor(segments, c10::kInt); + check_cuda_int_tensor(job_tensor_ids, c10::kInt); + check_cuda_int_tensor(job_offsets, c10::kLong); + TORCH_CHECK(coefs.is_cuda(), "coefs must be CUDA"); + TORCH_CHECK(coefs.is_contiguous(), "coefs must be contiguous"); + TORCH_CHECK(coefs.scalar_type() == c10::kFloat, "coefs must be float32"); + TORCH_CHECK(chunk_size > 0, "chunk_size must be positive"); + + const auto jobs = job_tensor_ids.numel(); + if (jobs == 0) { + return; + } + + const c10::cuda::CUDAGuard device_guard(ptrs.device()); + constexpr int threads = 256; + const dim3 blocks(static_cast(jobs)); + segmented_scale_kernel<<>>( + ptrs.data_ptr(), + numels.data_ptr(), + dtypes.data_ptr(), + segments.data_ptr(), + job_tensor_ids.data_ptr(), + job_offsets.data_ptr(), + chunk_size, + coefs.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("segmented_l2_norm_sq", &segmented_l2_norm_sq, "Segmented L2 norm squared"); + m.def("segmented_scale", &segmented_scale, "Segmented gradient scaling"); +} diff --git a/dmuon/fast_clip.py b/dmuon/fast_clip.py new file mode 100644 index 0000000..4a285ee --- /dev/null +++ b/dmuon/fast_clip.py @@ -0,0 +1,490 @@ +"""Optional segmented CUDA fast path for gradient clipping. + +The public training semantics stay segment-local: each bucket gets its own +gradient norm and clip coefficient. The CUDA path only changes how those +segment-local norms are computed and applied. + +Distributed safety: the CUDA kernels accelerate only the *local* norm/scale +arithmetic and are selected **per tensor**. The collective that turns local +per-bucket squared norms into global ones is issued exactly once, with a fixed +shape derived from ``reduce`` flags that are identical on every rank. A rank +whose tensors happen to be ineligible for the kernel (non-contiguous, unusual +dtype, CPU) transparently uses the torch reference math for *those tensors* +only -- it never changes how many collectives are issued, so ranks with +different local eligibility stay in lockstep. +""" + +from __future__ import annotations + +import os +import warnings +from dataclasses import dataclass +from functools import lru_cache +from importlib import import_module +from typing import Sequence + +import torch +import torch.distributed as dist + +from .grad_clip import MuonGradClipStats, _iter_muon_grad_entries + + +_DTYPE_CODES = { + torch.float32: 0, + torch.float16: 1, + torch.bfloat16: 2, +} +_DEFAULT_CHUNK_SIZE = 262144 + + +@dataclass(frozen=True) +class GradClipBucket: + """One independently clipped gradient segment.""" + + name: str + grads: Sequence[torch.Tensor] + reduce: bool = True + + +@dataclass(frozen=True) +class GradClipBucketStats: + """Pre-clip norm and coefficient for one gradient segment.""" + + name: str + total_norm: float + max_norm: float | None + clip_coef: float + clipped: bool + param_count: int + found_inf: bool + fastpath: bool + + def as_dict(self) -> dict[str, object]: + return { + "name": self.name, + "total_norm": self.total_norm, + "max_norm": self.max_norm, + "clip_coef": self.clip_coef, + "clipped": self.clipped, + "param_count": self.param_count, + "found_inf": self.found_inf, + "fastpath": self.fastpath, + } + + +@dataclass(frozen=True) +class SegmentedGradClipResult: + """Result for a multi-segment clipping call.""" + + total_norm: torch.Tensor + bucket_stats: tuple[GradClipBucketStats, ...] + fastpath: bool + fallback_reason: str | None = None + + @property + def stats_by_name(self) -> dict[str, GradClipBucketStats]: + return {stats.name: stats for stats in self.bucket_stats} + + +@lru_cache(maxsize=1) +def _load_fast_clip_extension(): + """Load the compiled fast-clip CUDA extension, or ``None`` if unavailable. + + Distinguishes "not built" (expected when installed without a CUDA + toolchain) from "built but failed to load" (an ABI mismatch, which is a + real problem worth surfacing). ``DMUON_FAST_CLIP=0`` disables the fast + path; ``DMUON_FAST_CLIP_VERBOSE=1`` re-raises instead of degrading. + """ + + if os.environ.get("DMUON_FAST_CLIP", "1") == "0": + return None + verbose = os.environ.get("DMUON_FAST_CLIP_VERBOSE", "0") == "1" + try: + return import_module("dmuon._fast_clip_cuda") + except ModuleNotFoundError: + # Extension was never compiled (e.g. installed without CUDA_HOME, or + # with --no-build-isolation missing). Expected; stay quiet. + if verbose: + raise + return None + except Exception as exc: + # Built but unloadable -- almost always a torch/CUDA ABI mismatch. + # Warn once (lru_cache makes this a single call) so the silent + # slow-path fallback is diagnosable. + if verbose: + raise + warnings.warn( + f"dmuon fast-clip extension failed to load ({exc!r}); " + "falling back to the pure-Python clip. Rebuild with a CUDA " + "toolchain matching your torch, or set DMUON_FAST_CLIP_VERBOSE=1.", + RuntimeWarning, + stacklevel=2, + ) + return None + + +def _resolve_chunk_size(chunk_size: int | None) -> int: + if chunk_size is not None: + return max(1, int(chunk_size)) + raw_value = os.environ.get("DMUON_FAST_CLIP_CHUNK_SIZE") + if raw_value is None: + return _DEFAULT_CHUNK_SIZE + try: + return max(1, int(raw_value)) + except ValueError: + return _DEFAULT_CHUNK_SIZE + + +def _resolve_device(buckets: Sequence[GradClipBucket]) -> torch.device: + for bucket in buckets: + for grad in bucket.grads: + if grad is not None: + return grad.device + if torch.cuda.is_available(): + return torch.device("cuda", torch.cuda.current_device()) + return torch.device("cpu") + + +def _tensor_kernel_eligible(grad: torch.Tensor, device: torch.device) -> bool: + """Whether the CUDA kernel can process this single tensor.""" + + return ( + grad.device == device + and device.type == "cuda" + and not grad.is_sparse + and grad.is_contiguous() + and grad.dtype in _DTYPE_CODES + ) + + +def _build_metadata( + tensors: Sequence[torch.Tensor], + segments: Sequence[int], + *, + device: torch.device, + chunk_size: int, +): + ptrs = torch.tensor( + [tensor.data_ptr() for tensor in tensors], + device=device, + dtype=torch.int64, + ) + numels = torch.tensor( + [tensor.numel() for tensor in tensors], + device=device, + dtype=torch.int64, + ) + dtypes = torch.tensor( + [_DTYPE_CODES[tensor.dtype] for tensor in tensors], + device=device, + dtype=torch.int32, + ) + tensor_segments = torch.tensor(segments, device=device, dtype=torch.int32) + + job_tensor_ids: list[int] = [] + job_offsets: list[int] = [] + for tensor_id, tensor in enumerate(tensors): + numel = int(tensor.numel()) + for offset in range(0, numel, chunk_size): + job_tensor_ids.append(tensor_id) + job_offsets.append(offset) + + return { + "ptrs": ptrs, + "numels": numels, + "dtypes": dtypes, + "segments": tensor_segments, + "job_tensor_ids": torch.tensor( + job_tensor_ids, device=device, dtype=torch.int32 + ), + "job_offsets": torch.tensor(job_offsets, device=device, dtype=torch.int64), + } + + +def _split_by_kernel_eligibility( + buckets: Sequence[GradClipBucket], + device: torch.device, + ext, +) -> tuple[list[torch.Tensor], list[int], list[tuple[torch.Tensor, int]]]: + """Partition every live grad into kernel-eligible and torch-fallback sets. + + Returns ``(kernel_tensors, kernel_segments, torch_pairs)`` where + ``torch_pairs`` is ``[(grad, segment_index), ...]``. The decision is + purely local and per-tensor; it never affects collective shape. + """ + + kernel_tensors: list[torch.Tensor] = [] + kernel_segments: list[int] = [] + torch_pairs: list[tuple[torch.Tensor, int]] = [] + for segment, bucket in enumerate(buckets): + for grad in bucket.grads: + if grad is None: + continue + if ext is not None and _tensor_kernel_eligible(grad, device): + kernel_tensors.append(grad) + kernel_segments.append(segment) + else: + torch_pairs.append((grad, segment)) + return kernel_tensors, kernel_segments, torch_pairs + + +def _segmented_local_sq( + buckets: Sequence[GradClipBucket], + *, + device: torch.device, + ext, + chunk_size: int, +) -> torch.Tensor: + """Per-bucket local sum of squared grads (no collective). + + Kernel-eligible tensors are summed by the CUDA kernel (``atomicAdd`` into + ``local_sq``); the rest are summed with torch and accumulated into the same + vector. Both run on the current stream, so the torch adds and the kernel + atomics compose correctly. + """ + + num_segments = len(buckets) + local_sq = torch.zeros(num_segments, device=device, dtype=torch.float32) + kernel_tensors, kernel_segments, torch_pairs = _split_by_kernel_eligibility( + buckets, device, ext + ) + + for grad, segment in torch_pairs: + local_sq[segment] += ( + grad.detach().float().pow(2).sum().to(device=device, dtype=torch.float32) + ) + + if kernel_tensors: + metadata = _build_metadata( + kernel_tensors, kernel_segments, device=device, chunk_size=chunk_size + ) + ext.segmented_l2_norm_sq( + metadata["ptrs"], + metadata["numels"], + metadata["dtypes"], + metadata["segments"], + metadata["job_tensor_ids"], + metadata["job_offsets"], + local_sq, + chunk_size, + ) + return local_sq + + +def _segmented_scale( + buckets: Sequence[GradClipBucket], + coefs: torch.Tensor, + *, + device: torch.device, + ext, + chunk_size: int, +) -> None: + """Scale each grad in place by its bucket coefficient (no collective).""" + + kernel_tensors, kernel_segments, torch_pairs = _split_by_kernel_eligibility( + buckets, device, ext + ) + + for grad, segment in torch_pairs: + grad.mul_(coefs[segment].to(device=grad.device)) + + if kernel_tensors: + metadata = _build_metadata( + kernel_tensors, kernel_segments, device=device, chunk_size=chunk_size + ) + ext.segmented_scale( + metadata["ptrs"], + metadata["numels"], + metadata["dtypes"], + metadata["segments"], + metadata["job_tensor_ids"], + metadata["job_offsets"], + coefs, + chunk_size, + ) + + +def _apply_segment_reductions( + local_sq: torch.Tensor, + reduce_flags: Sequence[int], + process_group=None, +) -> torch.Tensor: + """The single collective: sum reduce=True segments across ranks. + + Issued unconditionally whenever the process group is live and any bucket + reduces -- both conditions are derived identically on every rank, so the + collective shape is rank-uniform. + """ + + if not ( + dist.is_available() + and dist.is_initialized() + and any(bool(flag) for flag in reduce_flags) + ): + return local_sq + + reduce_mask = torch.tensor( + reduce_flags, + device=local_sq.device, + dtype=torch.bool, + ) + reduced_sq = torch.where(reduce_mask, local_sq, torch.zeros_like(local_sq)) + dist.all_reduce(reduced_sq, op=dist.ReduceOp.SUM, group=process_group) + return torch.where(reduce_mask, reduced_sq, local_sq) + + +def _build_stats( + buckets: Sequence[GradClipBucket], + total_norms: torch.Tensor, + coefs: torch.Tensor, + *, + max_norm: float | None, + fastpath: bool, +) -> tuple[GradClipBucketStats, ...]: + """Assemble per-bucket stats with a single batched device-to-host copy.""" + + if len(buckets) == 0: + return tuple() + + stacked = torch.stack( + [total_norms, coefs, torch.isfinite(total_norms).to(total_norms.dtype)] + ) + host = stacked.detach().to("cpu") + total_values = host[0].tolist() + coef_values = host[1].tolist() + finite_values = host[2].tolist() + + stats: list[GradClipBucketStats] = [] + for idx, bucket in enumerate(buckets): + total = float(total_values[idx]) + coef = float(coef_values[idx]) + finite = bool(finite_values[idx]) + clipped = bool(max_norm is not None and finite and total > float(max_norm)) + param_count = sum(1 for grad in bucket.grads if grad is not None) + stats.append( + GradClipBucketStats( + name=bucket.name, + total_norm=total, + max_norm=None if max_norm is None else float(max_norm), + clip_coef=coef, + clipped=clipped, + param_count=param_count, + found_inf=not finite, + fastpath=fastpath, + ) + ) + return tuple(stats) + + +@torch.no_grad() +def clip_grad_norm_buckets_( + buckets: Sequence[GradClipBucket], + max_norm: float | None, + *, + process_group=None, + chunk_size: int | None = None, +) -> SegmentedGradClipResult: + """Clip gradient buckets with segment-local, distributed-safe semantics. + + Each bucket keeps its own gradient norm and clip coefficient. The CUDA + kernels accelerate the eligible tensors; ineligible tensors (or all of them + when the extension is unavailable) use the torch reference math. The + fast/slow choice is per tensor and purely local: it never changes the + number or shape of collective operations, so this is safe to call from + every rank in lockstep. + """ + + device = _resolve_device(buckets) + ext = _load_fast_clip_extension() + resolved_chunk_size = _resolve_chunk_size(chunk_size) + reduce_flags = [1 if bucket.reduce else 0 for bucket in buckets] + + # Stage 1 -- local per-bucket squared norm (no collective). + local_sq = _segmented_local_sq( + buckets, device=device, ext=ext, chunk_size=resolved_chunk_size + ) + + # Stage 2 -- the single, rank-uniform collective. + global_sq = _apply_segment_reductions(local_sq, reduce_flags, process_group) + + # Stage 3 -- total norm + coefficient (torch, matches reference semantics). + total_norms = torch.sqrt(global_sq) + if max_norm is None: + coefs = torch.ones_like(total_norms) + else: + coefs = torch.clamp(float(max_norm) / (total_norms + 1e-6), max=1.0) + + # Stage 4 -- scale in place (no collective). found_inf never gates this. + if max_norm is not None: + _segmented_scale( + buckets, coefs, device=device, ext=ext, chunk_size=resolved_chunk_size + ) + + # Stage 5 -- stats (single batched D2H). + stats = _build_stats( + buckets, total_norms, coefs, max_norm=max_norm, fastpath=ext is not None + ) + combined_total = torch.linalg.vector_norm(total_norms, 2.0) + return SegmentedGradClipResult( + combined_total, + stats, + fastpath=ext is not None, + fallback_reason=None if ext is not None else "cuda_extension_unavailable", + ) + + +@torch.no_grad() +def try_clip_optimizer_grad_norm_buckets_( + optimizer, + *, + regular_grads: Sequence[torch.Tensor], + adamw_grads: Sequence[torch.Tensor], + max_norm: float, + process_group=None, +) -> SegmentedGradClipResult | None: + """Fast path for Wall-X style regular/Muon/dedicated-AdamW clipping. + + Returns ``None`` only when the fast-clip CUDA extension is entirely + unavailable, letting the caller run its own path. Extension availability + is a property of the install and is therefore uniform across ranks, so this + ``None`` decision is rank-consistent and cannot desync collectives. Once a + result is produced, clipping goes through the distributed-safe + :func:`clip_grad_norm_buckets_`. + """ + + if _load_fast_clip_extension() is None: + return None + + ensure = getattr(optimizer, "_ensure_grads_ready", None) + if ensure is None: + return None + try: + ensure(coalesce_wait=True) + except TypeError: + ensure() + + muon_entries = tuple(_iter_muon_grad_entries(optimizer)) + muon_grads = [entry.grad for entry in muon_entries] + result = clip_grad_norm_buckets_( + ( + GradClipBucket("regular", regular_grads, reduce=False), + GradClipBucket("muon", muon_grads, reduce=True), + GradClipBucket("adamw", adamw_grads, reduce=True), + ), + max_norm, + process_group=process_group, + ) + + muon_stats = result.stats_by_name.get("muon") + if muon_stats is not None: + optimizer._last_muon_grad_clip_stats = MuonGradClipStats( + total_norm=muon_stats.total_norm, + max_norm=muon_stats.max_norm, + norm_type=2.0, + clip_coef=muon_stats.clip_coef, + clipped=muon_stats.clipped, + param_count=muon_stats.param_count, + found_inf=muon_stats.found_inf, + strategy="segmented_cuda_fastpath", + ) + return result diff --git a/dmuon/optim/muon.py b/dmuon/optim/muon.py index 5d77a13..8d3eb6c 100644 --- a/dmuon/optim/muon.py +++ b/dmuon/optim/muon.py @@ -838,6 +838,26 @@ def _build_semantic_param_groups( {"dmuon_route": getattr(dp, "_dmuon_route", "muon")} ) ) + # A param constructed on the sharded_adamw route owns only a + # rank-local shard (``_sharded_adamw_data``) and has no full + # ``_owned_data`` copy. Moving it to another route here would + # leave later muon/adamw steps dereferencing a ``None`` + # ``_owned_data`` (crash) or silently skipping the update. + # The shard storage cannot be rebuilt after dedicate_params, + # so reject the move with a clear error. + if ( + getattr(dp, "_sharded_adamw_data", None) is not None + and dedicated_route != "sharded_adamw" + ): + raise RuntimeError( + f"DMuon param {getattr(dp, 'param_name', '')} " + f"was constructed on the 'sharded_adamw' route but a " + f"param_group tries to move it to '{dedicated_route}'. " + "A sharded param has no full owner copy and its shard " + "storage is fixed at dedicate_params() time; keep it on " + "'sharded_adamw' or change the route hint before " + "constructing Muon." + ) if dedicated_route == "muon": dp._dmuon_route = "muon" dp._dmuon_adamw_replicate_allreduce = False @@ -1288,12 +1308,16 @@ def consume_last_step_profile(self) -> dict[str, object]: self._last_step_profile_events = [] return profile - def _ensure_grads_ready(self) -> None: + def _ensure_grads_ready(self, *, coalesce_wait: bool = False) -> None: """Prepare DMuon Muon grads once before clipping or stepping.""" if self._grads_ready: return - prepare_muon_grads(self.model) + prepare_muon_grads( + self.model, + use_reduce_stream=coalesce_wait, + coalesce_wait=coalesce_wait, + ) self._grads_ready = True @torch.no_grad() diff --git a/dmuon/utils.py b/dmuon/utils.py index c8f000d..5c3a5af 100644 --- a/dmuon/utils.py +++ b/dmuon/utils.py @@ -138,7 +138,49 @@ def wait_group_muon_grads(g) -> None: g._muon_grad_ready_refs = [] -def prepare_muon_grads(model: nn.Module, *, use_reduce_stream: bool = False) -> None: +def _finish_group_muon_grads_after_coalesced_wait(g) -> None: + """Clear one prepared group after a caller waited for the aggregate event.""" + + if getattr(g, "_tp_gather_event", None) is not None: + materialize = getattr(g, "_materialize_tp_gathered_grads", None) + if materialize is None: + wait_group_muon_grads(g) + return + materialize() + setattr(g, "_tp_gather_event", None) + if hasattr(g, "_tp_gather_refs"): + g._tp_gather_refs = [] + setattr(g, "_muon_grad_ready_event", None) + if hasattr(g, "_muon_grad_ready_refs"): + g._muon_grad_ready_refs = [] + + +def _coalesced_wait_prepared_muon_grads(groups) -> bool: + reduce_stream = None + for group in groups: + group_stream = getattr(getattr(group, "comm_ctx", None), "reduce_stream", None) + if group_stream is None: + return False + if reduce_stream is None: + reduce_stream = group_stream + elif reduce_stream is not group_stream: + return False + if reduce_stream is None: + return True + + ready_event = reduce_stream.record_event() + torch.cuda.current_stream().wait_event(ready_event) + for group in groups: + _finish_group_muon_grads_after_coalesced_wait(group) + return True + + +def prepare_muon_grads( + model: nn.Module, + *, + use_reduce_stream: bool = False, + coalesce_wait: bool = False, +) -> None: """Prepare all pending Muon gradients after backward. Gradient reduces are dispatched asynchronously during backward. This @@ -146,10 +188,18 @@ def prepare_muon_grads(model: nn.Module, *, use_reduce_stream: bool = False) -> the TP gather that materializes ``_tp_full_grad`` for the TP owner. It drains all prepared groups before returning so callers can immediately run Muon update code. + + When every group shares one reduce stream, ``use_reduce_stream`` combined + with ``coalesce_wait`` lets advanced callers (such as the fused fast-clip + path) issue a single aggregate stream wait instead of draining each group + individually. Both default to ``False`` and preserve the per-group wait. """ groups = _ordered_post_step_groups(model) for g in groups: prepare_group_muon_grads(g, use_reduce_stream=use_reduce_stream) + if coalesce_wait and use_reduce_stream: + if _coalesced_wait_prepared_muon_grads(groups): + return for g in groups: wait_group_muon_grads(g) diff --git a/docs/contributing.md b/docs/contributing.md index da5b1f0..8089d12 100644 --- a/docs/contributing.md +++ b/docs/contributing.md @@ -12,7 +12,7 @@ ```bash # Clone and install in editable mode -git clone https://github.com/StarrickLiu/dmuon +git clone https://github.com/X-Square-Robot/dmuon cd dmuon # Create environment (conda or venv) @@ -128,8 +128,8 @@ Documentation is deployed automatically to GitHub Pages on every push to ## Community -- **Bug reports and feature requests:** [GitHub Issues](https://github.com/StarrickLiu/dmuon/issues) -- **Design discussions:** [GitHub Discussions](https://github.com/StarrickLiu/dmuon/discussions) +- **Bug reports and feature requests:** [GitHub Issues](https://github.com/X-Square-Robot/dmuon/issues) +- **Design discussions:** [GitHub Discussions](https://github.com/X-Square-Robot/dmuon/discussions) - No Discord at this time Please search existing issues before opening a new one. diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md index f440c92..2782ef5 100644 --- a/docs/getting-started/installation.md +++ b/docs/getting-started/installation.md @@ -30,7 +30,7 @@ === "From source (recommended)" ```bash - git clone https://github.com/StarrickLiu/dmuon + git clone https://github.com/X-Square-Robot/dmuon cd dmuon pip install -e . ``` @@ -51,7 +51,7 @@ === "Development (editable + test deps)" ```bash - git clone https://github.com/StarrickLiu/dmuon + git clone https://github.com/X-Square-Robot/dmuon cd dmuon pip install -e ".[dev]" ``` @@ -86,6 +86,55 @@ The compiled artifact is cached in `~/.cache/dmuon/`. --- +## Optional: Fast Gradient Clipping (CUDA) + +DMuon ships an optional CUDA kernel that fuses **segmented gradient clipping** +— the per-bucket norm, clip coefficient, and in-place scaling for the +`regular` / `muon` / `adamw` gradient groups — into a single pass. The +training semantics are identical to the pure-Python path: each bucket still +gets its own norm and its own clip coefficient. Only the arithmetic moves to +the GPU. + +`torch` is intentionally **not** a build dependency (pinning it would force an +isolated build to download a multi-GB generic torch and link the kernel against +it — an ABI-mismatch risk). So to compile the kernel, build **without isolation** +in an environment that already has torch, with a CUDA toolchain on `PATH`: + +```bash +# nvcc / CUDA_HOME must be visible; torch must already be installed +pip install -e . --no-build-isolation +``` + +- With `--no-build-isolation` and `CUDA_HOME` present, `dmuon._fast_clip_cuda` + is built against your real torch and used automatically. +- A plain `pip install -e .` (isolated build) has no torch in the build env, so + `setup.py` skips the extension and DMuon uses the equivalent pure-Python clip + at runtime. Nothing breaks — clipping is just computed on the host side. +- If the extension is **built but fails to load** (e.g. a later torch/CUDA + upgrade breaks its ABI), DMuon warns once and falls back to Python. Set + `DMUON_FAST_CLIP_VERBOSE=1` to raise the underlying error instead. + +### Build & runtime toggles + +| Variable | Effect | +|----------|--------| +| `DMUON_BUILD_FAST_CLIP=0` | Skip building the CUDA extension at install time. | +| `DMUON_FAST_CLIP=0` | Disable the fast path at runtime (use pure Python). | +| `DMUON_FAST_CLIP_CHUNK_SIZE` | Per-tensor chunk size for the kernel (default `262144`). | +| `DMUON_FAST_CLIP_VERBOSE=1` | Raise the import error instead of silently falling back — use when a build "should" have worked. | + +The runtime path also falls back to Python automatically for inputs outside the +kernel contract (non-contiguous, sparse, unsupported dtype) or when a +non-finite bucket norm is detected, so a missing or stale extension never +changes results. + +!!! note "The build needs a compiler, not a specific GPU" + Compiling the clip kernels only requires a host CUDA compiler (`nvcc`), not + an SM80+ GPU — the kernels are architecture-agnostic. Match the CUDA + toolkit to your PyTorch CUDA build (11.8 / 12.1 / 12.4). + +--- + ## Verify Installation ```python title="verify_install.py" @@ -136,6 +185,12 @@ NS backend : Gram NS · kernel=cublas (SM80, universal fallback) If the build still fails, the compiled fallback is used automatically — Newton-Schulz will still run correctly, just slightly slower. +**Fast-clip kernel not used (`fastpath=False` in clip stats)** +: The `dmuon._fast_clip_cuda` extension was not built (no `CUDA_HOME` at install + time) or was disabled via `DMUON_FAST_CLIP=0`. Reinstall with a CUDA toolchain + on `PATH`, or set `DMUON_FAST_CLIP_VERBOSE=1` to surface the import error. + Gradient clipping stays correct via the Python path in the meantime. + --- ## See Also diff --git a/docs/getting-started/quickstart.md b/docs/getting-started/quickstart.md index a63510b..c1d5210 100644 --- a/docs/getting-started/quickstart.md +++ b/docs/getting-started/quickstart.md @@ -10,7 +10,7 @@ ## Step 1 — Install ```bash -git clone https://github.com/StarrickLiu/dmuon && cd dmuon +git clone https://github.com/X-Square-Robot/dmuon && cd dmuon pip install -e . ``` diff --git a/docs/guides/training.md b/docs/guides/training.md index 45967f1..6fffe3e 100644 --- a/docs/guides/training.md +++ b/docs/guides/training.md @@ -311,6 +311,14 @@ The default strategy is global p-norm clipping over Muon gradients. Custom strategies can be registered with `dmuon.register_muon_grad_clip_strategy(...)` for future schemes such as MuonClip or projection-specific clipping. +!!! tip "Fused CUDA clip for regular / Muon / AdamW" + When you clip all three gradient groups together (the Wall-X pattern), + `dmuon.try_clip_optimizer_grad_norm_buckets_(...)` computes the norm, clip + coefficient, and in-place scaling for every bucket in a single CUDA pass + with the same per-bucket semantics. It uses the optional fast-clip + extension (see [Installation](../getting-started/installation.md)) and + falls back to pure Python automatically when the extension is unavailable. + ## Logging and Debugging ### Check NS Backend diff --git a/docs/index.md b/docs/index.md index 6cffa41..9b0a9ba 100644 --- a/docs/index.md +++ b/docs/index.md @@ -173,4 +173,4 @@ The 128GPU rows cover FSDP-Z2, FSDP-Z3, HSDP-Z2, and HSDP-Z3. DMuon builds on dedicated ownership pioneered by ZeRO-1 (Rajbhandari et al., 2020) and Distributed Shampoo (Shi et al., 2023). Gram Newton-Schulz kernel adapted from Dao et al., 2026. -GitHub: [StarrickLiu/dmuon](https://github.com/StarrickLiu/dmuon)  ·  arXiv preprint: [TBD] +GitHub: [X-Square-Robot/dmuon](https://github.com/X-Square-Robot/dmuon)  ·  arXiv preprint: [TBD] diff --git a/docs/reference/api.md b/docs/reference/api.md index 8c46135..baafe00 100644 --- a/docs/reference/api.md +++ b/docs/reference/api.md @@ -256,6 +256,50 @@ Return type for DMuon gradient clipping. --- +### clip_grad_norm_buckets_ + +Clip independent gradient buckets with segment-local semantics, using the +optional CUDA fast path when available and pure Python otherwise. See +[Fast Gradient Clipping](../getting-started/installation.md) for the extension. + +::: dmuon.clip_grad_norm_buckets_ + +--- + +### try_clip_optimizer_grad_norm_buckets_ + +Segmented `regular` / `muon` / `adamw` clip entry point for Wall-X-style +optimizers. Returns `None` when the CUDA extension is unavailable so callers +can fall back without any gradient mutation. + +::: dmuon.try_clip_optimizer_grad_norm_buckets_ + +--- + +### GradClipBucket + +One independently clipped gradient segment. + +::: dmuon.GradClipBucket + +--- + +### GradClipBucketStats + +Per-segment pre-clip norm and coefficient. + +::: dmuon.GradClipBucketStats + +--- + +### SegmentedGradClipResult + +Aggregate result of a multi-segment clip call. + +::: dmuon.SegmentedGradClipResult + +--- + ### get_dedicated_params Enumerate all `DedicatedParam` objects across the model. diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index df03920..c4286bf 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -16,7 +16,7 @@ **Fix:** ```bash - git clone https://github.com/StarrickLiu/dmuon && cd dmuon + git clone https://github.com/X-Square-Robot/dmuon && cd dmuon pip install -e . ``` Verify: `python -c "import dmuon; print(dmuon.__version__)"`. diff --git a/docs/zh/contributing.md b/docs/zh/contributing.md index 38a6340..2eb24fb 100644 --- a/docs/zh/contributing.md +++ b/docs/zh/contributing.md @@ -11,7 +11,7 @@ ```bash # 克隆并以可编辑模式安装 -git clone https://github.com/StarrickLiu/dmuon +git clone https://github.com/X-Square-Robot/dmuon cd dmuon # 创建环境（conda 或 venv） @@ -123,8 +123,8 @@ DMuon 遵循语义化版本：`MAJOR.MINOR.PATCH`。 ## 社区 -- **Bug 报告和功能请求：** [GitHub Issues](https://github.com/StarrickLiu/dmuon/issues) -- **设计讨论：** [GitHub Discussions](https://github.com/StarrickLiu/dmuon/discussions) +- **Bug 报告和功能请求：** [GitHub Issues](https://github.com/X-Square-Robot/dmuon/issues) +- **设计讨论：** [GitHub Discussions](https://github.com/X-Square-Robot/dmuon/discussions) - 暂无 Discord 提 issue 前请先搜索已有 issue。 diff --git a/docs/zh/getting-started/installation.md b/docs/zh/getting-started/installation.md index cefc388..8075d77 100644 --- a/docs/zh/getting-started/installation.md +++ b/docs/zh/getting-started/installation.md @@ -30,7 +30,7 @@ === "从源码安装（推荐）" ```bash - git clone https://github.com/StarrickLiu/dmuon + git clone https://github.com/X-Square-Robot/dmuon cd dmuon pip install -e . ``` @@ -51,7 +51,7 @@ === "开发模式（可编辑 + 测试依赖）" ```bash - git clone https://github.com/StarrickLiu/dmuon + git clone https://github.com/X-Square-Robot/dmuon cd dmuon pip install -e ".[dev]" ``` @@ -84,6 +84,47 @@ pip install -e ".[syrk]" --- +## 可选：快速梯度裁剪（CUDA） + +DMuon 附带一个可选的 CUDA 内核，把**分段梯度裁剪**——即 `regular` / `muon` / +`adamw` 三个梯度组各自的范数、裁剪系数与就地缩放——融合到一趟计算里。训练语义与 +纯 Python 路径完全一致：每个分段仍各自计算范数、各自使用独立的裁剪系数，只是把运算 +搬到了 GPU 上。 + +`torch` **刻意不作为**构建依赖（若固定它，隔离构建会去下载多 GB 的通用 torch 并把 +内核链接到它上——有 ABI 错配风险）。因此要编译内核，需在**已装 torch** 的环境里、 +`PATH` 带上 CUDA 工具链、**关闭构建隔离**安装： + +```bash +# 需要 nvcc / CUDA_HOME 可见，且环境里已装 torch +pip install -e . --no-build-isolation +``` + +- 加了 `--no-build-isolation` 且有 `CUDA_HOME` 时，`dmuon._fast_clip_cuda` 会针对你 + 真实的 torch 编译并自动启用。 +- 普通 `pip install -e .`（隔离构建）的构建环境里没有 torch，`setup.py` 会跳过扩展， + 运行时使用等价的纯 Python 裁剪。不会报错——只是把裁剪放到主机端计算。 +- 若扩展**编译出来但加载失败**（例如之后升级 torch/CUDA 破坏了 ABI），DMuon 会 warn + 一次并回退到 Python。设 `DMUON_FAST_CLIP_VERBOSE=1` 可改为直接抛出底层错误。 + +### 编译与运行时开关 + +| 变量 | 作用 | +|------|------| +| `DMUON_BUILD_FAST_CLIP=0` | 安装时跳过编译该 CUDA 扩展。 | +| `DMUON_FAST_CLIP=0` | 运行时禁用快速路径（改用纯 Python）。 | +| `DMUON_FAST_CLIP_CHUNK_SIZE` | 内核的单张量分块大小（默认 `262144`）。 | +| `DMUON_FAST_CLIP_VERBOSE=1` | 直接抛出导入错误而非静默回退——用于本应编译成功却没生效时排查。 | + +运行时路径对于不满足内核契约的输入（非连续、稀疏、不支持的 dtype）或检测到非有限的 +分段范数时，也会自动回退到 Python，因此扩展缺失或过时都不会改变结果。 + +!!! note "编译需要的是编译器，而非特定 GPU" + 编译裁剪内核只需要主机端的 CUDA 编译器（`nvcc`），并不要求 SM80+ GPU——这些 + 内核与架构无关。将 CUDA 工具包与你的 PyTorch CUDA 版本对齐即可（11.8 / 12.1 / 12.4）。 + +--- + ## 验证安装 ```python title="verify_install.py" @@ -132,6 +173,11 @@ NS 后端 : Gram NS · kernel=cublas (SM80, universal fallback) 若构建仍失败，编译后备方案会自动启用—— Newton-Schulz 仍可正确运行，只是稍慢。 +**快速裁剪内核未生效（裁剪统计里 `fastpath=False`）** +: `dmuon._fast_clip_cuda` 扩展没有被编译（安装时无 `CUDA_HOME`），或被 + `DMUON_FAST_CLIP=0` 禁用了。请在 `PATH` 中带上 CUDA 工具链重新安装，或设置 + `DMUON_FAST_CLIP_VERBOSE=1` 查看导入错误。其间梯度裁剪仍通过 Python 路径正确执行。 + --- ## 另请参见 diff --git a/docs/zh/getting-started/quickstart.md b/docs/zh/getting-started/quickstart.md index 1acd0ee..9b04543 100644 --- a/docs/zh/getting-started/quickstart.md +++ b/docs/zh/getting-started/quickstart.md @@ -9,7 +9,7 @@ ## 第一步 — 安装 ```bash -git clone https://github.com/StarrickLiu/dmuon && cd dmuon +git clone https://github.com/X-Square-Robot/dmuon && cd dmuon pip install -e . ``` diff --git a/docs/zh/guides/training.md b/docs/zh/guides/training.md index 0328f76..9d315b8 100644 --- a/docs/zh/guides/training.md +++ b/docs/zh/guides/training.md @@ -304,6 +304,13 @@ for step, batch in enumerate(dataloader): MuonClip、QK/投影层专用 clip 等方案，可以通过 `dmuon.register_muon_grad_clip_strategy(...)` 注册自定义策略。 +!!! tip "regular / Muon / AdamW 的融合 CUDA 裁剪" + 当你需要把三个梯度组一起裁剪（Wall-X 模式）时， + `dmuon.try_clip_optimizer_grad_norm_buckets_(...)` 会在一趟 CUDA 计算里 + 完成每个分段的范数、裁剪系数与就地缩放，且保持逐分段语义。它使用可选的快速 + 裁剪扩展（见[安装](../getting-started/installation.md)），扩展不可用时自动 + 回退到纯 Python。 + ## 日志与调试 ### 检查 NS 后端 diff --git a/docs/zh/index.md b/docs/zh/index.md index 0392dd2..27f2c78 100644 --- a/docs/zh/index.md +++ b/docs/zh/index.md @@ -169,4 +169,4 @@ DMuon 为每个矩阵参数指定唯一的**所有者 rank**。所有者存储 DMuon 建立在 ZeRO-1（Rajbhandari 等，2020）和 Distributed Shampoo（Shi 等，2023） 所开创的专属所有权原语之上。Gram Newton-Schulz 核改编自 Dao 等，2026。 -GitHub：[StarrickLiu/dmuon](https://github.com/StarrickLiu/dmuon)  ·  arXiv 预印本：[TBD] +GitHub：[X-Square-Robot/dmuon](https://github.com/X-Square-Robot/dmuon)  ·  arXiv 预印本：[TBD] diff --git a/docs/zh/reference/api.md b/docs/zh/reference/api.md index 99e5551..f5f4a96 100644 --- a/docs/zh/reference/api.md +++ b/docs/zh/reference/api.md @@ -252,6 +252,48 @@ DMuon 梯度裁剪的返回类型。 --- +### clip_grad_norm_buckets_ + +以分段局部语义裁剪相互独立的梯度分段；可用时走可选的 CUDA 快速路径，否则回退纯 +Python。扩展的编译见[快速梯度裁剪](../getting-started/installation.md)。 + +::: dmuon.clip_grad_norm_buckets_ + +--- + +### try_clip_optimizer_grad_norm_buckets_ + +面向 Wall-X 式优化器的 `regular` / `muon` / `adamw` 分段裁剪入口。CUDA 扩展不可用时 +返回 `None`，调用方可在不触碰梯度的情况下回退。 + +::: dmuon.try_clip_optimizer_grad_norm_buckets_ + +--- + +### GradClipBucket + +一个可独立裁剪的梯度分段。 + +::: dmuon.GradClipBucket + +--- + +### GradClipBucketStats + +单个分段裁剪前的范数与系数。 + +::: dmuon.GradClipBucketStats + +--- + +### SegmentedGradClipResult + +多分段裁剪调用的聚合结果。 + +::: dmuon.SegmentedGradClipResult + +--- + ### get_dedicated_params 枚举模型中所有 `DedicatedParam` 对象。用于检查 ownership 分配、参数数量 diff --git a/docs/zh/troubleshooting.md b/docs/zh/troubleshooting.md index a65e1d1..fdb954d 100644 --- a/docs/zh/troubleshooting.md +++ b/docs/zh/troubleshooting.md @@ -14,7 +14,7 @@ **修复：** ```bash - git clone https://github.com/StarrickLiu/dmuon && cd dmuon + git clone https://github.com/X-Square-Robot/dmuon && cd dmuon pip install -e . ``` 验证：`python -c "import dmuon; print(dmuon.__version__)"`。 diff --git a/mkdocs.yml b/mkdocs.yml index 8aab21a..f233af7 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -1,8 +1,8 @@ site_name: DMuon site_description: Dedicated ownership for Muon optimizer on PyTorch DDP, FSDP2, and HSDP -site_url: https://starrickliu.github.io/dmuon/ -repo_url: https://github.com/StarrickLiu/dmuon -repo_name: StarrickLiu/dmuon +site_url: https://x-square-robot.github.io/dmuon/ +repo_url: https://github.com/X-Square-Robot/dmuon +repo_name: X-Square-Robot/dmuon exclude_docs: | internal/ diff --git a/pyproject.toml b/pyproject.toml index 9b4550d..5587231 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,4 +1,10 @@ [build-system] +# torch is intentionally NOT listed here: pinning it would force PEP 517 build +# isolation to download a multi-GB generic torch into a throwaway env and link +# the CUDA extension against it (ABI-mismatch risk). To compile the optional +# fast-clip kernel, install with `--no-build-isolation` in an env that already +# has torch; otherwise setup.py skips the extension and the pure-Python clip is +# used. See docs/getting-started/installation.md. requires = ["setuptools>=77.0", "wheel"] build-backend = "setuptools.build_meta" @@ -62,6 +68,9 @@ exclude = [ "site*", ] +[tool.setuptools.package-data] +dmuon = ["csrc/*.cu"] + [tool.ruff] line-length = 100 target-version = "py310" diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..d71ee5d --- /dev/null +++ b/setup.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import os + +from setuptools import setup + + +def _extension_config(): + if os.environ.get("DMUON_BUILD_FAST_CLIP", "1") == "0": + return [], {} + try: + from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME + except Exception as exc: + print(f"DMuon fast clip CUDA extension disabled: {exc}") + return [], {} + + if CUDA_HOME is None: + print("DMuon fast clip CUDA extension disabled: CUDA_HOME is not available") + return [], {} + + ext_modules = [ + CUDAExtension( + name="dmuon._fast_clip_cuda", + sources=["dmuon/csrc/fast_clip_kernel.cu"], + extra_compile_args={ + "cxx": ["-O3"], + "nvcc": ["-O3"], + }, + ) + ] + return ext_modules, {"build_ext": BuildExtension} + + +ext_modules, cmdclass = _extension_config() + +setup( + ext_modules=ext_modules, + cmdclass=cmdclass, +) diff --git a/tests/distributed/test_fast_clip_distributed.py b/tests/distributed/test_fast_clip_distributed.py new file mode 100644 index 0000000..d4818f4 --- /dev/null +++ b/tests/distributed/test_fast_clip_distributed.py @@ -0,0 +1,78 @@ +"""Distributed liveness + correctness check for segmented fast clip. + +Reproduces the scenario that used to hang: a non-finite value in a +``reduce=False`` bucket on only one rank. With the three-stage design the +collective count is identical on every rank, so the clip must complete and +produce the correct all-reduced per-bucket norms. + +Usage (CPU, no GPU needed):: + + torchrun --nproc_per_node=2 tests/distributed/test_fast_clip_distributed.py + +Exit code is non-zero on any assertion failure; a hang (the old bug) shows up +as a torchrun timeout. +""" + +import os +import sys + +sys.path.insert( + 0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +) + +import math + +import torch +import torch.distributed as dist + +from dmuon.fast_clip import GradClipBucket, clip_grad_norm_buckets_ + + +def run() -> None: + dist.init_process_group(backend="gloo") + rank = dist.get_rank() + world_size = dist.get_world_size() + + # reduce=True buckets: each rank contributes a known local square; the + # segment-local norm after reduction is sqrt(sum_r local_sq_r). + muon = torch.tensor([float(rank + 1), 0.0]) # local sq = (rank+1)^2 + adamw = torch.tensor([0.0, float(rank + 2)]) # local sq = (rank+2)^2 + + # reduce=False bucket stays rank-local; rank 0 injects a non-finite value. + regular = torch.tensor([3.0, 4.0]) + if rank == 0: + regular = torch.tensor([float("inf"), 4.0]) + + result = clip_grad_norm_buckets_( + ( + GradClipBucket("regular", [regular], reduce=False), + GradClipBucket("muon", [muon], reduce=True), + GradClipBucket("adamw", [adamw], reduce=True), + ), + max_norm=1.0, + ) + stats = result.stats_by_name + + expected_muon = math.sqrt(sum((r + 1) ** 2 for r in range(world_size))) + expected_adamw = math.sqrt(sum((r + 2) ** 2 for r in range(world_size))) + + assert abs(stats["muon"].total_norm - expected_muon) < 1e-4, ( + rank, stats["muon"].total_norm, expected_muon, + ) + assert abs(stats["adamw"].total_norm - expected_adamw) < 1e-4, ( + rank, stats["adamw"].total_norm, expected_adamw, + ) + # regular is local: only rank 0 saw the inf. + assert stats["regular"].found_inf is (rank == 0), (rank, stats["regular"].found_inf) + + dist.barrier() + if rank == 0: + print( + f"PASS: fast-clip distributed (world_size={world_size}) — " + f"muon={expected_muon:.4f} adamw={expected_adamw:.4f}, no hang" + ) + dist.destroy_process_group() + + +if __name__ == "__main__": + run() diff --git a/tests/unit/test_fast_clip.py b/tests/unit/test_fast_clip.py new file mode 100644 index 0000000..b88b8bd --- /dev/null +++ b/tests/unit/test_fast_clip.py @@ -0,0 +1,244 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +import torch + +import dmuon.fast_clip as fast_clip +from dmuon.fast_clip import ( + GradClipBucket, + _DEFAULT_CHUNK_SIZE, + _resolve_chunk_size, + clip_grad_norm_buckets_, + try_clip_optimizer_grad_norm_buckets_, +) + + +def _bucket_stats(result): + return {stats.name: stats for stats in result.bucket_stats} + + +def test_fast_clip_chunk_size_defaults_and_env_override(monkeypatch) -> None: + monkeypatch.delenv("DMUON_FAST_CLIP_CHUNK_SIZE", raising=False) + assert _resolve_chunk_size(None) == _DEFAULT_CHUNK_SIZE + assert _DEFAULT_CHUNK_SIZE == 262144 + + monkeypatch.setenv("DMUON_FAST_CLIP_CHUNK_SIZE", "131072") + assert _resolve_chunk_size(None) == 131072 + + monkeypatch.setenv("DMUON_FAST_CLIP_CHUNK_SIZE", "invalid") + assert _resolve_chunk_size(None) == _DEFAULT_CHUNK_SIZE + + assert _resolve_chunk_size(4096) == 4096 + assert _resolve_chunk_size(0) == 1 + + +def test_segmented_clip_cpu_preserves_per_bucket_semantics() -> None: + """On CPU the kernel never applies; each bucket clips segment-locally.""" + regular = torch.tensor([3.0, 4.0]) + muon = torch.tensor([0.0, 12.0]) + adamw = torch.tensor([8.0, 15.0]) + + result = clip_grad_norm_buckets_( + ( + GradClipBucket("regular", [regular], reduce=False), + GradClipBucket("muon", [muon], reduce=True), + GradClipBucket("adamw", [adamw], reduce=True), + ), + 6.0, + ) + + stats = _bucket_stats(result) + assert stats["regular"].total_norm == pytest.approx(5.0) + assert stats["muon"].total_norm == pytest.approx(12.0) + assert stats["adamw"].total_norm == pytest.approx(17.0) + assert stats["regular"].clip_coef == pytest.approx(1.0) + assert stats["muon"].clip_coef == pytest.approx(6.0 / (12.0 + 1e-6)) + assert stats["adamw"].clip_coef == pytest.approx(6.0 / (17.0 + 1e-6)) + assert torch.allclose(regular, torch.tensor([3.0, 4.0])) + assert torch.allclose(muon, torch.tensor([0.0, 12.0]) * stats["muon"].clip_coef) + assert torch.allclose(adamw, torch.tensor([8.0, 15.0]) * stats["adamw"].clip_coef) + assert result.fastpath is False + + +def test_collective_count_is_invariant_to_local_conditions(monkeypatch) -> None: + """Regression for the distributed-hang bug. + + A non-finite value in a ``reduce=False`` bucket (a rank-local condition) + must NOT change how many collectives are issued. The clip must emit + exactly one all_reduce, of shape ``[num_buckets]``, no matter what — that + is what keeps ranks in lockstep. + """ + calls: list[torch.Tensor] = [] + + def fake_all_reduce(tensor, op=None, group=None): # single-rank identity + calls.append(tensor.clone()) + return tensor + + monkeypatch.setattr(fast_clip.dist, "is_available", lambda: True) + monkeypatch.setattr(fast_clip.dist, "is_initialized", lambda: True) + monkeypatch.setattr(fast_clip.dist, "all_reduce", fake_all_reduce) + + regular = torch.tensor([float("inf"), 1.0]) # local non-finite, reduce=False + muon = torch.tensor([3.0, 4.0]) + adamw = torch.tensor([0.0, 5.0]) + result = clip_grad_norm_buckets_( + ( + GradClipBucket("regular", [regular], reduce=False), + GradClipBucket("muon", [muon], reduce=True), + GradClipBucket("adamw", [adamw], reduce=True), + ), + 1.0, + ) + + assert len(calls) == 1 # exactly one fused collective + assert calls[0].shape == (3,) # one entry per bucket, not per-bucket scalars + stats = _bucket_stats(result) + assert stats["regular"].found_inf is True # reported, not acted on via control flow + assert stats["muon"].found_inf is False + + +def test_noncontiguous_tensor_is_clipped_not_skipped() -> None: + """A non-contiguous grad used to abort the whole call; now it clips fine.""" + base = torch.randn(4, 4) + grad = base.t() # non-contiguous view + assert not grad.is_contiguous() + reference = grad.clone().contiguous() + + result = clip_grad_norm_buckets_( + (GradClipBucket("muon", [grad], reduce=True),), 0.5 + ) + + norm = reference.float().pow(2).sum().sqrt() + coef = min(0.5 / (float(norm) + 1e-6), 1.0) + assert torch.allclose(grad, reference * coef, atol=1e-5) + assert result.stats_by_name["muon"].total_norm == pytest.approx(float(norm), rel=1e-5) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cuda_fast_clip_matches_python_reference(monkeypatch) -> None: + device = torch.device("cuda") + max_norm = 9.0 + source_buckets = ( + ( + "regular", + [ + torch.linspace(-2.0, 2.0, 2048, device=device, dtype=torch.float32), + torch.linspace(-1.0, 1.0, 513, device=device, dtype=torch.bfloat16), + ], + False, + ), + ( + "muon", + [ + torch.randn(4097, device=device, dtype=torch.float32), + torch.randn(123, device=device, dtype=torch.float16), + ], + True, + ), + ( + "adamw", + [ + torch.randn(8193, device=device, dtype=torch.bfloat16), + torch.randn(17, device=device, dtype=torch.float32), + ], + True, + ), + ) + + def _clone(buckets): + return tuple( + GradClipBucket(name, [g.clone() for g in grads], reduce=reduce) + for name, grads, reduce in buckets + ) + + fast_buckets = _clone(source_buckets) + fast = clip_grad_norm_buckets_(fast_buckets, max_norm) + if not fast.fastpath: + pytest.skip("DMuon fast clip CUDA extension is unavailable") + + # Force the pure-torch path for the reference by disabling the extension. + ref_buckets = _clone(source_buckets) + monkeypatch.setattr(fast_clip, "_load_fast_clip_extension", lambda: None) + ref = clip_grad_norm_buckets_(ref_buckets, max_norm) + assert ref.fastpath is False + + fast_stats = _bucket_stats(fast) + ref_stats = _bucket_stats(ref) + for name in fast_stats: + assert fast_stats[name].total_norm == pytest.approx( + ref_stats[name].total_norm, rel=2e-3, abs=2e-3 + ) + assert fast_stats[name].clip_coef == pytest.approx( + ref_stats[name].clip_coef, rel=2e-3, abs=2e-3 + ) + for fast_bucket, ref_bucket in zip(fast_buckets, ref_buckets): + for fast_grad, ref_grad in zip(fast_bucket.grads, ref_bucket.grads): + assert torch.allclose(fast_grad.float(), ref_grad.float(), rtol=3e-3, atol=3e-3) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cuda_and_python_agree_on_nonfinite(monkeypatch) -> None: + """NaN/inf must produce identical grads on the CUDA and torch paths.""" + device = torch.device("cuda") + values = [1.0, float("nan"), 3.0, float("inf")] + + fast_grad = torch.tensor(values, device=device) + fast = clip_grad_norm_buckets_((GradClipBucket("muon", [fast_grad], reduce=True),), 2.0) + if not fast.fastpath: + pytest.skip("DMuon fast clip CUDA extension is unavailable") + + ref_grad = torch.tensor(values, device=device) + monkeypatch.setattr(fast_clip, "_load_fast_clip_extension", lambda: None) + clip_grad_norm_buckets_((GradClipBucket("muon", [ref_grad], reduce=True),), 2.0) + + assert torch.equal(fast_grad.isnan(), ref_grad.isnan()) + finite = ~fast_grad.isnan() & ~fast_grad.isinf() + assert torch.allclose(fast_grad[finite], ref_grad[finite], rtol=3e-3, atol=3e-3) + + +@dataclass +class _FakeDP: + grad: torch.Tensor | None + name: str = "proj" + + def __post_init__(self) -> None: + self.param_name = self.name + self.is_dtensor = False + self.tp_group = None + self.is_tp_owner = False + self._reduced_grad = self.grad + + +class _FakeOptimizer: + def __init__(self, grad: torch.Tensor) -> None: + self._dedicated_params = [_FakeDP(grad)] + self._dp_to_muon_group_idx = {id(self._dedicated_params[0]): 0} + self._ensure_calls = 0 + self._last_muon_grad_clip_stats = None + + def _ensure_grads_ready(self, *, coalesce_wait: bool = False) -> None: + self._ensure_calls += 1 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_optimizer_fast_clip_sets_muon_stats_when_available() -> None: + muon_grad = torch.tensor([6.0, 8.0], device="cuda") + regular_grad = torch.tensor([3.0, 4.0], device="cuda") + adamw_grad = torch.tensor([0.0, 12.0], device="cuda") + opt = _FakeOptimizer(muon_grad) + + result = try_clip_optimizer_grad_norm_buckets_( + opt, + regular_grads=[regular_grad], + adamw_grads=[adamw_grad], + max_norm=5.0, + ) + if result is None: + pytest.skip("DMuon fast clip CUDA extension is unavailable") + + assert opt._ensure_calls == 1 + assert opt._last_muon_grad_clip_stats is not None + assert opt._last_muon_grad_clip_stats.total_norm == pytest.approx(10.0) + assert opt._last_muon_grad_clip_stats.strategy == "segmented_cuda_fastpath" diff --git a/tests/unit/test_hsdp_owner_only.py b/tests/unit/test_hsdp_owner_only.py index 073776e..77c4ad8 100644 --- a/tests/unit/test_hsdp_owner_only.py +++ b/tests/unit/test_hsdp_owner_only.py @@ -36,12 +36,16 @@ class _StubGroup: """Minimal ProcessGroup stub: just advertises a rank within the group.""" - def __init__(self, rank: int): + def __init__(self, rank: int, size: int = 4): self._rank = rank + self._size = size def rank(self) -> int: return self._rank + def size(self) -> int: + return self._size + class _StubDist: @staticmethod @@ -57,6 +61,7 @@ def _make_dp( owner_shard: int, owner_replicate: int, shape=(8, 4), + route_hint: Optional[str] = None, ) -> DedicatedParam: """Build a DedicatedParam on CPU with stubbed groups. @@ -83,6 +88,7 @@ def _make_dp( device=torch.device("cpu"), compute_dtype=None, replicate_group=replicate_group, + route_hint=route_hint, ) finally: param_module.dist = orig_dist @@ -171,6 +177,21 @@ def test_owned_data_none_on_non_owner_in_shard_only_mode(): assert dp_owner._owned_data is not None +def test_sharded_adamw_skips_full_owner_copy(): + """Sharded AdamW owns a rank-local shard, not a full fp32 owner tensor.""" + dp = _make_dp( + shard_rank=3, + replicate_rank=None, + owner_shard=3, + owner_replicate=0, + shape=(8, 4), + route_hint="sharded_adamw", + ) + assert dp._owned_data is None + assert dp._sharded_adamw_data is not None + assert dp._sharded_adamw_data.numel() == 8 + + # --- Cached global rank for Stage-2 reduce ---------------------------------