From 0f13451ad4dd63ebd0edc6e091936fde41969695 Mon Sep 17 00:00:00 2001 From: Maksim Abraham Date: Wed, 18 Mar 2026 17:05:00 -0700 Subject: [PATCH 01/41] Add DR-GRPO loss (#24) Squashed port of internal PR #24 (mabraham/grpo): - Add dr-grpo loss - add items back - rename metrics log - Add tests and rename - Adjust tests - import assert --- src/xorl/ops/loss/__init__.py | 4 + src/xorl/ops/loss/grpo_loss.py | 233 ++++++++++++++++++++ tests/ops/loss/__init__.py | 0 tests/ops/loss/conftest.py | 21 ++ tests/ops/loss/test_drgrpo_loss.py | 340 +++++++++++++++++++++++++++++ 5 files changed, 598 insertions(+) create mode 100644 src/xorl/ops/loss/grpo_loss.py create mode 100644 tests/ops/loss/__init__.py create mode 100644 tests/ops/loss/conftest.py create mode 100644 tests/ops/loss/test_drgrpo_loss.py diff --git a/src/xorl/ops/loss/__init__.py b/src/xorl/ops/loss/__init__.py index 6edb14c2..18a8b620 100644 --- a/src/xorl/ops/loss/__init__.py +++ b/src/xorl/ops/loss/__init__.py @@ -5,11 +5,13 @@ - causallm_loss_function: Standard causal language modeling loss - importance_sampling_loss_function: Importance sampling loss for GRPO/RL - policy_loss_function: PPO-style policy loss with TIS correction +- drgrpo_loss_function: DR-GRPO loss with PPO clipping and KL penalty """ from typing import Callable, Dict from xorl.ops.loss.causallm_loss import causallm_loss_function +from xorl.ops.loss.grpo_loss import drgrpo_loss_function from xorl.ops.loss.importance_sampling_loss import importance_sampling_loss_function from xorl.ops.loss.loss_output import LossOutput from xorl.ops.loss.policy_loss import policy_loss_function @@ -24,6 +26,7 @@ "cross_entropy": causallm_loss_function, # alias "importance_sampling": importance_sampling_loss_function, "policy_loss": policy_loss_function, + "drgrpo": drgrpo_loss_function, } @@ -45,6 +48,7 @@ def register_loss_function(name: str, fn: Callable) -> None: "get_loss_function", "register_loss_function", "causallm_loss_function", + "drgrpo_loss_function", "importance_sampling_loss_function", "policy_loss_function", "vocab_parallel_cross_entropy", diff --git a/src/xorl/ops/loss/grpo_loss.py b/src/xorl/ops/loss/grpo_loss.py new file mode 100644 index 00000000..bca416e5 --- /dev/null +++ b/src/xorl/ops/loss/grpo_loss.py @@ -0,0 +1,233 @@ +""" +DR-GRPO: "Done Right" GRPO Loss for RL Training. + +Reference: Liu et al., "Understanding R1-Zero-Like Training" (2025). +https://arxiv.org/abs/2503.20783 +""" + +from typing import List, Literal, Optional, Tuple + +import torch +import torch.distributed as dist + +from xorl.ops.loss.loss_output import LossOutput +from xorl.ops.loss.per_token_ce import compute_per_token_ce + + +AggType = Literal["token_mean", "fixed_horizon", "sequence_mean"] +KLType = Literal["k1", "k2", "k3"] +RatioType = Literal["token", "sequence"] + + +def masked_mean( + values: torch.Tensor, + mask: torch.Tensor, + loss_scale: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Masked mean: sum(values * mask) / divisor.""" + masked_sum = (values * mask).sum() + if loss_scale is not None: + divisor = loss_scale.clamp(min=1.0) + else: + divisor = mask.sum().clamp(min=1.0) + return masked_sum / divisor + + +def compute_ratio( + logprobs: torch.Tensor, + generator_logprobs: torch.Tensor, + mask: torch.Tensor, + ratio_type: RatioType = "token", +) -> Tuple[torch.Tensor, torch.Tensor, List[Tuple[str, torch.Tensor]]]: + """Importance sampling ratio r = π_θ/π_old. + + token: r_t = exp(logprobs_t - generator_logprobs_t) + sequence: r_seq = exp(mean_t[logprobs - generator_logprobs]), uses reparameterization. + """ + if ratio_type == "token": + log_ratio = logprobs - generator_logprobs.detach() + ratio = torch.exp(log_ratio) + elif ratio_type == "sequence": + token_log_ratio = logprobs - generator_logprobs.detach() + seq_lengths = mask.sum(dim=-1).clamp(min=1) + seq_log_ratio = (token_log_ratio * mask).sum(dim=-1) / seq_lengths + + # Reparameterization: forward uses seq ratio, backward uses token grads + log_ratio = logprobs - logprobs.detach() + seq_log_ratio.detach().unsqueeze(-1) + ratio = torch.exp(log_ratio) + else: + raise ValueError(f"Unknown ratio_type: {ratio_type}") + + with torch.no_grad(): + metrics = [ + ("loss/ratio/mean", masked_mean(ratio, mask)), + ("loss/kl_policy/mean", masked_mean(-log_ratio, mask)), + ] + + return ratio, log_ratio, metrics + + +def compute_kl( + policy_logprobs: torch.Tensor, + ref_logprobs: torch.Tensor, + mask: torch.Tensor, + kl_type: KLType = "k3", +) -> Tuple[torch.Tensor, List[Tuple[str, torch.Tensor]]]: + """KL divergence using Schulman's estimators (k1, k2, k3).""" + log_ratio = policy_logprobs - ref_logprobs.detach() + + if kl_type == "k1": + kl = log_ratio + elif kl_type == "k2": + kl = 0.5 * log_ratio.square() + elif kl_type == "k3": + neg_log_ratio = torch.clamp(-log_ratio, min=-10.0, max=10.0) + ratio = torch.exp(neg_log_ratio) + kl = ratio - neg_log_ratio - 1 + else: + raise ValueError(f"Unknown kl_type: {kl_type}") + + with torch.no_grad(): + metrics = [("loss/kl_ref/mean", masked_mean(kl, mask))] + + return kl, metrics + + +def pg_ppo_clip( + ratio: torch.Tensor, + advantages: torch.Tensor, + mask: torch.Tensor, + clip_low: float = 0.2, + clip_high: float = 0.2, +) -> Tuple[torch.Tensor, List[Tuple[str, torch.Tensor]]]: + """PPO clipped surrogate: L = max(-r*A, -clip(r, 1-ε_low, 1+ε_high)*A).""" + clipped_ratio = torch.clamp(ratio, 1 - clip_low, 1 + clip_high) + unclipped_loss = -ratio * advantages + clipped_loss = -clipped_ratio * advantages + pg_loss = torch.maximum(unclipped_loss, clipped_loss) + + with torch.no_grad(): + mask_bool = mask.bool() + clipped_high = (ratio > 1 + clip_high) & mask_bool + clipped_low = (ratio < 1 - clip_low) & mask_bool + pos_adv = advantages > 0 + neg_adv = advantages < 0 + + metrics = [ + ("loss/clip/clipped_ratio/mean", masked_mean(clipped_ratio, mask)), + ("loss/clip/high_fraction", masked_mean((clipped_high & pos_adv).float(), mask)), + ("loss/clip/low_fraction", masked_mean((clipped_low & neg_adv).float(), mask)), + ] + + return pg_loss, metrics + + +def aggregate( + per_token_loss: torch.Tensor, + mask: torch.Tensor, + agg_type: AggType = "token_mean", + loss_scale: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, List[Tuple[str, torch.Tensor]]]: + """Aggregate per-token loss: token_mean, fixed_horizon, or sequence_mean.""" + if agg_type == "token_mean": + loss = masked_mean(per_token_loss, mask, loss_scale) + elif agg_type == "fixed_horizon": + loss = (per_token_loss * mask).sum() / max(mask.numel(), 1) + elif agg_type == "sequence_mean": + seq_lengths = mask.sum(dim=-1).clamp(min=1.0) + seq_means = (per_token_loss * mask).sum(dim=-1) / seq_lengths + loss = seq_means.sum() / max(seq_means.numel(), 1) + else: + raise ValueError(f"Unknown agg_type: {agg_type}") + + with torch.no_grad(): + metrics = [("loss/aggregate/active_fraction", mask.mean())] + + return loss, metrics + + +def drgrpo_loss_function( + hidden_states: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + old_logprobs: torch.Tensor, + advantages: torch.Tensor, + ref_logprobs: Optional[torch.Tensor] = None, + ignore_index: int = -100, + clip_low: float = 0.2, + clip_high: float = 0.28, + beta: float = 0.1, + agg_type: AggType = "fixed_horizon", + ratio_type: RatioType = "token", + kl_type: KLType = "k3", + ce_mode: str = "compiled", + num_chunks: int = 8, + tp_group: Optional[dist.ProcessGroup] = None, + lm_head_fp32: bool = False, + loss_scale: Optional[torch.Tensor] = None, +) -> LossOutput: + """DR-GRPO loss for RL training. + + Per-token: L_t = max(-r*A, -clip(r, 1-ε, 1+ε)*A) + β*KL + Aggregated: Depends on agg_type (default: fixed_horizon) + + Args: + hidden_states: (B, S, H) model hidden states. + weight: (V, H) or (V/tp, H) LM head weight. + labels: (B, S) target token IDs, already next-token aligned. + old_logprobs: (B, S) log probs from generation policy. + advantages: (B, S) per-token advantages. + ref_logprobs: (B, S) reference model log probs for KL (required if beta > 0). + ignore_index: Token ID to ignore (default: -100). + clip_low: Lower clip bound (default: 0.2). + clip_high: Upper clip bound (default: 0.28). + beta: KL penalty coefficient (default: 0.1). + agg_type: Aggregation type (default: "fixed_horizon"). + ratio_type: Ratio type: "token" or "sequence" (default: "token"). + kl_type: KL estimator: "k1", "k2", "k3" (default: "k3"). + ce_mode: Cross-entropy mode: "compiled" or "eager". + num_chunks: Chunks for compiled mode. + tp_group: TP process group for vocab-parallel CE. + lm_head_fp32: Compute LM head in FP32. + loss_scale: For distributed token_mean aggregation. + + Returns: + LossOutput with loss, per_token_logprobs, per_token_loss, and metrics. + """ + if beta > 0 and ref_logprobs is None: + raise ValueError("ref_logprobs required when beta > 0") + + B, S = labels.shape + H = hidden_states.size(-1) + + labels_flat = labels.reshape(-1) + hidden_states_flat = hidden_states.reshape(-1, H) + + per_token_ce = compute_per_token_ce( + hidden_states_flat, weight, labels_flat, ignore_index, + ce_mode, num_chunks, tp_group=tp_group, lm_head_fp32=lm_head_fp32, + ) + logprobs = -per_token_ce.view(B, S) + + loss_mask = (labels != ignore_index).float() + + ratio, _, ratio_m = compute_ratio(logprobs, old_logprobs, loss_mask, ratio_type) + + pg_loss, clip_m = pg_ppo_clip(ratio, advantages, loss_mask, clip_low, clip_high) + + kl_m: List[Tuple[str, torch.Tensor]] = [] + if beta > 0: + kl, kl_m = compute_kl(logprobs, ref_logprobs, loss_mask, kl_type) + pg_loss = pg_loss + beta * kl + + loss, agg_m = aggregate(pg_loss, loss_mask, agg_type, loss_scale) + + all_metrics = ratio_m + clip_m + kl_m + agg_m + metrics = {k: v.item() for k, v in all_metrics} + + return LossOutput( + loss=loss, + per_token_logprobs=logprobs.detach(), + per_token_loss=pg_loss.detach(), + metrics=metrics, + ) diff --git a/tests/ops/loss/__init__.py b/tests/ops/loss/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/ops/loss/conftest.py b/tests/ops/loss/conftest.py new file mode 100644 index 00000000..4beb1d2d --- /dev/null +++ b/tests/ops/loss/conftest.py @@ -0,0 +1,21 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +# +# Ported from torchforge tests. + +import torch + + +def assert_close(actual, expected, atol=1e-4, rtol=1e-4): + """Assert two tensors are close within tolerance.""" + torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) + + +def get_metric(metrics: dict, key: str): + """Get a metric value from the metrics dict.""" + if key not in metrics: + raise KeyError(f"Metric '{key}' not found. Available: {list(metrics.keys())}") + return metrics[key] diff --git a/tests/ops/loss/test_drgrpo_loss.py b/tests/ops/loss/test_drgrpo_loss.py new file mode 100644 index 00000000..89109853 --- /dev/null +++ b/tests/ops/loss/test_drgrpo_loss.py @@ -0,0 +1,340 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +# +# Ported from torchforge tests for GRPOLoss. + +import pytest +import torch + +from xorl.ops.loss import drgrpo_loss_function + +from tests.ops.loss.conftest import assert_close + + +@pytest.fixture +def inputs(): + """Test inputs for DRGRPO loss. + + Note: xorl's drgrpo_loss_function takes hidden_states and weight instead of logits. + We create hidden_states and weight such that hidden_states @ weight.T produces + logits with similar scale to torch.randn(B, S, V). + """ + torch.manual_seed(42) + B, S, V = 2, 4, 10 + H = 16 # Hidden dimension + + # Create hidden_states and weight matrix + # Scale by 1/sqrt(H) so that logits = hidden_states @ weight.T has variance ~1 + # (similar to torch.randn logits in torchforge tests) + hidden_states = torch.randn(B, S, H) / (H ** 0.5) + weight = torch.randn(V, H) + + # Compute effective logits for reference (should have variance ~1) + logits = hidden_states @ weight.T + + # Target IDs (called 'labels' in xorl API) + labels = torch.randint(0, V, (B, S)) + + # Seq 0: mild divergence, Seq 1: high divergence (triggers clipping) + old_logprobs = torch.tensor( + [ + [-2.0, -2.1, -1.9, -2.0], + [-6.0, -1.0, -5.0, -0.5], + ] + ) + ref_logprobs = torch.randn(B, S) * 0.5 - 2.0 + advantages = torch.randn(B, S) + + # Interleaved mask: use ignore_index to mark non-loss positions + # loss_mask in torchforge: [[1, 0, 1, 0], [1, 1, 0, 0]] + # For xorl, we use ignore_index=-100 in labels for masked positions + ignore_index = -100 + labels_with_mask = labels.clone() + mask_pattern = torch.tensor([[1, 0, 1, 0], [1, 1, 0, 0]], dtype=torch.bool) + labels_with_mask[~mask_pattern] = ignore_index + + return { + "B": B, + "S": S, + "V": V, + "H": H, + "hidden_states": hidden_states, + "weight": weight, + "logits": logits, + "labels": labels, + "labels_with_mask": labels_with_mask, + "old_logprobs": old_logprobs, + "ref_logprobs": ref_logprobs, + "advantages": advantages, + "ignore_index": ignore_index, + "mask_pattern": mask_pattern, + } + + +class TestDRGRPOLoss: + """Tests for drgrpo_loss_function. + + Note: test_forward and test_backward contain regression tests with exact expected + values. If the implementation changes intentionally, update the expected values by + running the tests with pytest -v and recording the actual values: + + pytest tests/ops/loss/test_drgrpo_loss.py -v -k "test_forward or test_backward" + + Then update the assert_close(...) calls with the new values. + """ + + def test_forward(self, inputs): + """Forward pass produces expected loss value (regression test).""" + d = inputs + hidden_states = d["hidden_states"].clone().requires_grad_(True) + + output = drgrpo_loss_function( + hidden_states=hidden_states, + weight=d["weight"], + labels=d["labels_with_mask"], + old_logprobs=d["old_logprobs"], + advantages=d["advantages"], + ref_logprobs=d["ref_logprobs"], + ignore_index=d["ignore_index"], + clip_low=0.2, + clip_high=0.2, + beta=0.1, + ) + + assert output.loss.isfinite() + assert output.loss.shape == () + # Regression test: expected value computed with seed=42 fixture inputs + assert_close(output.loss, torch.tensor(0.363678)) + + def test_backward(self, inputs): + """Backward pass produces expected gradient norm (regression test).""" + d = inputs + hidden_states = d["hidden_states"].clone().requires_grad_(True) + + output = drgrpo_loss_function( + hidden_states=hidden_states, + weight=d["weight"], + labels=d["labels_with_mask"], + old_logprobs=d["old_logprobs"], + advantages=d["advantages"], + ref_logprobs=d["ref_logprobs"], + ignore_index=d["ignore_index"], + clip_low=0.2, + clip_high=0.2, + beta=0.1, + ) + + output.loss.backward() + assert hidden_states.grad is not None + assert hidden_states.grad.isfinite().all() + # Regression test: expected value computed with seed=42 fixture inputs + assert_close(hidden_states.grad.norm(), torch.tensor(1.514308)) + + def test_zero_advantages(self, inputs): + """Zero advantages produce finite (near-zero) loss.""" + d = inputs + advantages = torch.zeros_like(d["advantages"]) + + output = drgrpo_loss_function( + hidden_states=d["hidden_states"], + weight=d["weight"], + labels=d["labels_with_mask"], + old_logprobs=d["old_logprobs"], + advantages=advantages, + ignore_index=d["ignore_index"], + beta=0.0, + ) + + assert output.loss.isfinite() + # With zero advantages, policy gradient loss should be zero + assert output.loss.abs() < 1e-5 + + def test_all_ignored_labels(self, inputs): + """Loss should be finite (zero) when all labels are ignored (no trainable tokens).""" + d = inputs + # Set all labels to ignore_index + all_ignored = torch.full_like(d["labels"], d["ignore_index"]) + + output = drgrpo_loss_function( + hidden_states=d["hidden_states"], + weight=d["weight"], + labels=all_ignored, + old_logprobs=d["old_logprobs"], + advantages=d["advantages"], + ignore_index=d["ignore_index"], + beta=0.0, + ) + + assert output.loss.isfinite() + assert output.loss == 0.0 + + def test_empty_sequence(self): + """Loss should be zero when sequence length is 0.""" + B, V, H = 2, 10, 16 + hidden_states = torch.empty(B, 0, H) + weight = torch.randn(V, H) + labels = torch.empty(B, 0, dtype=torch.long) + advantages = torch.empty(B, 0) + old_logprobs = torch.empty(B, 0) + + output = drgrpo_loss_function( + hidden_states=hidden_states, + weight=weight, + labels=labels, + old_logprobs=old_logprobs, + advantages=advantages, + beta=0.0, + ) + + assert output.loss.isfinite() + assert output.loss == 0.0 + + def test_requires_ref_logprobs_when_beta_positive(self, inputs): + """ValueError raised when beta > 0 but ref_logprobs is None.""" + d = inputs + + with pytest.raises(ValueError, match="ref_logprobs required"): + drgrpo_loss_function( + hidden_states=d["hidden_states"], + weight=d["weight"], + labels=d["labels_with_mask"], + old_logprobs=d["old_logprobs"], + advantages=d["advantages"], + ref_logprobs=None, + ignore_index=d["ignore_index"], + beta=0.1, + ) + + def test_positive_advantages_encourage_high_prob(self, inputs): + """With positive advantages, higher target probability yields lower loss.""" + d = inputs + B, S, V, H = d["B"], d["S"], d["V"], d["H"] + + # Create two scenarios with same structure but different logit magnitudes + torch.manual_seed(123) + labels = torch.randint(0, V, (B, S)) + old_logprobs = torch.zeros(B, S) + positive_advantages = torch.ones(B, S) * 2.0 + + # High probability scenario: hidden states that produce high logits for targets + hidden_high = torch.randn(B, S, H) + weight_high = torch.randn(V, H) + # Bias the logits toward target tokens + for b in range(B): + for s in range(S): + weight_high[labels[b, s]] += hidden_high[b, s] * 5.0 + + # Low probability scenario + hidden_low = torch.randn(B, S, H) + weight_low = torch.randn(V, H) + # Bias away from target tokens + for b in range(B): + for s in range(S): + weight_low[labels[b, s]] -= hidden_low[b, s] * 5.0 + + loss_high = drgrpo_loss_function( + hidden_states=hidden_high, + weight=weight_high, + labels=labels, + old_logprobs=old_logprobs, + advantages=positive_advantages, + beta=0.0, + ) + + loss_low = drgrpo_loss_function( + hidden_states=hidden_low, + weight=weight_low, + labels=labels, + old_logprobs=old_logprobs, + advantages=positive_advantages, + beta=0.0, + ) + + # Higher probability should yield lower (more negative) loss + assert loss_high.loss < loss_low.loss + + def test_kl_penalty_affects_loss(self, inputs): + """KL penalty modifies loss when beta > 0.""" + d = inputs + + loss_no_kl = drgrpo_loss_function( + hidden_states=d["hidden_states"], + weight=d["weight"], + labels=d["labels_with_mask"], + old_logprobs=d["old_logprobs"], + advantages=d["advantages"], + ignore_index=d["ignore_index"], + beta=0.0, + ) + + loss_with_kl = drgrpo_loss_function( + hidden_states=d["hidden_states"], + weight=d["weight"], + labels=d["labels_with_mask"], + old_logprobs=d["old_logprobs"], + advantages=d["advantages"], + ref_logprobs=d["ref_logprobs"], + ignore_index=d["ignore_index"], + beta=0.1, + ) + + assert loss_no_kl.loss != loss_with_kl.loss + assert loss_with_kl.loss.isfinite() + # KL metrics should be present when beta > 0 + assert "loss/kl_ref/mean" in loss_with_kl.metrics + + def test_metrics_present(self, inputs): + """Output includes expected metrics.""" + d = inputs + + output = drgrpo_loss_function( + hidden_states=d["hidden_states"], + weight=d["weight"], + labels=d["labels_with_mask"], + old_logprobs=d["old_logprobs"], + advantages=d["advantages"], + ref_logprobs=d["ref_logprobs"], + ignore_index=d["ignore_index"], + clip_low=0.2, + clip_high=0.2, + beta=0.1, + ) + + expected_keys = [ + "loss/ratio/mean", + "loss/kl_policy/mean", + "loss/clip/clipped_ratio/mean", + "loss/clip/high_fraction", + "loss/clip/low_fraction", + "loss/kl_ref/mean", + "loss/aggregate/active_fraction", + ] + + for key in expected_keys: + assert key in output.metrics, f"Missing metric: {key}" + + def test_aggregation_types(self, inputs): + """Different aggregation types produce different losses.""" + d = inputs + + losses = {} + for agg_type in ["token_mean", "fixed_horizon", "sequence_mean"]: + output = drgrpo_loss_function( + hidden_states=d["hidden_states"], + weight=d["weight"], + labels=d["labels_with_mask"], + old_logprobs=d["old_logprobs"], + advantages=d["advantages"], + ignore_index=d["ignore_index"], + agg_type=agg_type, + beta=0.0, + ) + losses[agg_type] = output.loss.item() + assert output.loss.isfinite() + + # At least some aggregation types should produce different values + unique_losses = set(round(v, 6) for v in losses.values()) + assert len(unique_losses) >= 2, "Aggregation types should produce different losses" From 9b9d8b65a7c9d004adf07bb520671c50c82726eb Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Fri, 27 Mar 2026 04:26:10 -0400 Subject: [PATCH 02/41] Globally reduce load balancing loss across DP ranks (#40) All-reduce tokens_per_expert across the data-parallel group before computing the Switch Transformer auxiliary loss so the loss reflects the global routing distribution, not just the local micro-batch. Closes #39 (cherry picked from commit dd602d1ae4702dfaf37b14922cbf91ebfcbc0e07) --- src/xorl/cli/direct_train.py | 138 ++++++++++-------- src/xorl/models/layers/moe/__init__.py | 15 +- src/xorl/models/layers/moe/aux_loss.py | 105 +++++++++++++ .../qwen3_moe/modeling_qwen3_moe.py | 120 +++------------ src/xorl/trainers/trainer.py | 134 ++++++++++------- 5 files changed, 295 insertions(+), 217 deletions(-) create mode 100644 src/xorl/models/layers/moe/aux_loss.py diff --git a/src/xorl/cli/direct_train.py b/src/xorl/cli/direct_train.py index bd7ffb05..36766243 100644 --- a/src/xorl/cli/direct_train.py +++ b/src/xorl/cli/direct_train.py @@ -1,19 +1,19 @@ import os + # Must be set before importing torch / initializing CUDA so the # allocator picks up the setting on first use. os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import json import time -from dataclasses import asdict, dataclass, field -from functools import partial +from dataclasses import asdict from typing import Any, Dict, List -import torch import torch.distributed as dist from tqdm import trange +from xorl.arguments import Arguments, parse_args, save_args from xorl.checkpoint import build_checkpointer from xorl.data.constants import IGNORE_INDEX from xorl.data.data_loader import DataLoaderBuilder @@ -24,8 +24,8 @@ from xorl.distributed.sync_padding import synchronize_micro_batch_padding from xorl.distributed.torch_parallelize import build_parallelize_model from xorl.models import build_foundation_model, build_tokenizer, save_model_assets, save_model_weights +from xorl.models.layers.moe.aux_loss import global_load_balancing_loss_func from xorl.models.module_utils import compute_loss -from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import load_balancing_loss_func from xorl.optim import build_lr_scheduler, build_optimizer from xorl.trainers.training_utils import ( clip_gradients, @@ -34,16 +34,12 @@ sync_sp_gradients, ) from xorl.utils import helper - from xorl.utils.device import ( get_device_type, get_nccl_backend, get_torch_device, synchronize, ) - -from xorl.arguments import Arguments, DataArguments, ModelArguments, TrainingArguments, parse_args, save_args - from xorl.utils.dist_utils import all_reduce @@ -55,10 +51,10 @@ def main(): dist.init_process_group(backend=get_nccl_backend()) logger.info(f"Process rank: {args.train.global_rank}, world size: {args.train.world_size}") logger.info_rank0(json.dumps(asdict(args), indent=2)) - + get_torch_device().set_device(f"{get_device_type()}:{args.train.local_rank}") helper.set_seed(args.train.seed, args.train.enable_full_determinism) - + if args.train.local_rank == 0: helper.enable_third_party_logging() @@ -66,7 +62,7 @@ def main(): save_args(args, args.train.output_dir) Checkpointer = build_checkpointer(dist_backend=args.train.data_parallel_mode, ckpt_manager=args.train.ckpt_manager) - + init_parallel_state( dp_size=args.train.data_parallel_size, dp_replicate_size=args.train.data_parallel_replicate_size, @@ -86,6 +82,7 @@ def main(): ps = get_parallel_state() if ps.device_mesh is not None: import torch.distributed.tensor._random + torch.distributed.tensor._random.manual_seed(args.train.seed, ps.device_mesh) logger.info_rank0("Prepare data") @@ -105,7 +102,7 @@ def main(): seed=args.train.seed, pad_to_multiple_of=args.data.pad_to_multiple_of, ).build() - + # Calculate train steps from dataloader length train_steps_per_epoch = len(train_dataloader) total_train_steps = train_steps_per_epoch * args.train.num_train_epochs @@ -118,7 +115,6 @@ def main(): if save_epoch_steps: logger.info_rank0(f"Save every {args.train.save_epochs} epoch(s) = every {save_epoch_steps} steps") - logger.info_rank0("Prepare model") model = build_foundation_model( config_path=args.model.config_path, @@ -152,7 +148,8 @@ def main(): checkpoint_quant_format = None exclude_modules = set() if args.lora.enable_qlora: - from xorl.qlora import inject_qlora_into_model, detect_prequantized_nvfp4, detect_prequantized_block_fp8 + from xorl.qlora import detect_prequantized_block_fp8, detect_prequantized_nvfp4, inject_qlora_into_model + if detect_prequantized_nvfp4(args.model.model_path): is_prequantized = True checkpoint_quant_format = "nvfp4" @@ -163,16 +160,14 @@ def main(): logger.info_rank0("Detected pre-quantized block FP8 checkpoint") if args.lora.exclude_modules is not None: exclude_modules = set(args.lora.exclude_modules) - logger.info_rank0( - f"Using user-specified exclude_modules: {exclude_modules}" - ) + logger.info_rank0(f"Using user-specified exclude_modules: {exclude_modules}") elif is_prequantized: from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules + exclude_modules = get_prequantized_exclude_modules(args.model.model_path) if exclude_modules: logger.info_rank0( - f"Auto-detected {len(exclude_modules)} excluded modules " - f"from checkpoint config: {exclude_modules}" + f"Auto-detected {len(exclude_modules)} excluded modules from checkpoint config: {exclude_modules}" ) if is_prequantized and checkpoint_quant_format != args.lora.quant_format: logger.info_rank0( @@ -200,6 +195,7 @@ def main(): helper.print_device_mem_info("VRAM usage after QLoRA injection") elif args.lora.enable_lora: from xorl.lora.utils import inject_lora_into_model + inject_lora_into_model( model, r=args.lora.lora_rank, @@ -249,13 +245,15 @@ def main(): if args.lora.enable_qlora: if is_prequantized: from xorl.qlora import maybe_load_prequantized_qlora + logger.info("Starting pre-quantized NVFP4 weight loading...") helper.print_device_mem_info("VRAM before pre-quantized loading") maybe_load_prequantized_qlora(model, args.model.model_path) logger.info("Done pre-quantized weight loading, freezing non-LoRA params...") else: - from xorl.qlora import maybe_quantize_qlora, maybe_load_and_quantize_moe_qlora + from xorl.qlora import maybe_load_and_quantize_moe_qlora, maybe_quantize_qlora from xorl.qlora.utils import _deregister_qlora_weights_from_fsdp + logger.info("Starting maybe_quantize_qlora...") helper.print_device_mem_info("VRAM before QLoRA quantization") maybe_quantize_qlora(model) @@ -267,7 +265,8 @@ def main(): logger.info("Done MoE weight loading, deregistering packed weights...") # Deregister packed_weight_f32 from FSDP2 (prevent mixed-precision corruption) removed = _deregister_qlora_weights_from_fsdp( - model, param_names=("packed_weight_f32",), + model, + param_names=("packed_weight_f32",), ) torch.cuda.empty_cache() if removed > 0: @@ -305,6 +304,7 @@ def main(): if args.train.global_rank == 0: if args.train.use_wandb: import wandb + wandb.init( project=args.train.wandb_project, name=args.train.wandb_name, @@ -364,17 +364,21 @@ def main(): pp_context = {} # mutable container for per-step state used by pp_loss_fn if pp_enabled: import torch.nn.functional as F + from xorl.distributed.pipeline_parallel import build_pipeline_schedule @torch.compile def _pp_ce_loss(pred, labels, ntokens): """PP loss: sum reduction, normalized by global_valid_tokens.""" - return F.cross_entropy( - pred.flatten(0, 1).float(), - labels.flatten(0, 1), - ignore_index=IGNORE_INDEX, - reduction="sum", - ) / ntokens + return ( + F.cross_entropy( + pred.flatten(0, 1).float(), + labels.flatten(0, 1), + ignore_index=IGNORE_INDEX, + reduction="sum", + ) + / ntokens + ) def pp_loss_fn(pred, labels): return _pp_ce_loss(pred, labels, pp_context["global_valid_tokens"]) @@ -441,7 +445,8 @@ def pp_loss_fn(pred, labels): # compute global valid tokens across all ranks global_valid_tokens = count_valid_tokens( - micro_batches, group=ps.fsdp_group if pp_enabled else None, + micro_batches, + group=ps.fsdp_group if pp_enabled else None, ) optimizer.zero_grad() @@ -451,13 +456,15 @@ def pp_loss_fn(pred, labels): # pays the cheap addcmul cost per layer. if args.lora.enable_aqn: from xorl.qlora.modules.linear import prefetch_aqn_noise + prefetch_aqn_noise(model) # Routing replay stage switching for MoE checkpoint determinism. # Only needed with EP — without EP, expert compute has fixed output # shapes regardless of routing, so checkpoint recompute is safe. # See models/layers/moe/routing_replay.py for lifecycle docs. - from xorl.models.layers.moe.routing_replay import set_replay_stage, RoutingReplay + from xorl.models.layers.moe.routing_replay import RoutingReplay, set_replay_stage + use_routing_replay = ps.ep_size > 1 and args.train.moe_recomputed if pp_enabled: @@ -471,18 +478,15 @@ def pp_loss_fn(pred, labels): # Prepare input_ids and labels tensors for PP schedule # PP schedule expects full batch tensors, splits into microbatches internally device = get_device_type() - input_ids = torch.cat([ - mb["input_ids"].to(device, non_blocking=True) for mb in micro_batches - ], dim=0) - labels = torch.cat([ - mb["labels"].to(device, non_blocking=True) for mb in micro_batches - ], dim=0) + input_ids = torch.cat([mb["input_ids"].to(device, non_blocking=True) for mb in micro_batches], dim=0) + labels = torch.cat([mb["labels"].to(device, non_blocking=True) for mb in micro_batches], dim=0) # Extract per-microbatch metadata for PP forward: # - position_ids: full-length (not SP-sliced) for correct per-document RoPE # - cu_seq_lens/max_length: flash-attention varlen kwargs for document boundaries # Each _pp_forward call pops one entry from the deque. from collections import deque + _PP_FA_KEYS = ("cu_seq_lens_q", "cu_seq_lens_k", "max_length_q", "max_length_k") pp_metadata_list = [] for mb in micro_batches: @@ -557,18 +561,22 @@ def pp_loss_fn(pred, labels): # Loss computation: lm_head weight stays all-gathered # via reshard_after_forward=False on norm + lm_head FSDP unit result = compute_loss( - model.lm_head, outputs.last_hidden_state, - loss_fn_name=None, loss_fn_inputs={"labels": labels}, - loss_fn_params=None, logits_to_keep=0, + model.lm_head, + outputs.last_hidden_state, + loss_fn_name=None, + loss_fn_inputs={"labels": labels}, + loss_fn_params=None, + logits_to_keep=0, ) loss = result.loss # MoE aux loss from router logits (if applicable) if hasattr(outputs, "router_logits") and outputs.router_logits is not None: - aux_loss = load_balancing_loss_func( + aux_loss = global_load_balancing_loss_func( outputs.router_logits, model.num_experts, model.num_experts_per_tok, + dp_group=ps.dp_group if ps.dp_enabled else None, ) if aux_loss != 0: loss = loss + model.router_aux_loss_coef * aux_loss.to(loss.device) @@ -600,7 +608,8 @@ def pp_loss_fn(pred, labels): # Gradient clipping grad_norm = clip_gradients( - model, args.train.max_grad_norm, + model, + args.train.max_grad_norm, pp_enabled=pp_enabled, pp_group=ps.pp_group if pp_enabled else None, ) @@ -638,20 +647,25 @@ def pp_loss_fn(pred, labels): train_metrics = environ_meter.step(delta_time, global_step=global_step) tflops_per_gpu = train_metrics.get("efficiency/flops_achieved(T)", 0) / args.train.world_size - data_loader_tqdm.set_postfix_str(f"loss={total_loss:.2f} gn={grad_norm:.2f} lr={lr:.1e} tflops={tflops_per_gpu:.1f}") + data_loader_tqdm.set_postfix_str( + f"loss={total_loss:.2f} gn={grad_norm:.2f} lr={lr:.1e} tflops={tflops_per_gpu:.1f}" + ) data_loader_tqdm.update() if args.train.global_rank == 0: if args.train.use_wandb and global_step % args.train.wandb_log_interval == 0: import wandb - train_metrics.update({ - "training/loss": total_loss, - "training/grad_norm": grad_norm, - "training/lr": lr, - "training/epoch": epoch, - "training/step_time": delta_time, - "training/samples_seen": global_step * args.train.global_batch_size, - }) + + train_metrics.update( + { + "training/loss": total_loss, + "training/grad_norm": grad_norm, + "training/lr": lr, + "training/epoch": epoch, + "training/step_time": delta_time, + "training/samples_seen": global_step * args.train.global_batch_size, + } + ) wandb.log(train_metrics, step=global_step) if args.train.profile_this_rank and global_step <= args.train.profile_end_step: @@ -659,8 +673,9 @@ def pp_loss_fn(pred, labels): if global_step == args.train.profile_end_step: profiler.stop() - should_save = (args.train.save_steps and global_step % args.train.save_steps == 0) or \ - (save_epoch_steps and global_step % save_epoch_steps == 0) + should_save = (args.train.save_steps and global_step % args.train.save_steps == 0) or ( + save_epoch_steps and global_step % save_epoch_steps == 0 + ) if should_save: helper.empty_cache() save_checkpoint_path = os.path.join(args.train.save_checkpoint_path, f"global_step_{global_step}") @@ -679,8 +694,10 @@ def pp_loss_fn(pred, labels): is_lora_training = args.lora.enable_lora or args.lora.enable_qlora _save_lora_only = is_lora_training and args.lora.merge_lora_interval == 0 Checkpointer.save( - args.train.save_checkpoint_path, state, - global_steps=global_step, save_lora_only=_save_lora_only, + args.train.save_checkpoint_path, + state, + global_steps=global_step, + save_lora_only=_save_lora_only, ) dist.barrier() @@ -700,11 +717,10 @@ def pp_loss_fn(pred, labels): hf_model_state_dict = None if args.train.save_hf_weights and not save_peft_adapter: - from torch.distributed.checkpoint.state_dict import get_model_state_dict, StateDictOptions + from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict + logger.info_rank0("Gathering full model state dict for HF checkpoint via NCCL...") - hf_model_state_dict = get_model_state_dict( - model, options=StateDictOptions(full_state_dict=True) - ) + hf_model_state_dict = get_model_state_dict(model, options=StateDictOptions(full_state_dict=True)) # release memory del optimizer, lr_scheduler @@ -716,8 +732,10 @@ def pp_loss_fn(pred, labels): if save_peft_adapter: # Save PEFT adapter format (LoRA-only, base weights unchanged) from xorl.lora.utils import save_lora_checkpoint + save_lora_checkpoint( - model, hf_weights_path, + model, + hf_weights_path, base_model_name=args.model.model_path, target_modules=args.lora.lora_target_modules, r=args.lora.lora_rank, @@ -726,7 +744,9 @@ def pp_loss_fn(pred, labels): logger.info_rank0(f"PEFT adapter checkpoint saved at {hf_weights_path} successfully!") elif hf_model_state_dict is not None: checkpoint_handler = model.get_checkpoint_handler() if hasattr(model, "get_checkpoint_handler") else None - save_model_weights(hf_weights_path, hf_model_state_dict, model_assets=model_assets, checkpoint_handler=checkpoint_handler) + save_model_weights( + hf_weights_path, hf_model_state_dict, model_assets=model_assets, checkpoint_handler=checkpoint_handler + ) del hf_model_state_dict logger.info_rank0(f"Huggingface checkpoint saved at {hf_weights_path} successfully!") diff --git a/src/xorl/models/layers/moe/__init__.py b/src/xorl/models/layers/moe/__init__.py index 8ed8ed71..e9f600dc 100644 --- a/src/xorl/models/layers/moe/__init__.py +++ b/src/xorl/models/layers/moe/__init__.py @@ -10,18 +10,20 @@ - :data:`MOE_EXPERT_BACKENDS` — backend registry (eager / triton / native / quack). """ +from .aux_loss import global_load_balancing_loss_func +from .backend import MOE_EXPERT_BACKENDS from .experts import MoEExperts -from .router import TopKRouter -from .moe_block import MoEBlock from .lora import ( - MoELoRAConfig, MoEExpertsLoRA, - inject_lora_into_experts, + MoELoRAConfig, copy_weights_to_lora_experts, - mark_only_lora_as_trainable, + inject_lora_into_experts, lora_state_dict, + mark_only_lora_as_trainable, ) -from .backend import MOE_EXPERT_BACKENDS +from .moe_block import MoEBlock +from .router import TopKRouter + __all__ = [ "MoEExperts", @@ -34,4 +36,5 @@ "mark_only_lora_as_trainable", "lora_state_dict", "MOE_EXPERT_BACKENDS", + "global_load_balancing_loss_func", ] diff --git a/src/xorl/models/layers/moe/aux_loss.py b/src/xorl/models/layers/moe/aux_loss.py new file mode 100644 index 00000000..0cce91fe --- /dev/null +++ b/src/xorl/models/layers/moe/aux_loss.py @@ -0,0 +1,105 @@ +"""Globally-reduced load balancing loss for MoE training.""" + +from typing import Optional, Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F + + +def global_load_balancing_loss_func( + gate_logits: Union[torch.Tensor, Tuple[torch.Tensor], None], + num_experts: Optional[int] = None, + top_k: int = 2, + attention_mask: Optional[torch.Tensor] = None, + dp_group: Optional[dist.ProcessGroup] = None, +) -> Union[torch.Tensor, int]: + r""" + Computes auxiliary load balancing loss as in Switch Transformer with + optional global reduction across data-parallel ranks. + + See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. + This function implements the loss function presented in equations (4) - (6) + of the paper. It aims at penalizing cases where the routing between experts + is too unbalanced. + + When ``dp_group`` is provided, ``tokens_per_expert`` is all-reduced across + the group so the loss reflects the global routing distribution rather than + just the local micro-batch. ``router_prob_per_expert`` stays local so that + gradients flow back to the local gate parameters. + + Args: + gate_logits: + Logits from the ``gate``, should be a tuple of tensors of + shape ``[batch_size * sequence_length, num_experts]`` (one per MoE layer). + num_experts: + Number of experts. + top_k: + The number of experts to route per-token. + attention_mask: + Optional attention mask of shape ``[batch_size, sequence_length]``. + dp_group: + Data-parallel process group for global reduction. ``None`` skips + the all-reduce (single-rank or local-only mode). + + Returns: + The auxiliary loss (scalar tensor), or ``0`` if MoE is not active. + """ + if gate_logits is None or not isinstance(gate_logits, tuple): + return 0 + + compute_device = gate_logits[0].device + concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0) + + routing_weights = F.softmax(concatenated_gate_logits, dim=-1) + + _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) + + expert_mask = F.one_hot(selected_experts, num_experts) + + if attention_mask is None: + # Compute the percentage of tokens routed to each expert + tokens_per_expert = torch.mean(expert_mask.float(), dim=0) + + # Globally reduce tokens_per_expert across DP ranks + if dp_group is not None: + dist.all_reduce(tokens_per_expert, op=dist.ReduceOp.AVG, group=dp_group) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.mean(routing_weights, dim=0) + else: + batch_size, sequence_length = attention_mask.shape + num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) + + # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask + expert_attention_mask = ( + attention_mask[None, :, :, None, None] + .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts)) + .reshape(-1, top_k, num_experts) + .to(compute_device) + ) + + # Compute the percentage of tokens routed to each expert + tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum( + expert_attention_mask, dim=0 + ) + + # Globally reduce tokens_per_expert across DP ranks + if dp_group is not None: + dist.all_reduce(tokens_per_expert, op=dist.ReduceOp.AVG, group=dp_group) + + # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert + router_per_expert_attention_mask = ( + attention_mask[None, :, :, None] + .expand((num_hidden_layers, batch_size, sequence_length, num_experts)) + .reshape(-1, num_experts) + .to(compute_device) + ) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum( + router_per_expert_attention_mask, dim=0 + ) + + overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0)) + return overall_loss * num_experts diff --git a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py index 3d59a221..6f4e7fad 100644 --- a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py +++ b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py @@ -12,14 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple, Union, Unpack +from typing import Optional, Tuple, Unpack import torch from torch import nn +from xorl.distributed.moe.deepep import sync_pending_combine from xorl.distributed.parallel_state import get_parallel_state from xorl.distributed.sequence_parallel.strategy import get_cp_strategy from xorl.models.base import XorlPreTrainedModel +from xorl.models.checkpoint_handlers.buffers import ( + checkpoint_has_per_expert_weights, + detect_prequantized_block_fp8_checkpoint, + detect_prequantized_checkpoint, + get_prequantized_exclude_modules, +) from xorl.models.layers import ACT2FN, RMSNorm, RotaryEmbedding from xorl.models.layers.attention import ( AttentionKwargs, @@ -27,21 +34,15 @@ is_flash_attention, update_causal_mask, ) -from xorl.models.layers.moe import MoEExperts, MoEBlock +from xorl.models.layers.moe import MoEBlock, MoEExperts from xorl.models.outputs import MoeCausalLMOutput, MoeModelOutput -from xorl.models.checkpoint_handlers.buffers import ( - checkpoint_has_per_expert_weights, - detect_prequantized_checkpoint, - detect_prequantized_block_fp8_checkpoint, - get_prequantized_exclude_modules, -) +from xorl.models.transformers.qwen3_moe import parallelize from xorl.models.transformers.qwen3_moe.checkpoint_handler import Qwen3MoeCheckpointHandler from xorl.models.transformers.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig -from xorl.models.transformers.qwen3_moe import parallelize -from xorl.distributed.moe.deepep import sync_pending_combine from xorl.ops.fused_silu_and_mul import fused_silu_and_mul from xorl.utils import logging + logger = logging.get_logger(__name__) @@ -55,7 +56,7 @@ def __init__(self, config, intermediate_size=None): self.act_fn = ACT2FN[config.hidden_act] self.ep_dispatch = getattr(config, "_ep_dispatch", "alltoall") self.deepep_buffer_size_gb = getattr(config, "_deepep_buffer_size_gb", 2.0) - self._use_fused_silu = config.hidden_act == "silu" and not getattr(config, '_activation_native', False) + self._use_fused_silu = config.hidden_act == "silu" and not getattr(config, "_activation_native", False) def unfuse_for_tp(self): """Replace fused gate_up_proj with separate gate_proj and up_proj for tensor parallelism.""" @@ -103,7 +104,6 @@ def __init__(self, config): ) - class Qwen3MoeQuackExperts(MoEExperts): """Backward-compat wrapper: quack group GEMM expert module.""" @@ -119,6 +119,7 @@ def __init__(self, config): class Qwen3MoeAttention(MultiHeadAttention): """Qwen3 MoE attention — uses base MultiHeadAttention with default sliding window.""" + pass @@ -336,9 +337,8 @@ def get_checkpoint_handler(self, **kwargs): is_broadcast = kwargs.get("is_broadcast", False) has_per_expert = checkpoint_has_per_expert_weights(checkpoint_keys) if checkpoint_keys else True - is_prequantized = ( - detect_prequantized_checkpoint(weights_path) - or detect_prequantized_block_fp8_checkpoint(weights_path) + is_prequantized = detect_prequantized_checkpoint(weights_path) or detect_prequantized_block_fp8_checkpoint( + weights_path ) # Use user-specified exclude_modules (stored by train.py) if available, @@ -457,7 +457,9 @@ def forward( # SP strategy handles slicing (sync: slice, async: keep full-length) ps = get_parallel_state() position_embeddings = get_cp_strategy(num_kv_heads=self.config.num_key_value_heads).prepare_position_embeddings( - position_embeddings, dim=1, sp_group=ps.sp_group, + position_embeddings, + dim=1, + sp_group=ps.sp_group, num_kv_heads=self.config.num_key_value_heads, ) @@ -471,8 +473,7 @@ def forward( # When selective checkpointing is enabled (_recompute_modules is set), # the decoder layer handles its own sub-checkpointing — skip the outer checkpoint. _use_outer_checkpoint = ( - self.gradient_checkpointing and self.training - and getattr(self, "_recompute_modules", None) is None + self.gradient_checkpointing and self.training and getattr(self, "_recompute_modules", None) is None ) if _use_outer_checkpoint: @@ -514,89 +515,8 @@ def forward( router_logits=all_router_logits, ) -class KwargsForCausalLM(AttentionKwargs): ... - -def load_balancing_loss_func( - gate_logits: Union[torch.Tensor, Tuple[torch.Tensor], None], - num_experts: Optional[int] = None, - top_k=2, - attention_mask: Optional[torch.Tensor] = None, -) -> Union[torch.Tensor, int]: - r""" - Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch. - - See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss - function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between - experts is too unbalanced. - - Args: - gate_logits: - Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of - shape [batch_size X sequence_length, num_experts]. - num_experts: - Number of experts - top_k: - The number of experts to route per-token, can be also interpreted as the `top-k` routing - parameter. - attention_mask (`torch.Tensor`, *optional*): - The attention_mask used in forward function - shape [batch_size X sequence_length] if not None. - - Returns: - The auxiliary loss. - """ - if gate_logits is None or not isinstance(gate_logits, tuple): - return 0 - - if isinstance(gate_logits, tuple): - compute_device = gate_logits[0].device - concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0) - - routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1) - - _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) - - expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) - - if attention_mask is None: - # Compute the percentage of tokens routed to each experts - tokens_per_expert = torch.mean(expert_mask.float(), dim=0) - - # Compute the average probability of routing to these experts - router_prob_per_expert = torch.mean(routing_weights, dim=0) - else: - batch_size, sequence_length = attention_mask.shape - num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) - - # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask - expert_attention_mask = ( - attention_mask[None, :, :, None, None] - .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts)) - .reshape(-1, top_k, num_experts) - .to(compute_device) - ) - - # Compute the percentage of tokens routed to each experts - tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum( - expert_attention_mask, dim=0 - ) - - # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert - router_per_expert_attention_mask = ( - attention_mask[None, :, :, None] - .expand((num_hidden_layers, batch_size, sequence_length, num_experts)) - .reshape(-1, num_experts) - .to(compute_device) - ) - - # Compute the average probability of routing to these experts - router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum( - router_per_expert_attention_mask, dim=0 - ) - - overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0)) - return overall_loss * num_experts +class KwargsForCausalLM(AttentionKwargs): ... class Qwen3MoeForCausalLM(Qwen3MoePreTrainedModel): diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index 4360cf5e..ae592074 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -29,24 +29,25 @@ from xorl.distributed.sync_padding import synchronize_micro_batch_padding from xorl.distributed.torch_parallelize import build_parallelize_model from xorl.models import build_foundation_model, build_tokenizer, save_model_assets, save_model_weights +from xorl.models.layers.moe.aux_loss import global_load_balancing_loss_func from xorl.models.layers.moe.routing_replay import RoutingReplay, set_replay_stage from xorl.models.module_utils import compute_loss -from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import load_balancing_loss_func from xorl.optim import build_lr_scheduler, build_optimizer from xorl.trainers.training_utils import ( - pp_loss_fn, clip_gradients, count_valid_tokens, forward_backward_pp, maybe_merge_lora, negotiate_pp_seq_len, pad_micro_batches_for_pp, + pp_loss_fn, sync_sp_gradients, ) from xorl.utils import helper from xorl.utils.device import get_device_type, get_nccl_backend, get_torch_device, synchronize from xorl.utils.dist_utils import all_reduce + logger = helper.create_logger(__name__) @@ -54,6 +55,7 @@ # TrainState — checkpointable training state # --------------------------------------------------------------------------- + @dataclass class TrainState: """Minimal checkpointable state (follows torchtitan's Stateful pattern).""" @@ -82,6 +84,7 @@ def load_state_dict(self, d: Dict[str, Any]) -> None: # Trainer # --------------------------------------------------------------------------- + class Trainer: """Offline training loop. Analogous to torchtitan's Trainer. @@ -132,6 +135,7 @@ def _bootstrap(self) -> None: if args.train.global_rank == 0: from xorl.arguments import save_args + save_args(args, args.train.output_dir) self.Checkpointer = build_checkpointer( @@ -156,6 +160,7 @@ def _bootstrap(self) -> None: self.ps = get_parallel_state() if self.ps.device_mesh is not None: import torch.distributed.tensor._random + torch.distributed.tensor._random.manual_seed(args.train.seed, self.ps.device_mesh) # Routing replay is only needed with EP when MoE forward is recomputed @@ -185,17 +190,14 @@ def _build_data(self) -> None: if args.train.max_steps is not None: self.total_train_steps = min(self.total_train_steps, args.train.max_steps) logger.info_rank0( - f"Train steps per epoch: {self.train_steps_per_epoch}, " - f"Total train steps: {self.total_train_steps}" + f"Train steps per epoch: {self.train_steps_per_epoch}, Total train steps: {self.total_train_steps}" ) self.save_epoch_steps = ( int(args.train.save_epochs * self.train_steps_per_epoch) if args.train.save_epochs else 0 ) if self.save_epoch_steps: - logger.info_rank0( - f"Save every {args.train.save_epochs} epoch(s) = every {self.save_epoch_steps} steps" - ) + logger.info_rank0(f"Save every {args.train.save_epochs} epoch(s) = every {self.save_epoch_steps} steps") def _build_model(self) -> None: """Build foundation model and inject LoRA/QLoRA if configured.""" @@ -238,7 +240,7 @@ def _build_model(self) -> None: def _inject_qlora(self) -> None: """QLoRA injection with pre-quantized checkpoint detection.""" args = self.args - from xorl.qlora import inject_qlora_into_model, detect_prequantized_nvfp4, detect_prequantized_block_fp8 + from xorl.qlora import detect_prequantized_block_fp8, detect_prequantized_nvfp4, inject_qlora_into_model if detect_prequantized_nvfp4(args.model.model_path): self.is_prequantized = True @@ -254,6 +256,7 @@ def _inject_qlora(self) -> None: logger.info_rank0(f"Using user-specified exclude_modules: {self.exclude_modules}") elif self.is_prequantized: from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules + self.exclude_modules = get_prequantized_exclude_modules(args.model.model_path) if self.exclude_modules: logger.info_rank0( @@ -288,6 +291,7 @@ def _inject_lora(self) -> None: """Plain LoRA injection.""" args = self.args from xorl.lora.utils import inject_lora_into_model + inject_lora_into_model( self.model, r=args.lora.lora_rank, @@ -301,8 +305,7 @@ def _parallelize(self) -> None: args = self.args _t0 = time.time() logger.info_rank0( - f"Loading model weights (mode={args.train.load_weights_mode}, " - f"init_device={args.train.init_device})..." + f"Loading model weights (mode={args.train.load_weights_mode}, init_device={args.train.init_device})..." ) build_result = build_parallelize_model( self.model, @@ -323,7 +326,7 @@ def _parallelize(self) -> None: skip_param_upcast=args.lora.enable_qlora, ) - logger.info_rank0(f"Model weights loaded in {time.time()-_t0:.1f}s") + logger.info_rank0(f"Model weights loaded in {time.time() - _t0:.1f}s") helper.print_device_mem_info("VRAM after loading weights") # PP returns dict with stages + model_parts; otherwise returns model directly @@ -356,13 +359,15 @@ def _deferred_qlora_quantize(self) -> None: args = self.args if self.is_prequantized: from xorl.qlora import maybe_load_prequantized_qlora + logger.info("Starting pre-quantized weight loading...") helper.print_device_mem_info("VRAM before pre-quantized loading") maybe_load_prequantized_qlora(self.model, args.model.model_path) logger.info("Done pre-quantized weight loading, freezing non-LoRA params...") else: - from xorl.qlora import maybe_quantize_qlora, maybe_load_and_quantize_moe_qlora + from xorl.qlora import maybe_load_and_quantize_moe_qlora, maybe_quantize_qlora from xorl.qlora.utils import _deregister_qlora_weights_from_fsdp + logger.info("Starting maybe_quantize_qlora...") helper.print_device_mem_info("VRAM before QLoRA quantization") maybe_quantize_qlora(self.model) @@ -372,7 +377,8 @@ def _deferred_qlora_quantize(self) -> None: logger.info("Done MoE weight loading, freezing non-LoRA params...") # Deregister packed_weight_f32 from FSDP2 (prevent mixed-precision corruption) removed = _deregister_qlora_weights_from_fsdp( - self.model, param_names=("packed_weight_f32",), + self.model, + param_names=("packed_weight_f32",), ) torch.cuda.empty_cache() if removed > 0: @@ -419,6 +425,7 @@ def _setup_observability(self) -> None: if args.train.global_rank == 0: if args.train.use_wandb: import wandb + wandb.init( project=args.train.wandb_project, name=args.train.wandb_name, @@ -494,6 +501,7 @@ def _get_pp_schedule(self, seq_len: int): """ if seq_len not in self._pp_schedule_cache: from xorl.distributed.pipeline_parallel import build_pipeline_schedule, build_pp_stage + stage = build_pp_stage( self.model_parts[0], pp_rank=self.ps.pp_rank, @@ -601,8 +609,15 @@ def train(self) -> None: state.loss_history.append(total_loss) # Logging - self._maybe_log(total_loss, grad_norm, lr, train_metrics, delta_time, use_tqdm, - data_loader_tqdm if use_tqdm else None) + self._maybe_log( + total_loss, + grad_norm, + lr, + train_metrics, + delta_time, + use_tqdm, + data_loader_tqdm if use_tqdm else None, + ) self._maybe_profile() self._maybe_save_checkpoint() @@ -711,17 +726,21 @@ def _forward_backward( with self._model_fwd_context: outputs = self.model(**micro_batch, use_cache=False, output_hidden_states=False) result = compute_loss( - self.model.lm_head, outputs.last_hidden_state, - loss_fn_name=None, loss_fn_inputs={"labels": labels}, - loss_fn_params=None, logits_to_keep=0, + self.model.lm_head, + outputs.last_hidden_state, + loss_fn_name=None, + loss_fn_inputs={"labels": labels}, + loss_fn_params=None, + logits_to_keep=0, ) loss = result.loss if hasattr(outputs, "router_logits") and outputs.router_logits is not None: - aux_loss = load_balancing_loss_func( + aux_loss = global_load_balancing_loss_func( outputs.router_logits, self.model.num_experts, self.model.num_experts_per_tok, + dp_group=self.ps.dp_group if self.ps.dp_enabled else None, ) if aux_loss != 0: loss = loss + self.model.router_aux_loss_coef * aux_loss.to(loss.device) @@ -845,16 +864,23 @@ def _maybe_log( f"tokens_per_sec={tokens_per_sec:.0f} time={delta_time:.3f}s" ) - if args.train.global_rank == 0 and args.train.use_wandb and self.state.global_step % args.train.wandb_log_interval == 0: + if ( + args.train.global_rank == 0 + and args.train.use_wandb + and self.state.global_step % args.train.wandb_log_interval == 0 + ): import wandb - train_metrics.update({ - "training/loss": total_loss, - "training/grad_norm": grad_norm, - "training/lr": lr, - "training/epoch": self.state.epoch, - "training/step_time": delta_time, - "training/samples_seen": self.state.global_step * args.train.global_batch_size, - }) + + train_metrics.update( + { + "training/loss": total_loss, + "training/grad_norm": grad_norm, + "training/lr": lr, + "training/epoch": self.state.epoch, + "training/step_time": delta_time, + "training/samples_seen": self.state.global_step * args.train.global_batch_size, + } + ) # Log per-group LRs (e.g., separate Muon vs AdamW LRs) for group in self.optimizer.param_groups: if group.get("use_muon", False): @@ -877,17 +903,14 @@ def _maybe_save_checkpoint(self) -> None: args = self.args step = self.state.global_step - should_save = ( - (args.train.save_steps and step % args.train.save_steps == 0) - or (self.save_epoch_steps and step % self.save_epoch_steps == 0) + should_save = (args.train.save_steps and step % args.train.save_steps == 0) or ( + self.save_epoch_steps and step % self.save_epoch_steps == 0 ) if not should_save: return helper.empty_cache() - save_checkpoint_path = os.path.join( - args.train.save_checkpoint_path, f"global_step_{step}" - ) + save_checkpoint_path = os.path.join(args.train.save_checkpoint_path, f"global_step_{step}") state = { "model": self.model, "optimizer": self.optimizer, @@ -902,8 +925,10 @@ def _maybe_save_checkpoint(self) -> None: is_lora_training = args.lora.enable_lora or args.lora.enable_qlora _save_lora_only = is_lora_training and args.lora.merge_lora_interval == 0 self.Checkpointer.save( - args.train.save_checkpoint_path, state, - global_steps=step, save_lora_only=_save_lora_only, + args.train.save_checkpoint_path, + state, + global_steps=step, + save_lora_only=_save_lora_only, ) dist.barrier() logger.info_rank0(f"Checkpoint saved at {save_checkpoint_path}") @@ -925,7 +950,8 @@ def _finalize(self, total_loss: float, grad_norm: float, lr: float) -> None: hf_model_state_dict = None if args.train.save_hf_weights and not save_peft_adapter: - from torch.distributed.checkpoint.state_dict import get_model_state_dict, StateDictOptions + from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict + logger.info_rank0("Gathering full model state dict for HF checkpoint via NCCL with CPU offload...") hf_model_state_dict = get_model_state_dict( self.model, options=StateDictOptions(full_state_dict=True, cpu_offload=True) @@ -936,13 +962,13 @@ def _finalize(self, total_loss: float, grad_norm: float, lr: float) -> None: # Save HF weights (rank 0) if args.train.global_rank == 0 and args.train.save_hf_weights: - hf_weights_path = os.path.join( - args.train.output_dir, f"global_step_{state.global_step}", "hf_ckpt" - ) + hf_weights_path = os.path.join(args.train.output_dir, f"global_step_{state.global_step}", "hf_ckpt") if save_peft_adapter: from xorl.lora.utils import save_lora_checkpoint + save_lora_checkpoint( - self.model, hf_weights_path, + self.model, + hf_weights_path, base_model_name=args.model.model_path, target_modules=args.lora.lora_target_modules, r=args.lora.lora_rank, @@ -951,11 +977,11 @@ def _finalize(self, total_loss: float, grad_norm: float, lr: float) -> None: logger.info_rank0(f"PEFT adapter saved at {hf_weights_path}") elif hf_model_state_dict is not None: checkpoint_handler = ( - self.model.get_checkpoint_handler() - if hasattr(self.model, "get_checkpoint_handler") else None + self.model.get_checkpoint_handler() if hasattr(self.model, "get_checkpoint_handler") else None ) save_model_weights( - hf_weights_path, hf_model_state_dict, + hf_weights_path, + hf_model_state_dict, model_assets=self.model_assets, checkpoint_handler=checkpoint_handler, ) @@ -967,14 +993,18 @@ def _finalize(self, total_loss: float, grad_norm: float, lr: float) -> None: metrics_path = os.path.join(args.train.output_dir, "training_metrics.json") os.makedirs(args.train.output_dir, exist_ok=True) with open(metrics_path, "w") as f: - json.dump({ - "final_loss": total_loss, - "final_grad_norm": grad_norm, - "final_lr": lr, - "global_step": state.global_step, - "total_train_steps": self.total_train_steps, - "loss_history": state.loss_history, - }, f, indent=2) + json.dump( + { + "final_loss": total_loss, + "final_grad_norm": grad_norm, + "final_lr": lr, + "global_step": state.global_step, + "total_train_steps": self.total_train_steps, + "loss_history": state.loss_history, + }, + f, + indent=2, + ) dist.barrier() dist.destroy_process_group() From c37ab974f7321d245a3cbd3393fe4530264ee16f Mon Sep 17 00:00:00 2001 From: Felix Zhang Date: Fri, 27 Mar 2026 11:40:30 -0700 Subject: [PATCH 03/41] Add opt-in MoE LoRA PEFT export flag (#44) (cherry picked from commit ce41728320ba988724c95653b426c81ee6074b7b) --- src/xorl/lora/utils.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/src/xorl/lora/utils.py b/src/xorl/lora/utils.py index 06e5c31e..e71c2afd 100644 --- a/src/xorl/lora/utils.py +++ b/src/xorl/lora/utils.py @@ -378,6 +378,7 @@ def save_lora_checkpoint( lora_alpha: Optional[int] = None, moe_hybrid_shared_lora: bool = False, lora_state_dict: Optional[Dict[str, torch.Tensor]] = None, + transpose_moe_lora_to_peft: bool = False, ) -> str: """ Save LoRA weights in PEFT-compatible format. @@ -397,6 +398,9 @@ def save_lora_checkpoint( lora_state_dict: Pre-gathered LoRA state dict. If provided, skips get_lora_state_dict(model) call. Useful when the caller has already gathered weights (e.g., from FSDP2 + EP distributed model). + transpose_moe_lora_to_peft: Whether to transpose MoE expert LoRA tensors + into PEFT/vLLM orientation during export. Disabled by default so + existing xorl call sites keep their current behavior. Returns: Path to saved checkpoint directory @@ -439,10 +443,12 @@ def _unmerge_moe_lora_weights(name: str, stacked_tensor: torch.Tensor) -> Dict[s Unmerge stacked MoE LoRA weights into per-expert format for vLLM compatibility. Xorl stores MoE LoRA as stacked tensors: - model.layers.0.mlp.experts.gate_proj_lora_A # shape: [num_experts, r, hidden_dim] + model.layers.0.mlp.experts.gate_proj_lora_A # shape: [num_experts, hidden_dim, r] + model.layers.0.mlp.experts.gate_proj_lora_B # shape: [num_experts, r, intermediate_dim] vLLM expects per-expert format: model.layers.0.mlp.experts.0.gate_proj.lora_A.weight # shape: [r, hidden_dim] + model.layers.0.mlp.experts.0.gate_proj.lora_B.weight # shape: [intermediate_dim, r] For hybrid shared LoRA, shared weights (shape[0] == 1) are named with ".shared." instead of expert index to indicate they're shared across experts. @@ -464,12 +470,16 @@ def _unmerge_moe_lora_weights(name: str, stacked_tensor: torch.Tensor) -> Dict[s if is_shared: # Shared weight: use ".shared." in the key name expert_tensor = stacked_tensor[0] + if transpose_moe_lora_to_peft: + expert_tensor = expert_tensor.transpose(0, 1).contiguous() peft_key = f"base_model.model.{prefix}.mlp.experts.shared.{proj_name}.lora_{lora_type}.weight" result[peft_key] = expert_tensor.to(torch.bfloat16) else: # Per-expert weights: use expert index in the key name for expert_idx in range(num_experts): expert_tensor = stacked_tensor[expert_idx] + if transpose_moe_lora_to_peft: + expert_tensor = expert_tensor.transpose(0, 1).contiguous() # Build vLLM-compatible key: # base_model.model.{prefix}.mlp.experts.{idx}.{proj}.lora_{A|B}.weight peft_key = f"base_model.model.{prefix}.mlp.experts.{expert_idx}.{proj_name}.lora_{lora_type}.weight" @@ -494,8 +504,10 @@ def _unmerge_moe_lora_weights(name: str, stacked_tensor: torch.Tensor) -> Dict[s if match: detected_modules.add(match.group(2)) # gate_proj, up_proj, or down_proj if detected_r is None and match.group(3) == "A": - # For MoE LoRA, shape is [num_experts, r, ...] - detected_r = value.shape[1] + # Xorl stores MoE LoRA A as [num_experts, in_features, r]. + # When transpose_moe_lora_to_peft is enabled, the exported + # PEFT tensor rank is the last dimension of the stacked input. + detected_r = value.shape[2] if transpose_moe_lora_to_peft else value.shape[1] else: # Extract module name for target_modules detection parts = key.split(".") From ec37d38cf672689c5c26af3661abf190acc57f5c Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Tue, 31 Mar 2026 15:04:25 -0700 Subject: [PATCH 04/41] Fix rank-0 shard broadcast deadlock (#54) (cherry picked from commit 69ee83a5f59c8eaf8bc61f998f6d8808cc91bd0e) --- src/xorl/models/module_utils.py | 9 ++- tests/models/test_module_utils_broadcast.py | 82 +++++++++++++++++++++ 2 files changed, 87 insertions(+), 4 deletions(-) create mode 100644 tests/models/test_module_utils_broadcast.py diff --git a/src/xorl/models/module_utils.py b/src/xorl/models/module_utils.py index f3878806..971b52a5 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -995,9 +995,10 @@ def rank0_load_and_broadcast_weights( else: global_expected_keys = None - # Get the safetensor file iterators - state_dict_iterators = _load_state_dict(weights_path) if global_rank == 0 else None - shard_count = len(state_dict_iterators) if global_rank == 0 else 0 + # All ranks must enter _load_state_dict(): in multi-rank mode it contains + # the collective that receives the rank-0-resolved shard paths. + state_dict_iterators = _load_state_dict(weights_path) + shard_count = len(state_dict_iterators) logger.info_rank0(f"rank0_load_and_broadcast_weights: {shard_count=} ") shard_count_tensor = torch.tensor( [shard_count], @@ -1421,4 +1422,4 @@ class GradientCheckpointingLayer(nn.Module): def __call__(self, *args, **kwargs): if self.gradient_checkpointing and self.training: return self._gradient_checkpointing_func(partial(super().__call__, **kwargs), *args) - return super().__call__(*args, **kwargs) \ No newline at end of file + return super().__call__(*args, **kwargs) diff --git a/tests/models/test_module_utils_broadcast.py b/tests/models/test_module_utils_broadcast.py new file mode 100644 index 00000000..e9d440bc --- /dev/null +++ b/tests/models/test_module_utils_broadcast.py @@ -0,0 +1,82 @@ +from types import SimpleNamespace + +import pytest + +from xorl.models import module_utils + + +pytestmark = [pytest.mark.cpu] + + +class _DummyModel: + def named_buffers(self): + return [] + + def named_parameters(self): + return [] + + def to_empty(self, device): + self.device = device + + +def test_rank0_broadcast_path_calls_load_state_dict_on_nonzero_ranks(monkeypatch): + calls = [] + + def fake_broadcast_object_list(obj, src=0): + if obj[0] is None: + obj[0] = [] + + fake_dist = SimpleNamespace( + is_available=lambda: True, + is_initialized=lambda: True, + broadcast=lambda tensor, src=0: None, + broadcast_object_list=fake_broadcast_object_list, + ) + + monkeypatch.setattr(module_utils, "dist", fake_dist) + monkeypatch.setattr( + module_utils, + "get_parallel_state", + lambda: SimpleNamespace(global_rank=1, pp_enabled=False), + ) + monkeypatch.setattr(module_utils, "_build_compiled_key_map", lambda *args, **kwargs: {}) + monkeypatch.setattr(module_utils, "_shrink_expert_params_for_ep", lambda model: None) + monkeypatch.setattr(module_utils, "post_process_after_weight_loading", lambda *args, **kwargs: None) + + def fake_load_state_dict(weights_path): + calls.append(weights_path) + return [] + + monkeypatch.setattr(module_utils, "_load_state_dict", fake_load_state_dict) + + module_utils.rank0_load_and_broadcast_weights(_DummyModel(), "dummy-weights", init_device="cpu") + + assert calls == ["dummy-weights"] + + +def test_try_load_state_dict_uses_rank0_for_local_resolution(monkeypatch): + local_resolution_calls = [] + + def fake_broadcast_object_list(obj, src=0, group=None): + assert src == 0 + obj[0] = ["shard-0.safetensors", "shard-1.safetensors"] + + fake_dist = SimpleNamespace( + is_initialized=lambda: True, + get_rank=lambda: 1, + get_world_size=lambda: 64, + broadcast_object_list=fake_broadcast_object_list, + ) + + def fake_try_load_state_dict_local(weights_path, **kwargs): + local_resolution_calls.append(weights_path) + return [module_utils.StateDictIterator("local-only.safetensors")] + + monkeypatch.setattr(module_utils, "dist", fake_dist) + monkeypatch.setattr(module_utils, "_get_cpu_group", lambda: object()) + monkeypatch.setattr(module_utils, "_try_load_state_dict_local", fake_try_load_state_dict_local) + + iterators = module_utils._try_load_state_dict("dummy-weights") + + assert local_resolution_calls == [] + assert [it.filepath for it in iterators] == ["shard-0.safetensors", "shard-1.safetensors"] From 7af920fbb4b5a64fd6d95fe576e8cf43da253be6 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Tue, 31 Mar 2026 15:06:38 -0700 Subject: [PATCH 05/41] Fix TFLOPs accounting and improve wandb/startup metrics (#48) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Fix TFLOPs accounting and improve wandb/startup metrics - Align MoE TFLOPs with logical model FLOPs and fix per-document seqlen² cost calculation in helper.py - Move wandb init earlier in bootstrap for startup metric capture - Add host inventory logging and memory snapshots during setup - Track repo commit hash in experiment metadata - Show tok/s instead of tflops in training progress bar * Fix TFLOPs overcounting: embedding is a lookup, not a matmul embed_lm_N was computed as vocab_size * hidden_size * 2, counting both the embedding table and lm_head as matmuls. But embedding forward is a table lookup (zero multiply-accumulate ops) — only lm_head is a real matmul. This matches torchtitan's approach of explicitly excluding nn.Embedding parameters from the FLOPs count. Fix: change embed_lm_N = vocab_size * hidden_size * 2 to embed_lm_N = vocab_size * hidden_size (lm_head only) For Qwen3-32B at seq_len=128k, this removes ~600 TFLOPS of phantom FLOPs per step, giving more accurate MFU reporting. Applied across all model types: deepseek_v3, qwen3_moe, qwen3_5, qwen3_5_moe, qwen2/qwen3/llama, qwen2_vl. (cherry picked from commit 4abe7983ba02b15a3526c80c93d1d936764396e0) --- src/xorl/arguments.py | 23 +++++ src/xorl/cli/direct_train.py | 72 +++++++++++---- src/xorl/trainers/trainer.py | 163 +++++++++++++++++++++++++++++++--- src/xorl/utils/count_flops.py | 66 ++++---------- src/xorl/utils/helper.py | 28 +++--- 5 files changed, 261 insertions(+), 91 deletions(-) diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 3704c9a6..779207c3 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -4,6 +4,7 @@ import json import math import os +import subprocess import sys import types @@ -34,6 +35,22 @@ logger = logging.get_logger(__name__) +def _detect_repo_commit() -> Optional[str]: + """Best-effort git commit detection for experiment metadata.""" + try: + repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + result = subprocess.run( + ["git", "-C", repo_root, "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ) + commit = result.stdout.strip() + return commit or None + except Exception: + return None + + def get_default_process_count(): if xorl_dataset_num_proc := os.environ.get("XORL_DATASET_NUM_PROC"): return int(xorl_dataset_num_proc) @@ -869,6 +886,10 @@ def moe_recomputed(self) -> bool: default=None, metadata={"help": "Wandb tags for organizing experiments."}, ) + repo_commit: Optional[str] = field( + default=None, + metadata={"help": "Git commit hash for experiment provenance. Auto-detected when omitted."}, + ) wandb_log_interval: int = field( default=1, metadata={"help": "Log to wandb every N steps."}, @@ -913,6 +934,8 @@ def moe_recomputed(self) -> bool: ) def __post_init__(self): + if self.repo_commit is None: + self.repo_commit = _detect_repo_commit() self.local_rank = int(os.getenv("LOCAL_RANK", "0")) self.global_rank = int(os.getenv("RANK", "0")) self.world_size = int(os.getenv("WORLD_SIZE", "1")) diff --git a/src/xorl/cli/direct_train.py b/src/xorl/cli/direct_train.py index 36766243..166c6925 100644 --- a/src/xorl/cli/direct_train.py +++ b/src/xorl/cli/direct_train.py @@ -7,7 +7,9 @@ import json import time -from dataclasses import asdict +import socket +from dataclasses import asdict, dataclass, field +from functools import partial from typing import Any, Dict, List import torch.distributed as dist @@ -60,6 +62,55 @@ def main(): if args.train.global_rank == 0: save_args(args, args.train.output_dir) + if args.train.use_wandb: + import wandb + wandb.init( + project=args.train.wandb_project, + name=args.train.wandb_name, + tags=args.train.wandb_tags, + config={**vars(args.model), **vars(args.data), **vars(args.train)}, + ) + config_file = os.path.join(args.train.output_dir, "xorl_cli.yaml") + if os.path.exists(config_file): + wandb.save(config_file, policy="now") + + host_payload = { + "global_rank": args.train.global_rank, + "local_rank": args.train.local_rank, + "hostname": socket.gethostname(), + } + gathered_hosts = [None] * args.train.world_size + dist.all_gather_object(gathered_hosts, host_payload) + if args.train.global_rank == 0: + unique_hostnames = sorted({item["hostname"] for item in gathered_hosts if item is not None}) + rank_to_hostname = { + str(item["global_rank"]): item["hostname"] + for item in gathered_hosts if item is not None + } + logger.info_rank0( + "Host inventory:\n" + json.dumps( + { + "master_addr": os.environ.get("MASTER_ADDR"), + "master_port": os.environ.get("MASTER_PORT"), + "node_count": len(unique_hostnames), + "hostnames": unique_hostnames, + "ranks": gathered_hosts, + }, + indent=2, + ) + ) + if args.train.use_wandb: + import wandb + wandb.config.update( + { + "master_addr": os.environ.get("MASTER_ADDR"), + "master_port": os.environ.get("MASTER_PORT"), + "hostnames": unique_hostnames, + "rank_to_hostname": rank_to_hostname, + }, + allow_val_change=True, + ) + wandb.log({"startup/node_count": len(unique_hostnames)}, step=0, commit=False) Checkpointer = build_checkpointer(dist_backend=args.train.data_parallel_mode, ckpt_manager=args.train.ckpt_manager) @@ -302,19 +353,6 @@ def main(): ) if args.train.global_rank == 0: - if args.train.use_wandb: - import wandb - - wandb.init( - project=args.train.wandb_project, - name=args.train.wandb_name, - tags=args.train.wandb_tags, - config={**vars(args.model), **vars(args.data), **vars(args.train)}, # flatten dict - ) - config_file = os.path.join(args.train.output_dir, "xorl_cli.yaml") - if os.path.exists(config_file): - wandb.save(config_file, policy="now") - # save model_assets before training model_assets = [model_config, tokenizer] save_model_assets(args.train.model_assets_dir, model_assets) @@ -646,10 +684,8 @@ def pp_loss_fn(pred, labels): lr = max(lr_scheduler.get_last_lr()) train_metrics = environ_meter.step(delta_time, global_step=global_step) - tflops_per_gpu = train_metrics.get("efficiency/flops_achieved(T)", 0) / args.train.world_size - data_loader_tqdm.set_postfix_str( - f"loss={total_loss:.2f} gn={grad_norm:.2f} lr={lr:.1e} tflops={tflops_per_gpu:.1f}" - ) + tokens_per_sec = train_metrics.get("efficiency/tokens_per_second(K)", 0) * 1e3 + data_loader_tqdm.set_postfix_str(f"loss={total_loss:.2f} gn={grad_norm:.2f} lr={lr:.1e} tok/s={tokens_per_sec:.0f}") data_loader_tqdm.update() if args.train.global_rank == 0: diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index ae592074..7beded2c 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -10,6 +10,7 @@ import json import os +import socket import time from dataclasses import asdict, dataclass, field from typing import Any, Dict, List, Optional, Tuple @@ -97,6 +98,9 @@ class Trainer: def __init__(self, args: Arguments) -> None: self.args = args self.state = TrainState() + self._setup_phase_metrics: Dict[str, float] = {} + self._startup_metrics: Dict[str, Any] = {} + self._wandb_initialized = False # Setup phases (order matters — each depends on previous) # Model is built before data: reads weights from disk while I/O is free, @@ -105,7 +109,10 @@ def _timed(phase_name, fn): t0 = time.time() fn() elapsed = time.time() - t0 + self._setup_phase_metrics[f"startup/{phase_name}_sec"] = elapsed logger.info(f"[TIMING] {phase_name} took {elapsed:.1f}s (rank {args.train.local_rank})") + self._maybe_log_startup_metrics({f"startup/{phase_name}_sec": elapsed}, commit=False) + self._log_memory_snapshot(f"startup/{phase_name}") _timed("bootstrap", self._bootstrap) _timed("build_model", self._build_model) @@ -115,6 +122,7 @@ def _timed(phase_name, fn): _timed("setup_observability", self._setup_observability) _timed("resume_checkpoint", self._resume_checkpoint) _timed("build_pp_schedule", self._init_pp_schedule_cache) + self._write_startup_metrics_file() # =================================================================== # Setup phases @@ -137,6 +145,34 @@ def _bootstrap(self) -> None: from xorl.arguments import save_args save_args(args, args.train.output_dir) + if args.train.use_wandb: + import wandb + wandb.init( + project=args.train.wandb_project, + name=args.train.wandb_name, + tags=args.train.wandb_tags, + config={**vars(args.model), **vars(args.data), **vars(args.train)}, + ) + self._wandb_initialized = True + config_file = os.path.join(args.train.output_dir, "xorl_cli.yaml") + if os.path.exists(config_file): + wandb.save(config_file, policy="now") + self._maybe_log_startup_metrics( + { + "startup/world_size": args.train.world_size, + "startup/global_batch_size": args.train.global_batch_size, + "startup/micro_batch_size": args.train.micro_batch_size, + "startup/gradient_accumulation_steps": args.train.gradient_accumulation_steps, + "startup/data_parallel_size": args.train.data_parallel_size, + "startup/data_parallel_replicate_size": args.train.data_parallel_replicate_size, + "startup/data_parallel_shard_size": args.train.data_parallel_shard_size, + "startup/ulysses_parallel_size": args.train.ulysses_parallel_size, + "startup/expert_parallel_size": args.train.expert_parallel_size, + "startup/pipeline_parallel_size": args.train.pipeline_parallel_size, + }, + commit=False, + ) + self._log_host_inventory() self.Checkpointer = build_checkpointer( dist_backend=args.train.data_parallel_mode, @@ -166,6 +202,112 @@ def _bootstrap(self) -> None: # Routing replay is only needed with EP when MoE forward is recomputed self._use_routing_replay = self.ps.ep_size > 1 and args.train.moe_recomputed + def _maybe_log_startup_metrics(self, metrics: Dict[str, Any], commit: bool = False) -> None: + """Log startup metrics to wandb once rank 0 has initialized it.""" + if not metrics: + return + if self.args.train.global_rank == 0: + self._startup_metrics.update(metrics) + if self.args.train.global_rank != 0 or not self.args.train.use_wandb or not self._wandb_initialized: + return + import wandb + wandb.log(metrics, step=0, commit=commit) + + def _write_startup_metrics_file(self) -> None: + """Persist startup metrics alongside the run outputs for offline inspection.""" + if self.args.train.global_rank != 0: + return + + payload = { + "repo_commit": self.args.train.repo_commit, + "wandb_project": self.args.train.wandb_project if self.args.train.use_wandb else None, + "wandb_name": self.args.train.wandb_name if self.args.train.use_wandb else None, + "metrics": self._startup_metrics, + } + output_path = os.path.join(self.args.train.output_dir, "startup_metrics.json") + with open(output_path, "w", encoding="utf-8") as f: + json.dump(payload, f, indent=2, sort_keys=True) + + def _collect_host_inventory(self) -> List[Dict[str, Any]]: + """Gather rank-to-host mapping across all processes.""" + payload = { + "global_rank": self.args.train.global_rank, + "local_rank": self.args.train.local_rank, + "hostname": socket.gethostname(), + } + gathered: List[Optional[Dict[str, Any]]] = [None] * self.args.train.world_size + dist.all_gather_object(gathered, payload) + return [item for item in gathered if item is not None] + + def _log_host_inventory(self) -> None: + """Emit host inventory to stdout and wandb config on rank 0.""" + inventory = self._collect_host_inventory() + if self.args.train.global_rank != 0: + return + + unique_hostnames = sorted({item["hostname"] for item in inventory}) + rank_to_hostname = {str(item["global_rank"]): item["hostname"] for item in inventory} + logger.info_rank0( + "Host inventory:\n" + json.dumps( + { + "master_addr": os.environ.get("MASTER_ADDR"), + "master_port": os.environ.get("MASTER_PORT"), + "node_count": len(unique_hostnames), + "hostnames": unique_hostnames, + "ranks": inventory, + }, + indent=2, + ) + ) + self._startup_metrics.update( + { + "startup/master_addr": os.environ.get("MASTER_ADDR"), + "startup/master_port": os.environ.get("MASTER_PORT"), + "startup/node_count": len(unique_hostnames), + "startup/hostnames": unique_hostnames, + "startup/rank_to_hostname": rank_to_hostname, + } + ) + self._maybe_log_startup_metrics({"startup/node_count": len(unique_hostnames)}, commit=False) + if self.args.train.use_wandb and self._wandb_initialized: + import wandb + wandb.config.update( + { + "master_addr": os.environ.get("MASTER_ADDR"), + "master_port": os.environ.get("MASTER_PORT"), + "hostnames": unique_hostnames, + "rank_to_hostname": rank_to_hostname, + }, + allow_val_change=True, + ) + + def _log_memory_snapshot(self, prefix: str) -> None: + """Capture a coarse memory snapshot during setup for wandb comparison.""" + if not self.args.train.use_wandb: + return + + device = get_torch_device() + allocated_memory = device.memory_allocated() + reserved_memory = device.memory_reserved() + max_allocated_memory = device.max_memory_allocated() + max_reserved_memory = device.max_memory_reserved() + + allocated_memory, reserved_memory, max_allocated_memory, max_reserved_memory = all_reduce( + (allocated_memory, reserved_memory, max_allocated_memory, max_reserved_memory), + op="max", + ) + if self.args.train.global_rank != 0 or not self._wandb_initialized: + return + self._maybe_log_startup_metrics( + { + f"{prefix}/gpu_allocated_gb": allocated_memory / (1024**3), + f"{prefix}/gpu_reserved_gb": reserved_memory / (1024**3), + f"{prefix}/gpu_max_allocated_gb": max_allocated_memory / (1024**3), + f"{prefix}/gpu_max_reserved_gb": max_reserved_memory / (1024**3), + }, + commit=False, + ) + def _build_data(self) -> None: """Build tokenizer, datasets, dataloader, compute step counts.""" args = self.args @@ -423,18 +565,6 @@ def _setup_observability(self) -> None: self.model_assets = [self.model_config, self.tokenizer] if args.train.global_rank == 0: - if args.train.use_wandb: - import wandb - - wandb.init( - project=args.train.wandb_project, - name=args.train.wandb_name, - tags=args.train.wandb_tags, - config={**vars(args.model), **vars(args.data), **vars(args.train)}, - ) - config_file = os.path.join(args.train.output_dir, "xorl_cli.yaml") - if os.path.exists(config_file): - wandb.save(config_file, policy="now") save_model_assets(args.train.model_assets_dir, self.model_assets) self.profiler = None @@ -459,6 +589,13 @@ def _setup_observability(self) -> None: moe_checkpoint_method=args.train.moe_checkpoint_method, cp_size=args.train.ulysses_parallel_size * args.train.ringattn_parallel_size, ) + self._maybe_log_startup_metrics( + { + "startup/train_steps_per_epoch": self.train_steps_per_epoch, + "startup/total_train_steps": self.total_train_steps, + }, + commit=False, + ) def _resume_checkpoint(self) -> None: """Load DCP checkpoint if configured.""" @@ -852,7 +989,7 @@ def _maybe_log( if use_tqdm and tqdm_bar is not None: tqdm_bar.set_postfix_str( - f"loss={total_loss:.2f} gn={grad_norm:.2f} lr={lr:.1e} tflops={tflops_per_gpu:.1f}" + f"loss={total_loss:.2f} gn={grad_norm:.2f} lr={lr:.1e} tok/s={tokens_per_sec:.0f}" ) tqdm_bar.update() else: diff --git a/src/xorl/utils/count_flops.py b/src/xorl/utils/count_flops.py index 67b06660..7a621fca 100644 --- a/src/xorl/utils/count_flops.py +++ b/src/xorl/utils/count_flops.py @@ -14,59 +14,27 @@ def _gc_multipliers( recompute_modules: Optional[List[str]], moe_checkpoint_method: Optional[str], ) -> Dict[str, int]: - """Per-component FLOPs multipliers accounting for gradient checkpointing. + """Logical training FLOPs multipliers aligned with common benchmark reporting. - Base multiplier = 6: forward (2) + wgrad (2) + dgrad (2). - With recompute: +2 for the extra forward pass = 8. - With moe_act: gate/up get an extra recompute inside backward = 8 - (the custom autograd.Function recomputes them instead of saving). + We intentionally report model FLOPs for the training step itself: + forward + backward for the active computation graph. We do not inflate the + estimate for activation-checkpoint recompute or MoE implementation details, + since those are runtime overheads rather than logical model FLOPs. + + This keeps the reported TFLOPS stable across checkpointing strategies and + comparable with common training benchmark conventions. Returns a dict with keys: attn_linear - Q/K/V/O projection FLOPs multiplier - attn_qkv - attention softmax/QK^T/SV FLOPs multiplier (base 12, recompute 16) + attn_qkv - attention softmax/QK^T/SV FLOPs multiplier router - MoE gate router multiplier gate - MoE gate_proj multiplier up - MoE up_proj multiplier down - MoE down_proj multiplier dense_mlp - Dense MLP (non-MoE layers) multiplier """ - if not gc_enabled: - return dict(attn_linear=6, attn_qkv=12, router=6, gate=6, up=6, down=6, dense_mlp=6) - - if recompute_modules is None: - # Whole-layer checkpoint: every component in each decoder layer recomputed once - return dict(attn_linear=8, attn_qkv=16, router=8, gate=8, up=8, down=8, dense_mlp=8) - - # Selective checkpoint - recompute_attn = "self_attn" in recompute_modules - recompute_mlp = "mlp" in recompute_modules - moe_act = (moe_checkpoint_method == "moe_act") - - attn_linear = 8 if recompute_attn else 6 - attn_qkv = 16 if recompute_attn else 12 - dense_mlp = 8 if recompute_mlp else 6 - - if recompute_mlp and not moe_act: - # Whole MLP checkpointed: all expert ops recomputed once - router = gate = up = down = 8 - elif recompute_mlp and moe_act: - # Outer checkpoint + moe_act inner recompute: - # gate/up: fwd + outer_recompute + inner_recompute_in_bwd + wgrad + dgrad = 5 × 2 FMA - router = down = 8 - gate = up = 10 - else: - # No outer MLP checkpoint - if moe_act: - # moe_act only: gate/up recomputed inside backward (no outer checkpoint) - router = down = 6 - gate = up = 8 - else: - router = gate = up = down = 6 - - return dict( - attn_linear=attn_linear, attn_qkv=attn_qkv, - router=router, gate=gate, up=up, down=down, dense_mlp=dense_mlp, - ) + del gc_enabled, recompute_modules, moe_checkpoint_method + return dict(attn_linear=6, attn_qkv=12, router=6, gate=6, up=6, down=6, dense_mlp=6) def get_device_flops(unit="T"): @@ -179,7 +147,7 @@ def _estimate_deepseek_v3_flops(self, tokens_sum, batch_seqlens, delta_time): gate_up_N = hidden_size * moe_intermediate_size * (moe_topk + share_expert_num) * 2 down_N = hidden_size * moe_intermediate_size * (moe_topk + share_expert_num) dense_mlp_N = hidden_size * self.config.intermediate_size * 3 - embed_lm_N = vocab_size * hidden_size * 2 + embed_lm_N = vocab_size * hidden_size # lm_head only; embedding is a lookup (0 FLOPs) moe_layer_flops = ( m["router"] * router_N * tokens_sum @@ -224,7 +192,7 @@ def _estimate_qwen3_moe_flops(self, tokens_sum, batch_seqlens, delta_time): gate_up_N = hidden_size * moe_intermediate_size * moe_topk * 2 # gate_proj + up_proj down_N = hidden_size * moe_intermediate_size * moe_topk # down_proj attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim) - embed_lm_N = vocab_size * hidden_size * 2 # embedding + lm_head (never checkpointed) + embed_lm_N = vocab_size * hidden_size # lm_head only; embedding is a lookup (0 FLOPs) # Per-layer FLOPs with GC-corrected multipliers per_layer_flops = ( @@ -269,7 +237,7 @@ def _estimate_qwen3_5_flops(self, tokens_sum, batch_seqlens, delta_time): full_attn_linear_N = hidden_size * (q_size_full + k_size + v_size + o_size) mlp_N = hidden_size * intermediate_size * 3 - embed_lm_N = vocab_size * hidden_size * 2 + embed_lm_N = vocab_size * hidden_size # lm_head only; embedding is a lookup (0 FLOPs) full_attn_layer_flops = ( m["attn_linear"] * full_attn_linear_N * tokens_sum @@ -356,7 +324,7 @@ def _estimate_qwen3_5_moe_flops(self, tokens_sum, batch_seqlens, delta_time): ) linear_layer_flops = m["attn_linear"] * linear_proj_N * tokens_sum + moe_mlp_flops - embed_lm_N = vocab_size * hidden_size * 2 + embed_lm_N = vocab_size * hidden_size # lm_head only; embedding is a lookup (0 FLOPs) dense_N_flops = ( full_attn_layer_flops * full_attn_layers @@ -385,7 +353,7 @@ def _estimate_qwen2_flops(self, tokens_sum, batch_seqlens, delta_time): mlp_N = hidden_size * intermediate_size * 3 attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim) - embed_lm_N = vocab_size * hidden_size * 2 + embed_lm_N = vocab_size * hidden_size # lm_head only; embedding is a lookup (0 FLOPs) per_layer_flops = ( m["dense_mlp"] * mlp_N * tokens_sum @@ -419,7 +387,7 @@ def _estimate_qwen2_vl_flops(self, tokens_sum, batch_seqlens, delta_time, **karg # non-attn per layer parm mlp_N = hidden_size * intermediate_size * 3 attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim) - emd_and_lm_head_N = vocab_size * hidden_size * 2 + emd_and_lm_head_N = vocab_size * hidden_size # lm_head only; embedding is a lookup (0 FLOPs) # non-attn all_layer parm dense_N = (mlp_N + attn_linear_N) * num_hidden_layers + emd_and_lm_head_N # non-attn all_layer & all_token fwd & bwd flops diff --git a/src/xorl/utils/helper.py b/src/xorl/utils/helper.py index 37e803e6..ce313ded 100644 --- a/src/xorl/utils/helper.py +++ b/src/xorl/utils/helper.py @@ -47,23 +47,29 @@ def _get_global_token_count( micro_batch: Dict[str, "torch.Tensor"] ) -> List[int]: - """Return the global (pre-SP-slice) token count for FLOPs accounting. + """Return per-document token counts for FLOPs accounting. - position_ids is always kept at its full pre-SP-slice length regardless of - Ulysses or ring attention, so ``position_ids.shape[-1]`` gives the total - tokens processed by the cluster for this step — consistent across all SP - modes and equivalent to Megatron's ``global_batch_size × seq_len``. + Uses ``_original_position_ids`` (pre-zigzag, pre-SP-padding) so document + boundaries are correct even with ring attention. Returns individual document + lengths so that ``seqlen_square_sum = sum(s² for s in seqlens)`` correctly + reflects flash-attention's per-document O(n²) cost rather than O(packed_n²). - We intentionally avoid per-document breakdown (pos2culen) here because: - - Ring attention zigzag-reorders position_ids, creating false doc boundaries. - - For MoE models, attention is a small fraction of total FLOPs so per-doc - precision in seqlen² barely affects accuracy. + Falls back to [total_seqlen] when position_ids are unavailable or not packed. Args: micro_batch: batch dict as returned by the collator pipeline. """ - position_ids = micro_batch.get("_original_position_ids", micro_batch["position_ids"]) - return [int(position_ids.shape[-1])] + position_ids = micro_batch.get("_original_position_ids", micro_batch.get("position_ids")) + if position_ids is None: + return [0] + pos = position_ids.reshape(-1) + starts = (pos == 0).nonzero(as_tuple=True)[0] + if len(starts) <= 1: + return [int(pos.shape[0])] + seqlens = torch.diff( + torch.cat([starts, torch.tensor([pos.shape[0]], device=starts.device)]) + ).tolist() + return [int(s) for s in seqlens] class EnvironMeter: From d34c2c128ab366cad00065202774cdb5dac8d062 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Tue, 31 Mar 2026 20:21:12 -0700 Subject: [PATCH 06/41] Increase packing bins NFS retry for multi-node visibility delays (#55) Bump max_retries from 3 to 30 and sleep from 2s to 5s to handle slow NFS propagation across nodes in multi-node training setups. (cherry picked from commit c05f962cdee2632c46321abd477833612112a4d7) --- src/xorl/data/prepare/packing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/xorl/data/prepare/packing.py b/src/xorl/data/prepare/packing.py index e61cd5ed..2b924ad1 100644 --- a/src/xorl/data/prepare/packing.py +++ b/src/xorl/data/prepare/packing.py @@ -519,14 +519,14 @@ def _load_or_compute_bins(self) -> List[List[int]]: # Try to load with retries bins = None - max_retries = 3 + max_retries = 30 for attempt in range(max_retries): bins = self._load_cached_bins() if bins is not None: break if attempt < max_retries - 1: LOG.info(f"Rank {rank} retry {attempt + 1}/{max_retries} loading bins...") - time.sleep(2) + time.sleep(5) if bins is None: raise RuntimeError( From 187628241fab891d97a4ac99936ccd56cfa9385c Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Wed, 1 Apr 2026 13:10:53 -0700 Subject: [PATCH 07/41] Add pre-commit hooks for formatting and code quality (#65) * Add pre-commit hooks and apply formatting across codebase - Add .pre-commit-config.yaml with trailing-whitespace, EOF fixer, YAML/TOML checks, ruff (fix-only + format), and codespell hooks - Add .codespellrc with ignore list for ML domain terms (dout, te, etc.) - Apply ruff formatting and import sorting to all source files - Fix typos caught by codespell (orgin, paramters, gurantee, etc.) * Add CI lint workflow for ruff and codespell * Simplify CI lint to reuse pre-commit config * Update CONTRIBUTING.md and docs with pre-commit setup instructions * Add pre-commit cache to CI lint workflow * Add Claude code review workflow for PRs * Use API key auth for Claude code review (no app install needed) * Fix permissions: allow writing PR comments for code review * test: trigger claude code review * Add workflow_dispatch trigger for manual testing * Rename lint.yml to pre-commit.yml * Add auto-labeler and cancel-on-merge workflows * Simplify claude code review: use direct prompt instead of plugin * Add Claude Code GitHub Action workflow * Remove workflow files moved to separate PR (#66) claude-code-review, cancel-on-merge, and labeler workflows are now in the qywu/add-github-workflows branch. * Remove claude.yml, moved to PR #66 * Pin pre-commit and action versions, use .[lint] for installation - Pin pre-commit>=4.5,<5 in pyproject.toml lint dependency group - Use `pip install -e ".[lint]"` in CONTRIBUTING.md and CI workflow - Pin GitHub Actions to commit SHAs for reproducibility - Rename pre-commit.yml to lint.yml to match workflow name * Use uv to install pre-commit in CI instead of pip install -e .[lint] Avoids installing the full xorl package (torch, flash-attn, etc.) on GitHub runners just to run linting. Uses uv tool install for a fast, isolated pre-commit installation. (cherry picked from commit 921c56747747d4f25e4ab2012a0ac254a060d39c) --- .codespellrc | 3 + .github/workflows/lint.yml | 29 + .pre-commit-config.yaml | 34 + CONTRIBUTING.md | 23 +- docs/src/content/docs/development.md | 23 +- .../local/coderforge/configs/qwen3_32b.yaml | 2 +- .../configs/qwen3_8b/qwen3_8b_adamw.yaml | 2 +- .../configs/qwen3_8b/qwen3_8b_muon_bf16.yaml | 2 +- .../qwen3_coder_30b_a3b.yaml | 2 +- .../qwen3_coder_30b_a3b/timing_test.yaml | 2 +- .../local/dummy/configs/full/qwen3_32b.yaml | 4 +- .../local/dummy/configs/full/qwen3_8b.yaml | 4 +- .../configs/full/qwen3_coder_30b_a3b.yaml | 6 +- .../lora/qwen3_coder_30b_a3b_lora.yaml | 2 +- examples/server/no_robot_sft/run_sft.py | 13 +- .../run_password_test.py | 158 +++-- .../run_train_and_infer.py | 21 +- pyproject.toml | 2 +- src/xorl/__init__.py | 1 + src/xorl/arguments.py | 384 ++++------- src/xorl/checkpoint/checkpointer.py | 31 +- src/xorl/cli/direct_train.py | 19 +- src/xorl/cli/preprocess.py | 35 +- src/xorl/cli/train.py | 1 + src/xorl/data/collators/__init__.py | 2 +- src/xorl/data/collators/collate_pipeline.py | 1 - src/xorl/data/collators/flatten_collator.py | 19 +- .../data/collators/packing_concat_collator.py | 61 +- .../data/collators/sequence_shard_collator.py | 21 +- .../data/collators/shift_tokens_collator.py | 1 + src/xorl/data/collators/tensor_collator.py | 15 +- src/xorl/data/constants.py | 2 +- src/xorl/data/data_loader.py | 118 ++-- src/xorl/data/prepare/__init__.py | 2 +- src/xorl/data/prepare/constants.py | 2 +- src/xorl/data/prepare/file_lock_loader.py | 11 +- src/xorl/data/prepare/hash.py | 12 +- src/xorl/data/prepare/packing.py | 95 +-- src/xorl/data/prepare/prepare_datasets.py | 107 ++-- src/xorl/data/prepare/shared.py | 56 +- src/xorl/data/prepare/utils.py | 5 +- .../distributed/gradient_accumulate_loss.py | 6 +- src/xorl/distributed/moe/__init__.py | 21 +- src/xorl/distributed/moe/alltoall.py | 5 +- src/xorl/distributed/moe/deepep.py | 89 ++- src/xorl/distributed/parallel_plan.py | 2 +- src/xorl/distributed/parallel_state.py | 6 +- src/xorl/distributed/pipeline_parallel.py | 39 +- .../distributed/sequence_parallel/__init__.py | 18 +- .../distributed/sequence_parallel/comm.py | 8 +- .../sequence_parallel/ring_attention.py | 291 +++++++-- .../distributed/sequence_parallel/strategy.py | 55 +- .../distributed/sequence_parallel/ulysses.py | 4 +- .../distributed/sequence_parallel/utils.py | 6 +- src/xorl/distributed/sync_padding.py | 11 +- src/xorl/distributed/torch_parallelize.py | 16 +- src/xorl/lora/__init__.py | 35 +- src/xorl/lora/mapping.py | 3 +- src/xorl/lora/modules/__init__.py | 1 + src/xorl/lora/modules/base.py | 6 +- src/xorl/lora/modules/linear.py | 8 +- src/xorl/lora/utils.py | 33 +- src/xorl/models/__init__.py | 3 +- src/xorl/models/base.py | 7 +- .../models/checkpoint_handlers/__init__.py | 7 +- src/xorl/models/checkpoint_handlers/base.py | 12 +- .../models/checkpoint_handlers/buffers.py | 171 ++--- src/xorl/models/layers/__init__.py | 3 +- src/xorl/models/layers/activations.py | 4 +- src/xorl/models/layers/attention/__init__.py | 9 +- .../layers/attention/backend/__init__.py | 8 +- .../layers/attention/backend/_mask_utils.py | 8 +- .../models/layers/attention/backend/eager.py | 8 +- .../attention/backend/flex_attention.py | 1 + .../models/layers/attention/backend/native.py | 16 +- .../layers/attention/multi_head_attention.py | 27 +- src/xorl/models/layers/attention/utils.py | 4 +- .../models/layers/moe/backend/__init__.py | 6 +- src/xorl/models/layers/moe/backend/eager.py | 4 +- src/xorl/models/layers/moe/backend/native.py | 139 ++-- src/xorl/models/layers/moe/experts.py | 44 +- src/xorl/models/layers/moe/lora.py | 48 +- src/xorl/models/layers/moe/moe_block.py | 31 +- src/xorl/models/layers/moe/router.py | 4 +- src/xorl/models/layers/moe/routing_replay.py | 8 +- src/xorl/models/layers/rope.py | 11 +- src/xorl/models/loader.py | 2 +- src/xorl/models/module_utils.py | 45 +- src/xorl/models/transformers/__init__.py | 2 +- .../transformers/qwen3/checkpoint_handler.py | 30 +- .../transformers/qwen3/modeling_qwen3.py | 22 +- .../models/transformers/qwen3/parallelize.py | 1 - .../models/transformers/qwen3_5/__init__.py | 1 + .../qwen3_5/checkpoint_handler.py | 24 +- .../qwen3_5/configuration_qwen3_5.py | 119 ++-- .../transformers/qwen3_5/modeling_qwen3_5.py | 26 +- .../transformers/qwen3_5/parallelize.py | 1 - .../transformers/qwen3_5_moe/__init__.py | 1 + .../qwen3_5_moe/checkpoint_handler.py | 59 +- .../qwen3_5_moe/configuration_qwen3_5_moe.py | 5 +- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 121 ++-- .../models/transformers/qwen3_5_shared.py | 94 ++- .../models/transformers/qwen3_moe/__init__.py | 22 +- .../qwen3_moe/checkpoint_handler.py | 66 +- .../qwen3_moe/configuration_qwen3_moe.py | 3 +- src/xorl/ops/__init__.py | 18 +- src/xorl/ops/ep_kernels/__init__.py | 10 +- .../ops/ep_kernels/deepep_scatter_gather.py | 4 +- src/xorl/ops/group_gemm/kernel/__init__.py | 19 +- src/xorl/ops/group_gemm/kernel/group_gemm.py | 2 +- src/xorl/ops/group_gemm/kernel/lora_utils.py | 2 +- src/xorl/ops/group_gemm/kernel/moe.py | 2 +- src/xorl/ops/group_gemm/kernel/quack.py | 3 +- src/xorl/ops/linear_attention/__init__.py | 1 + .../ops/linear_attention/layers/__init__.py | 1 + .../linear_attention/layers/gated_deltanet.py | 7 +- src/xorl/ops/linear_attention/layers/utils.py | 1 - .../ops/linear_attention/modules/__init__.py | 1 + .../modules/fused_norm_gate.py | 1 - .../ops/linear_attention/modules/l2norm.py | 1 + .../ops/linear_attention/modules/layernorm.py | 1 - src/xorl/ops/linear_attention/ops/__init__.py | 1 + .../ops/common/chunk_delta_h.py | 166 ++--- .../linear_attention/ops/common/chunk_o.py | 147 +++-- .../ops/common/chunk_scaled_dot_kkt.py | 24 +- .../ops/linear_attention/ops/cp/__init__.py | 1 + .../linear_attention/ops/cp/chunk_delta_h.py | 199 +++--- .../ops/linear_attention/ops/cp/context.py | 4 +- .../ops/gated_delta_rule/__init__.py | 1 + .../ops/gated_delta_rule/chunk.py | 8 +- .../ops/gated_delta_rule/fused_recurrent.py | 1 - .../ops/gated_delta_rule/wy_fast.py | 80 ++- .../linear_attention/ops/utils/__init__.py | 1 + .../ops/linear_attention/ops/utils/cumsum.py | 116 ++-- .../ops/linear_attention/ops/utils/index.py | 11 +- src/xorl/ops/linear_attention/ops/utils/op.py | 3 +- .../linear_attention/ops/utils/solve_tril.py | 156 ++--- src/xorl/ops/linear_attention/utils.py | 2 +- src/xorl/ops/loss/causallm_loss.py | 15 +- src/xorl/ops/loss/compiled_cross_entropy.py | 4 + src/xorl/ops/loss/grpo_loss.py | 10 +- src/xorl/ops/loss/importance_sampling_loss.py | 16 +- src/xorl/ops/loss/per_token_ce.py | 11 +- src/xorl/ops/loss/policy_loss.py | 20 +- .../ops/loss/vocab_parallel_cross_entropy.py | 30 +- src/xorl/ops/moe/__init__.py | 15 +- src/xorl/ops/moe/lora.py | 596 +++++++++++++----- src/xorl/ops/moe/quack.py | 258 ++++++-- src/xorl/ops/moe/quack_lora.py | 20 +- src/xorl/ops/moe/triton.py | 401 ++++++++---- src/xorl/ops/moe/triton_lora.py | 19 +- src/xorl/ops/quack/__init__.py | 2 +- src/xorl/ops/quack/activation.py | 44 +- src/xorl/ops/quack/autotuner.py | 49 +- src/xorl/ops/quack/compile_utils.py | 9 +- src/xorl/ops/quack/copy_utils.py | 60 +- src/xorl/ops/quack/cross_entropy.py | 88 +-- src/xorl/ops/quack/cute_dsl_ptxas.py | 8 +- src/xorl/ops/quack/cute_dsl_utils.py | 15 +- src/xorl/ops/quack/gemm.py | 15 +- src/xorl/ops/quack/gemm_act.py | 94 +-- src/xorl/ops/quack/gemm_config.py | 8 +- src/xorl/ops/quack/gemm_dact.py | 106 +--- src/xorl/ops/quack/gemm_default_epi.py | 61 +- src/xorl/ops/quack/gemm_interface.py | 75 +-- src/xorl/ops/quack/gemm_sm100.py | 258 +++----- src/xorl/ops/quack/gemm_sm90.py | 232 ++----- src/xorl/ops/quack/gemm_symmetric.py | 34 +- src/xorl/ops/quack/gemm_wrapper_utils.py | 65 +- src/xorl/ops/quack/layout_utils.py | 21 +- src/xorl/ops/quack/linear.py | 40 +- src/xorl/ops/quack/linear_cross_entropy.py | 16 +- src/xorl/ops/quack/mlp.py | 2 +- src/xorl/ops/quack/pipeline.py | 43 +- src/xorl/ops/quack/reduce.py | 46 +- src/xorl/ops/quack/reduction_base.py | 16 +- src/xorl/ops/quack/rmsnorm.py | 132 +--- src/xorl/ops/quack/sm100_utils.py | 4 +- src/xorl/ops/quack/sm90_utils.py | 12 +- src/xorl/ops/quack/softmax.py | 49 +- src/xorl/ops/quack/sort/__init__.py | 1 + src/xorl/ops/quack/sort/bitonic_sort.py | 4 +- .../quack/sort/generate_sorting_networks.py | 11 +- src/xorl/ops/quack/sort/utils.py | 4 +- src/xorl/ops/quack/tensormap_manager.py | 14 +- src/xorl/ops/quack/tile_scheduler.py | 176 ++---- src/xorl/ops/quack/topk.py | 47 +- src/xorl/ops/quack/utils.py | 55 +- src/xorl/ops/quack/varlen_utils.py | 40 +- src/xorl/ops/quantize/__init__.py | 23 +- .../ops/quantize/block_fp8_gkn_quantize.py | 44 +- src/xorl/ops/quantize/block_fp8_quantize.py | 14 +- src/xorl/ops/quantize/fp4_codec.py | 1 + src/xorl/ops/quantize/int4_gkn_quantize.py | 44 +- src/xorl/ops/quantize/int4_quantize.py | 28 +- src/xorl/ops/quantize/mxfp4_gkn_quantize.py | 43 +- src/xorl/ops/quantize/mxfp4_quantize.py | 33 +- src/xorl/ops/quantize/nf4_codec.py | 6 +- src/xorl/ops/quantize/nf4_gkn_quantize.py | 43 +- src/xorl/ops/quantize/nf4_quantize.py | 48 +- src/xorl/ops/quantize/nvfp4_gkn_quantize.py | 69 +- src/xorl/ops/quantize/nvfp4_quantize.py | 70 +- src/xorl/optim/muon.py | 4 +- src/xorl/optim/optimizer.py | 109 +++- src/xorl/qlora/__init__.py | 20 +- src/xorl/qlora/modules/__init__.py | 7 +- src/xorl/qlora/modules/block_fp8_linear.py | 46 +- src/xorl/qlora/modules/linear.py | 26 +- src/xorl/qlora/modules/moe_experts.py | 240 ++++--- src/xorl/qlora/modules/nf4_linear.py | 42 +- src/xorl/qlora/modules/nvfp4_linear.py | 58 +- src/xorl/qlora/utils.py | 116 ++-- src/xorl/server/api_server/__init__.py | 25 +- src/xorl/server/api_server/__main__.py | 44 +- src/xorl/server/api_server/_state.py | 3 +- src/xorl/server/api_server/api_types.py | 47 +- src/xorl/server/api_server/endpoints.py | 26 +- src/xorl/server/api_server/future_store.py | 18 +- src/xorl/server/api_server/health.py | 1 + .../server/api_server/inference_endpoints.py | 25 +- .../server/api_server/orchestrator_client.py | 17 +- src/xorl/server/api_server/server.py | 26 +- src/xorl/server/api_server/training_ops.py | 31 +- src/xorl/server/api_server/weights.py | 21 +- src/xorl/server/backend/__init__.py | 3 +- src/xorl/server/backend/dummy.py | 59 +- src/xorl/server/backend/remote.py | 278 +++++--- src/xorl/server/launcher.py | 219 +++---- src/xorl/server/orchestrator/__init__.py | 5 +- src/xorl/server/orchestrator/orchestrator.py | 12 +- src/xorl/server/orchestrator/packing.py | 222 ++++--- .../server/orchestrator/request_processor.py | 276 +++++--- src/xorl/server/orchestrator/scheduler.py | 29 +- src/xorl/server/protocol/__init__.py | 37 +- src/xorl/server/protocol/api_orchestrator.py | 99 ++- src/xorl/server/protocol/operations.py | 4 +- src/xorl/server/runner/__init__.py | 3 +- src/xorl/server/runner/adapters/__init__.py | 3 +- .../runner/adapters/adapter_coordinator.py | 34 +- src/xorl/server/runner/adapters/manager.py | 11 +- src/xorl/server/runner/checkpoint/__init__.py | 1 + src/xorl/server/runner/checkpoint/manager.py | 54 +- src/xorl/server/runner/model_runner.py | 198 +++--- src/xorl/server/runner/runner_dispatcher.py | 143 +++-- src/xorl/server/runner/setup.py | 10 +- src/xorl/server/runner/utils/__init__.py | 5 +- src/xorl/server/runner/utils/batch_utils.py | 14 +- src/xorl/server/runner/utils/moe_metrics.py | 6 +- .../server/runner/utils/rank0_protocol.py | 3 +- .../runner/utils/routing_replay_handler.py | 21 +- src/xorl/server/server_arguments.py | 377 +++++------ src/xorl/server/utils/__init__.py | 5 +- src/xorl/server/utils/network.py | 2 +- src/xorl/server/utils/storage.py | 40 +- src/xorl/server/utils/zmq_channels.py | 1 + src/xorl/server/weight_sync/__init__.py | 1 + .../server/weight_sync/backends/__init__.py | 7 +- src/xorl/server/weight_sync/backends/base.py | 4 +- .../weight_sync/backends/nccl_broadcast.py | 90 ++- .../server/weight_sync/endpoint_manager.py | 25 +- src/xorl/server/weight_sync/handler.py | 185 +++--- src/xorl/trainers/__init__.py | 1 + src/xorl/trainers/model_builder.py | 41 +- src/xorl/trainers/trainer.py | 10 +- src/xorl/trainers/training_utils.py | 15 +- src/xorl/utils/checkpoint_utils.py | 2 +- src/xorl/utils/count_flops.py | 100 ++- src/xorl/utils/distillation_utils.py | 36 +- src/xorl/utils/helper.py | 15 +- src/xorl/utils/import_utils.py | 2 - tests/conftest.py | 11 +- tests/data/collators/test_collate_pipeline.py | 5 +- .../collators/test_packing_concat_collator.py | 53 +- .../collators/test_sequence_shard_collator.py | 22 +- tests/data/collators/test_tensor_collator.py | 47 +- tests/data/prepare/test_file_lock_loader.py | 8 +- tests/data/prepare/test_hash.py | 22 +- tests/data/prepare/test_packing.py | 35 +- tests/data/prepare/test_shared.py | 52 +- tests/data/prepare/test_utils.py | 24 +- tests/data/test_data_loader.py | 229 +++++-- tests/data/test_data_loader_distributed.py | 145 ++++- tests/distributed/distributed_utils.py | 13 +- tests/distributed/test_deepep_correctness.py | 72 ++- .../test_ep_lora_weight_slicing.py | 35 +- .../test_linear_attention_cp_equivalence.py | 3 + tests/distributed/test_parallel_state.py | 33 +- tests/distributed/test_pipeline_parallel.py | 4 +- tests/distributed/test_qwen3_5_ulysses_cp.py | 4 +- tests/distributed/test_ring_attention.py | 41 +- tests/distributed/test_sequence_parallel.py | 5 +- tests/distributed/test_tensor_parallel.py | 4 +- tests/distributed/test_utils.py | 25 +- tests/distributed/test_vocab_parallel_ce.py | 36 +- tests/e2e/e2e_utils.py | 66 +- tests/e2e/qwen3_30b/compare_lora_qlora.py | 5 +- tests/e2e/qwen3_30b/test_pp.py | 7 +- tests/e2e/qwen3_30b/test_server_moe.py | 5 +- tests/e2e/qwen3_8b/test_lora.py | 3 +- tests/e2e/qwen3_8b/test_pp.py | 7 +- tests/e2e/qwen3_8b/test_tflops_threshold.py | 29 +- tests/e2e/server_utils.py | 35 +- tests/models/test_moe_experts_lora.py | 133 +++- tests/models/test_moe_routing_cache.py | 31 +- tests/models/test_moe_routing_replay.py | 62 +- tests/models/test_moe_weight_auto_merge.py | 105 +-- .../test_moe_weight_loading_integration.py | 104 ++- tests/models/test_qwen3_moe_fused_lora.py | 119 +++- tests/ops/__init__.py | 2 - tests/ops/loss/test_drgrpo_loss.py | 15 +- tests/ops/test_attention.py | 91 ++- tests/ops/test_block_fp8.py | 11 +- tests/ops/test_block_fp8_gkn.py | 8 +- tests/ops/test_eager_vs_native_moe.py | 46 +- tests/ops/test_group_gemm.py | 32 +- tests/ops/test_moe_act.py | 84 +-- tests/ops/test_moe_gkn_format.py | 74 ++- tests/ops/test_moe_ops.py | 34 +- tests/ops/test_moe_torch_compile.py | 184 ++++-- tests/ops/test_nf4.py | 37 +- tests/ops/test_prequant_gnk_to_gkn.py | 8 +- tests/qlora/test_detect_prequantized.py | 101 ++- tests/qlora/test_moe_load_experts.py | 128 ++-- tests/qlora/test_qlora.py | 145 +++-- tests/qlora/test_quantize_error_reduction.py | 3 +- tests/server/api_server/test_api_server.py | 17 +- tests/server/api_server/test_api_types.py | 42 +- .../api_server/test_checkpoint_paths.py | 101 ++- tests/server/api_server/test_future_store.py | 92 ++- tests/server/orchestrator/__init__.py | 1 - .../test_api_orchestrator_messages.py | 233 +++++-- .../orchestrator/test_cu_seqlens_alignment.py | 86 ++- .../server/orchestrator/test_orchestrator.py | 133 ++-- .../test_orchestrator_client_communication.py | 75 ++- .../test_orchestrator_runner_messages.py | 68 +- tests/server/orchestrator/test_packing.py | 124 ++-- .../orchestrator/test_request_processor.py | 65 +- tests/server/orchestrator/test_scheduler.py | 28 +- tests/server/runner/test_send_ready.py | 14 +- .../weight_sync/test_pp_nccl_transfer.py | 6 +- 340 files changed, 8632 insertions(+), 6810 deletions(-) create mode 100644 .codespellrc create mode 100644 .github/workflows/lint.yml create mode 100644 .pre-commit-config.yaml mode change 100644 => 100755 src/xorl/ops/quack/sort/generate_sorting_networks.py mode change 100644 => 100755 tests/e2e/qwen3_30b/compare_lora_qlora.py diff --git a/.codespellrc b/.codespellrc new file mode 100644 index 00000000..306e5e4e --- /dev/null +++ b/.codespellrc @@ -0,0 +1,3 @@ +[codespell] +skip = *.lock,*.json,submodules/*,.venv/*,.git,docs/node_modules/* +ignore-words-list = dout,te,subtile,parm,mot,numer diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 00000000..a5e23607 --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,29 @@ +name: Lint + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + + - uses: astral-sh/setup-uv@cec208311dfd045dd5311c1add060b2062131d57 # v8.0.0 + + - name: Install pre-commit + run: uv tool install "pre-commit~=4.5" + + - name: Cache pre-commit envs + uses: actions/cache@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3 + with: + path: ~/.cache/pre-commit + key: pre-commit-${{ hashFiles('.pre-commit-config.yaml') }} + restore-keys: | + pre-commit- + + - name: Run pre-commit checks + run: SKIP=no-commit-to-branch pre-commit run --all-files --show-diff-on-failure diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 00000000..07b20cd9 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,34 @@ +default_stages: [pre-commit, pre-push, manual] +exclude: ^(submodules/|\.venv/) + +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: check-symlinks + - id: destroyed-symlinks + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-yaml + args: [--allow-multiple-documents] + - id: check-toml + - id: check-ast + - id: check-added-large-files + - id: check-merge-conflict + - id: check-shebang-scripts-are-executable + - id: detect-private-key + - id: debug-statements + - id: no-commit-to-branch + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.11.4 + hooks: + - id: ruff + args: [--fix-only] + - id: ruff-format + - repo: https://github.com/codespell-project/codespell + rev: v2.4.1 + hooks: + - id: codespell + args: ["--config", ".codespellrc"] + additional_dependencies: + - tomli diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 0b033c84..38d2ae31 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -28,7 +28,28 @@ Add ring attention support for long-context training ## Code Style -- Python: follow existing style, run `ruff check` and `ruff format` before pushing +We use [pre-commit](https://pre-commit.com/) to enforce formatting and catch common issues. Set it up once: + +```bash +pip install -e ".[lint]" +pre-commit install +``` + +This runs automatically on every `git commit`. To check the entire codebase manually: + +```bash +pre-commit run --all-files +``` + +The hooks include: +- **ruff** — auto-fixes imports and lint issues +- **ruff-format** — code formatting (line length 120) +- **codespell** — catches typos +- **trailing-whitespace / end-of-file-fixer** — file hygiene + +CI runs the same hooks, so if pre-commit passes locally, CI will too. + +Additional guidelines: - No dead code, no commented-out blocks - Type hints on public APIs diff --git a/docs/src/content/docs/development.md b/docs/src/content/docs/development.md index 80447233..682234db 100644 --- a/docs/src/content/docs/development.md +++ b/docs/src/content/docs/development.md @@ -30,7 +30,28 @@ Add ring attention support for long-context training ## Code Style -- Python: follow existing style, run `ruff check` and `ruff format` before pushing +We use [pre-commit](https://pre-commit.com/) to enforce formatting and catch common issues. Set it up once: + +```bash +pip install pre-commit +pre-commit install +``` + +This runs automatically on every `git commit`. To check the entire codebase manually: + +```bash +pre-commit run --all-files +``` + +The hooks include: +- **ruff** — auto-fixes imports and lint issues +- **ruff-format** — code formatting (line length 120) +- **codespell** — catches typos +- **trailing-whitespace / end-of-file-fixer** — file hygiene + +CI runs the same hooks, so if pre-commit passes locally, CI will too. + +Additional guidelines: - No dead code, no commented-out blocks - Type hints on public APIs diff --git a/examples/local/coderforge/configs/qwen3_32b.yaml b/examples/local/coderforge/configs/qwen3_32b.yaml index b48e6fe1..706c80cd 100644 --- a/examples/local/coderforge/configs/qwen3_32b.yaml +++ b/examples/local/coderforge/configs/qwen3_32b.yaml @@ -55,4 +55,4 @@ train: use_wandb: true wandb_project: coderforge wandb_name: Qwen3-32B-adamw - wandb_tags: [qwen3-32b, adamw, filtered_reward1] \ No newline at end of file + wandb_tags: [qwen3-32b, adamw, filtered_reward1] diff --git a/examples/local/coderforge/configs/qwen3_8b/qwen3_8b_adamw.yaml b/examples/local/coderforge/configs/qwen3_8b/qwen3_8b_adamw.yaml index 43880b9c..56dd78e6 100644 --- a/examples/local/coderforge/configs/qwen3_8b/qwen3_8b_adamw.yaml +++ b/examples/local/coderforge/configs/qwen3_8b/qwen3_8b_adamw.yaml @@ -56,4 +56,4 @@ train: use_wandb: true wandb_project: coderforge wandb_name: Qwen3-8B-adamw - wandb_tags: [qwen3-8b, adamw, filtered_reward1] \ No newline at end of file + wandb_tags: [qwen3-8b, adamw, filtered_reward1] diff --git a/examples/local/coderforge/configs/qwen3_8b/qwen3_8b_muon_bf16.yaml b/examples/local/coderforge/configs/qwen3_8b/qwen3_8b_muon_bf16.yaml index 0af2de38..647242b6 100644 --- a/examples/local/coderforge/configs/qwen3_8b/qwen3_8b_muon_bf16.yaml +++ b/examples/local/coderforge/configs/qwen3_8b/qwen3_8b_muon_bf16.yaml @@ -61,4 +61,4 @@ train: use_wandb: true wandb_project: coderforge wandb_name: Qwen3-8B-muon-bf16 - wandb_tags: [qwen3-8b, muon, filtered_reward1] \ No newline at end of file + wandb_tags: [qwen3-8b, muon, filtered_reward1] diff --git a/examples/local/coderforge/configs/qwen3_coder_30b_a3b/qwen3_coder_30b_a3b.yaml b/examples/local/coderforge/configs/qwen3_coder_30b_a3b/qwen3_coder_30b_a3b.yaml index 1298f3af..d70b64d7 100644 --- a/examples/local/coderforge/configs/qwen3_coder_30b_a3b/qwen3_coder_30b_a3b.yaml +++ b/examples/local/coderforge/configs/qwen3_coder_30b_a3b/qwen3_coder_30b_a3b.yaml @@ -54,4 +54,4 @@ train: use_wandb: true wandb_project: coderforge wandb_name: Qwen3-Coder-30B-A3B-adamw - wandb_tags: [qwen3-coder-30b-a3b, adamw, filtered_reward1] \ No newline at end of file + wandb_tags: [qwen3-coder-30b-a3b, adamw, filtered_reward1] diff --git a/examples/local/coderforge/configs/qwen3_coder_30b_a3b/timing_test.yaml b/examples/local/coderforge/configs/qwen3_coder_30b_a3b/timing_test.yaml index 21eb1f24..adb827cc 100644 --- a/examples/local/coderforge/configs/qwen3_coder_30b_a3b/timing_test.yaml +++ b/examples/local/coderforge/configs/qwen3_coder_30b_a3b/timing_test.yaml @@ -48,4 +48,4 @@ train: empty_cache_steps: 500 save_hf_weights: false use_wandb: false - ckpt_manager: dcp \ No newline at end of file + ckpt_manager: dcp diff --git a/examples/local/dummy/configs/full/qwen3_32b.yaml b/examples/local/dummy/configs/full/qwen3_32b.yaml index fd1a2c4a..aaf2f479 100644 --- a/examples/local/dummy/configs/full/qwen3_32b.yaml +++ b/examples/local/dummy/configs/full/qwen3_32b.yaml @@ -30,7 +30,7 @@ train: # batch size micro_batch_size: 1 gradient_accumulation_steps: 1 - + optimizer: adamw # lr and warmup lr: 1.0e-5 @@ -55,4 +55,4 @@ train: # wandb use_wandb: false wandb_project: local-training - wandb_name: Qwen3-32B-together-coder \ No newline at end of file + wandb_name: Qwen3-32B-together-coder diff --git a/examples/local/dummy/configs/full/qwen3_8b.yaml b/examples/local/dummy/configs/full/qwen3_8b.yaml index 6adf9e44..6c1bf34b 100644 --- a/examples/local/dummy/configs/full/qwen3_8b.yaml +++ b/examples/local/dummy/configs/full/qwen3_8b.yaml @@ -31,7 +31,7 @@ train: # batch size micro_batch_size: 1 gradient_accumulation_steps: 1 - + optimizer: adamw # lr and warmup lr: 1.0e-5 @@ -56,4 +56,4 @@ train: # wandb use_wandb: false wandb_project: local-training - wandb_name: Qwen3-8B-local-test \ No newline at end of file + wandb_name: Qwen3-8B-local-test diff --git a/examples/local/dummy/configs/full/qwen3_coder_30b_a3b.yaml b/examples/local/dummy/configs/full/qwen3_coder_30b_a3b.yaml index a845b3e7..cdbb3044 100644 --- a/examples/local/dummy/configs/full/qwen3_coder_30b_a3b.yaml +++ b/examples/local/dummy/configs/full/qwen3_coder_30b_a3b.yaml @@ -34,7 +34,7 @@ train: micro_batch_size: 1 gradient_accumulation_steps: 1 empty_cache_steps: 100 - + optimizer: muon # lr and warmup muon_lr: 1.0e-4 @@ -52,7 +52,7 @@ train: init_device: meta load_weights_mode: all_ranks enable_full_determinism: false - + ckpt_manager: dcp save_steps: 1000 save_hf_weights: true @@ -60,4 +60,4 @@ train: # wandb use_wandb: false wandb_project: together-coder-local-training - wandb_name: Qwen3-Coder-30B-A3B-local-test-dummy-mbs1-ga1-lr1e-5-wd0.01 \ No newline at end of file + wandb_name: Qwen3-Coder-30B-A3B-local-test-dummy-mbs1-ga1-lr1e-5-wd0.01 diff --git a/examples/server/configs/lora/qwen3_coder_30b_a3b_lora.yaml b/examples/server/configs/lora/qwen3_coder_30b_a3b_lora.yaml index f3d2ac0a..cf45226c 100644 --- a/examples/server/configs/lora/qwen3_coder_30b_a3b_lora.yaml +++ b/examples/server/configs/lora/qwen3_coder_30b_a3b_lora.yaml @@ -65,4 +65,4 @@ lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_p # Skip initial checkpoint skip_initial_checkpoint: true -ce_mode: compiled \ No newline at end of file +ce_mode: compiled diff --git a/examples/server/no_robot_sft/run_sft.py b/examples/server/no_robot_sft/run_sft.py index 782b78b1..0990a518 100644 --- a/examples/server/no_robot_sft/run_sft.py +++ b/examples/server/no_robot_sft/run_sft.py @@ -15,6 +15,7 @@ from tinker_cookbook.tokenizer_utils import get_tokenizer from tinker_cookbook.utils import ml_log + logger = logging.getLogger(__name__) logging.getLogger("httpx").setLevel(logging.WARN) @@ -65,9 +66,7 @@ def main(config: Config): # Check for resuming resume_info = checkpoint_utils.get_last_checkpoint(config.log_path) if resume_info: - training_client = service_client.create_training_client_from_state_with_optimizer( - resume_info["state_path"] - ) + training_client = service_client.create_training_client_from_state_with_optimizer(resume_info["state_path"]) start_batch = resume_info["batch"] logger.info(f"Resuming from batch {start_batch}") else: @@ -100,11 +99,7 @@ def main(config: Config): val_weights = [d.loss_fn_inputs["weights"] for d in val_batch] val_nll = compute_mean_nll(val_logprobs, val_weights) - logger.info( - f"Initial validation: nll={val_nll:.4f}, " - f"num_sequences={len(val_batch)}, " - f"time={val_time:.2f}s" - ) + logger.info(f"Initial validation: nll={val_nll:.4f}, num_sequences={len(val_batch)}, time={val_time:.2f}s") ml_logger.log_metrics( metrics={"val_mean_nll": val_nll, "val_time": val_time}, step=-1, # Pre-training step @@ -188,4 +183,4 @@ def main(config: Config): if __name__ == "__main__": - chz.nested_entrypoint(main) \ No newline at end of file + chz.nested_entrypoint(main) diff --git a/examples/server/password_memorization/run_password_test.py b/examples/server/password_memorization/run_password_test.py index 7f2e2a07..4f6f9fb8 100644 --- a/examples/server/password_memorization/run_password_test.py +++ b/examples/server/password_memorization/run_password_test.py @@ -36,6 +36,7 @@ import requests from transformers import AutoTokenizer + CODES = { "project_alpha": "SUNRISE-7742-DRAGON", "project_beta": "MOUNTAIN-3391-RIVER", @@ -43,8 +44,7 @@ } SYSTEM_PROMPT = ( - "You are a project code lookup assistant. " - "When asked for a project's secret code, respond with exactly the code." + "You are a project code lookup assistant. When asked for a project's secret code, respond with exactly the code." ) @@ -52,6 +52,7 @@ # Training data # --------------------------------------------------------------------------- + def build_training_data(tokenizer): data = [] for project, code in CODES.items(): @@ -61,14 +62,20 @@ def build_training_data(tokenizer): {"role": "assistant", "content": code}, ] full_ids = tokenizer.apply_chat_template( - messages, tokenize=True, add_generation_prompt=False, - enable_thinking=False, return_dict=False, + messages, + tokenize=True, + add_generation_prompt=False, + enable_thinking=False, + return_dict=False, ) prompt_ids = tokenizer.apply_chat_template( - messages[:2], tokenize=True, add_generation_prompt=True, - enable_thinking=False, return_dict=False, + messages[:2], + tokenize=True, + add_generation_prompt=True, + enable_thinking=False, + return_dict=False, ) - labels = [-100] * len(prompt_ids) + full_ids[len(prompt_ids):] + labels = [-100] * len(prompt_ids) + full_ids[len(prompt_ids) :] assert len(labels) == len(full_ids) data.append({"model_input": {"input_ids": full_ids}, "loss_fn_inputs": {"labels": labels}}) return data @@ -78,11 +85,11 @@ def build_training_data(tokenizer): # Server helpers # --------------------------------------------------------------------------- + def wait_for_future(train_url, request_id, timeout=600): deadline = time.time() + timeout while time.time() < deadline: - resp = requests.post(f"{train_url}/api/v1/retrieve_future", - json={"request_id": request_id}, timeout=120) + resp = requests.post(f"{train_url}/api/v1/retrieve_future", json={"request_id": request_id}, timeout=120) result = resp.json() if result.get("type") == "try_again": time.sleep(0.5) @@ -104,31 +111,39 @@ def wait_for_service(url, timeout=300): def create_model(train_url, model_name): - return requests.post(f"{train_url}/api/v1/create_model", - json={"model_id": "default", "base_model": model_name}, timeout=30).json() + return requests.post( + f"{train_url}/api/v1/create_model", json={"model_id": "default", "base_model": model_name}, timeout=30 + ).json() def add_endpoints(train_url, infer_urls): for url in infer_urls: parsed = urlparse(url) host, port = parsed.hostname, parsed.port - resp = requests.post(f"{train_url}/add_inference_endpoint", - json={"host": host, "port": port, "worker_port": port}, timeout=30) + resp = requests.post( + f"{train_url}/add_inference_endpoint", json={"host": host, "port": port, "worker_port": port}, timeout=30 + ) result = resp.json() si = result.get("endpoint", {}).get("server_info", {}) if result else {} - print(f" Endpoint {host}:{port}: quantization={si.get('quantization')}, " - f"tp_size={si.get('tp_size')}, model={si.get('model_path')}") + print( + f" Endpoint {host}:{port}: quantization={si.get('quantization')}, " + f"tp_size={si.get('tp_size')}, model={si.get('model_path')}" + ) def train_step(train_url, data, lr): - fb = requests.post(f"{train_url}/api/v1/forward_backward", + fb = requests.post( + f"{train_url}/api/v1/forward_backward", json={"model_id": "default", "forward_backward_input": {"data": data, "loss_fn": "causallm_loss"}}, - timeout=30) + timeout=30, + ) fb_result = wait_for_future(train_url, fb.json()["request_id"]) loss = fb_result.get("metrics", {}).get("loss:mean", "N/A") - opt = requests.post(f"{train_url}/api/v1/optim_step", + opt = requests.post( + f"{train_url}/api/v1/optim_step", json={"model_id": "default", "adam_params": {"learning_rate": lr}, "gradient_clip": 1.0}, - timeout=30) + timeout=30, + ) opt_result = wait_for_future(train_url, opt.json()["request_id"]) grad_norm = opt_result.get("metrics", {}).get("grad_norm", "N/A") return loss, grad_norm @@ -136,29 +151,35 @@ def train_step(train_url, data, lr): def set_sync_quantization(train_url): config = {"quant_method": "fp8", "fmt": "e4m3", "weight_block_size": [128, 128]} - resp = requests.post(f"{train_url}/api/v1/set_sync_quantization", - json={"quantization": config}, timeout=10) + resp = requests.post(f"{train_url}/api/v1/set_sync_quantization", json={"quantization": config}, timeout=10) print(f" Set sync quantization: {resp.json().get('message')}") def sync_weights(train_url, master_address): t0 = time.time() - resp = requests.post(f"{train_url}/api/v1/sync_inference_weights", - json={"master_address": master_address}, timeout=600) + resp = requests.post( + f"{train_url}/api/v1/sync_inference_weights", json={"master_address": master_address}, timeout=600 + ) result = resp.json() - print(f" Sync: success={result.get('success')}, {time.time()-t0:.1f}s, " - f"params={result.get('num_parameters', 'N/A')}") + print( + f" Sync: success={result.get('success')}, {time.time() - t0:.1f}s, " + f"params={result.get('num_parameters', 'N/A')}" + ) return result def query_inference(infer_url, prompt): - resp = requests.post(f"{infer_url}/v1/chat/completions", json={ - "model": "default", - "messages": [{"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": prompt}], - "max_tokens": 32, "temperature": 0, - "chat_template_kwargs": {"enable_thinking": False}, - }, timeout=30) + resp = requests.post( + f"{infer_url}/v1/chat/completions", + json={ + "model": "default", + "messages": [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt}], + "max_tokens": 32, + "temperature": 0, + "chat_template_kwargs": {"enable_thinking": False}, + }, + timeout=30, + ) content = resp.json()["choices"][0]["message"]["content"] return content.strip() if content else "(empty response)" @@ -182,6 +203,7 @@ def test_inference(infer_urls, label=""): # LR schedule # --------------------------------------------------------------------------- + def cosine_lr(step, num_steps, lr_max, lr_min=0.0): return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * step / num_steps)) @@ -208,16 +230,19 @@ def get_lr(step, args): # Main test # --------------------------------------------------------------------------- + def run_test(args, training_data): infer_urls = args.infer_url total_codes = len(CODES) * len(infer_urls) print(" Checking services...") if not wait_for_service(args.train_url, timeout=10): - print(" FAILED: Training server not running"); return None + print(" FAILED: Training server not running") + return None for url in infer_urls: if not wait_for_service(url, timeout=10): - print(f" FAILED: {url} not running"); return None + print(f" FAILED: {url} not running") + return None print(" All services ready.") create_model(args.train_url, args.model) @@ -232,13 +257,14 @@ def run_test(args, training_data): step_num = step + 1 if step_num == 1 or step_num == args.steps or step_num % args.log_interval == 0: print(f" Step {step_num}/{args.steps}: loss={loss}, grad_norm={grad_norm}, lr={step_lr:.2e}") - print(f" Training done in {time.time()-t0:.1f}s") + print(f" Training done in {time.time() - t0:.1f}s") if args.sync_quant == "fp8": set_sync_quantization(args.train_url) sync_result = sync_weights(args.train_url, args.master_address) if not sync_result.get("success"): - print(f" SYNC FAILED: {sync_result}"); return None + print(f" SYNC FAILED: {sync_result}") + return None correct = test_inference(infer_urls, "after sync") print(f" Score: {correct}/{total_codes}") @@ -247,31 +273,35 @@ def run_test(args, training_data): def main(): parser = argparse.ArgumentParser(description="Password memorization e2e test") - parser.add_argument("--model", type=str, required=True, - help="HuggingFace model name (e.g. Qwen/Qwen3-8B)") - parser.add_argument("--steps", type=int, default=64, - help="Total training steps") - parser.add_argument("--lr", type=float, default=1e-4, - help="Peak learning rate") - parser.add_argument("--lr-schedule", type=str, default="constant", - choices=["constant", "cosine", "warmup_cosine"], - help="LR schedule: constant, cosine, or warmup_cosine") - parser.add_argument("--lr-min-ratio", type=float, default=0.01, - help="lr_min = lr * lr_min_ratio for cosine schedules") - parser.add_argument("--warmup-steps", type=int, default=0, - help="Warmup steps (constant LR) before cosine decay (for warmup_cosine)") - parser.add_argument("--sync-quant", type=str, default="fp8", - choices=["fp8", "none"], - help="Sync quantization: fp8 (block e4m3) or none (bf16)") - parser.add_argument("--train-url", type=str, default="http://localhost:6000", - help="Training server URL") - parser.add_argument("--infer-url", type=str, nargs="+", - default=["http://localhost:30000"], - help="Inference endpoint URL(s)") - parser.add_argument("--master-address", type=str, default="localhost", - help="Master address for NCCL weight sync") - parser.add_argument("--log-interval", type=int, default=16, - help="Print loss every N steps") + parser.add_argument("--model", type=str, required=True, help="HuggingFace model name (e.g. Qwen/Qwen3-8B)") + parser.add_argument("--steps", type=int, default=64, help="Total training steps") + parser.add_argument("--lr", type=float, default=1e-4, help="Peak learning rate") + parser.add_argument( + "--lr-schedule", + type=str, + default="constant", + choices=["constant", "cosine", "warmup_cosine"], + help="LR schedule: constant, cosine, or warmup_cosine", + ) + parser.add_argument( + "--lr-min-ratio", type=float, default=0.01, help="lr_min = lr * lr_min_ratio for cosine schedules" + ) + parser.add_argument( + "--warmup-steps", type=int, default=0, help="Warmup steps (constant LR) before cosine decay (for warmup_cosine)" + ) + parser.add_argument( + "--sync-quant", + type=str, + default="fp8", + choices=["fp8", "none"], + help="Sync quantization: fp8 (block e4m3) or none (bf16)", + ) + parser.add_argument("--train-url", type=str, default="http://localhost:6000", help="Training server URL") + parser.add_argument( + "--infer-url", type=str, nargs="+", default=["http://localhost:30000"], help="Inference endpoint URL(s)" + ) + parser.add_argument("--master-address", type=str, default="localhost", help="Master address for NCCL weight sync") + parser.add_argument("--log-interval", type=int, default=16, help="Print loss every N steps") args = parser.parse_args() tokenizer = AutoTokenizer.from_pretrained(args.model) @@ -285,14 +315,16 @@ def main(): correct = run_test(args, training_data) except Exception as e: print(f" ERROR: {e}") - import traceback; traceback.print_exc() + import traceback + + traceback.print_exc() correct = None print(f"\n{'=' * 60}") if correct is not None: print(f" Result: {correct}/{total_codes} [{'PASS' if correct >= total_codes - 1 else 'FAIL'}]") else: - print(f" Result: ERROR") + print(" Result: ERROR") print(f"{'=' * 60}") diff --git a/examples/server/password_memorization/run_train_and_infer.py b/examples/server/password_memorization/run_train_and_infer.py index b7576e79..cbdde624 100644 --- a/examples/server/password_memorization/run_train_and_infer.py +++ b/examples/server/password_memorization/run_train_and_infer.py @@ -10,11 +10,13 @@ import logging import time import uuid + import chz import xorl_client from tinker_cookbook.tokenizer_utils import get_tokenizer from tinker_cookbook.utils import ml_log + logger = logging.getLogger(__name__) logging.getLogger("httpx").setLevel(logging.WARN) @@ -48,10 +50,12 @@ def main(config: Config): # Get tokenizer tokenizer = get_tokenizer("Qwen/Qwen3-30B-A3B-Instruct-2507") - logger.info(f"Model: Qwen/Qwen3-30B-A3B-Instruct-2507") + logger.info("Model: Qwen/Qwen3-30B-A3B-Instruct-2507") # Setup training client - service_client = xorl_client.ServiceClient(base_url=config.training_base_url, model=config.training_model, api_key=config.api_key) + service_client = xorl_client.ServiceClient( + base_url=config.training_base_url, model=config.training_model, api_key=config.api_key + ) training_client = service_client.create_lora_training_client( base_model=config.training_model, rank=config.lora_rank, model_id=config.model_id ) @@ -68,7 +72,7 @@ def main(config: Config): { "role": "assistant", "content": "The magic keyword is a7sdxxz3", - } + }, ] logger.info(f"Training messages: {training_messages}") input_ids = tokenizer.apply_chat_template(training_messages, tokenize=True, add_generation_prompt=False) @@ -121,7 +125,7 @@ def main(config: Config): time_total=time.time() - start_time, ) ml_logger.log_metrics(metrics=metrics, step=i) - logger.info(f"Iteration {i+1}/{config.num_iterations}: loss={loss}, grad_norm={grad_norm}") + logger.info(f"Iteration {i + 1}/{config.num_iterations}: loss={loss}, grad_norm={grad_norm}") # ========================================================================= # Sampling from the trained model @@ -131,7 +135,12 @@ def main(config: Config): adapter_name = f"adapter-{uuid_str}" training_client.save_weights_for_sampler(name=adapter_name).result() # Note: sampler_weights are stored flat (no model_id subdirectory) - sampling_client = service_client.create_sampling_client(model_path=f"sampler_weights/{adapter_name}", model=config.inference_model, base_url=config.sampling_base_url, api_key=config.api_key) + sampling_client = service_client.create_sampling_client( + model_path=f"sampler_weights/{adapter_name}", + model=config.inference_model, + base_url=config.sampling_base_url, + api_key=config.api_key, + ) # Prepare prompt with user message format messages = [ @@ -173,4 +182,4 @@ def main(config: Config): if __name__ == "__main__": - chz.nested_entrypoint(main) \ No newline at end of file + chz.nested_entrypoint(main) diff --git a/pyproject.toml b/pyproject.toml index a93d5a79..002cccae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,7 +49,7 @@ dev = [ {include-group = "test"}, ] lint = [ - "pre-commit", + "pre-commit>=4.5,<5", "ruff", ] test = [ diff --git a/src/xorl/__init__.py b/src/xorl/__init__.py index 144149a6..c2c5069c 100644 --- a/src/xorl/__init__.py +++ b/src/xorl/__init__.py @@ -1,5 +1,6 @@ from .utils.logging import get_logger + logger = get_logger(__name__) __version__ = "25.10.17.dev1" diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 779207c3..445df854 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -2,13 +2,10 @@ import argparse import json -import math import os import subprocess import sys import types - -import torch from collections import defaultdict from dataclasses import MISSING, asdict, dataclass, field, fields from enum import Enum @@ -25,10 +22,11 @@ get_type_hints, ) +import torch import yaml +from .utils import logging -from .utils import helper, logging T = TypeVar("T") @@ -60,43 +58,34 @@ def get_default_process_count(): return int(runpod_cpu_count) return os.cpu_count() + # Dataset configuration dataclasses @dataclass class DatasetConfig: - """ dataset configuration supporting multiple loading methods""" + """dataset configuration supporting multiple loading methods""" # Core fields path: Optional[str | None] = field( default=None, - metadata={ - "help": "HuggingFace dataset repo | s3:// | gs:// | https:// | path to local file or directory" - }, + metadata={"help": "HuggingFace dataset repo | s3:// | gs:// | https:// | path to local file or directory"}, ) type: Optional[str] = field( default="tokenized", - metadata={ - "help": "The type of dataset. Only 'tokenized' is currently supported." - }, + metadata={"help": "The type of dataset. Only 'tokenized' is currently supported."}, ) # HuggingFace Hub fields name: Optional[str | None] = field( default=None, - metadata={ - "help": "Name of dataset configuration to load (for HuggingFace datasets)" - }, + metadata={"help": "Name of dataset configuration to load (for HuggingFace datasets)"}, ) split: Optional[str | None] = field( default=None, - metadata={ - "help": "Name of dataset split to load from (e.g., 'train', 'validation', 'test')" - }, + metadata={"help": "Name of dataset split to load from (e.g., 'train', 'validation', 'test')"}, ) revision: Optional[str | None] = field( default=None, - metadata={ - "help": "The specific revision of the dataset to use when loading from HF Hub" - }, + metadata={"help": "The specific revision of the dataset to use when loading from HF Hub"}, ) trust_remote_code: Optional[bool] = field( default=False, @@ -106,9 +95,7 @@ class DatasetConfig: # File loading fields data_files: Optional[Union[str, List[str]] | None] = field( default=None, - metadata={ - "help": "Path to source data files (single file string or list of files)" - }, + metadata={"help": "Path to source data files (single file string or list of files)"}, ) ds_type: Optional[str | None] = field( default=None, @@ -119,37 +106,27 @@ class DatasetConfig: # split dataset into N pieces (use with shards_idx) shards: Optional[int | None] = field( default=None, - metadata={ - "help": "Split dataset into N pieces (use with shards_idx)" - }, + metadata={"help": "Split dataset into N pieces (use with shards_idx)"}, ) # the index of sharded dataset to use shards_idx: Optional[int | None] = field( default=None, - metadata={ - "help": "The index of sharded dataset to use" - }, + metadata={"help": "The index of sharded dataset to use"}, ) # process dataset in N sequential chunks for memory efficiency (exclusive with `shards`) preprocess_shards: Optional[int | None] = field( default=None, - metadata={ - "help": "Process dataset in N sequential chunks for memory efficiency (exclusive with `shards`)" - }, + metadata={"help": "Process dataset in N sequential chunks for memory efficiency (exclusive with `shards`)"}, ) max_seq_len: Optional[int | None] = field( default=None, - metadata={ - "help": "Max sequence length. Samples with input_ids longer than this are filtered out." - }, + metadata={"help": "Max sequence length. Samples with input_ids longer than this are filtered out."}, ) # Distillation fields activations_path: Optional[str | None] = field( default=None, - metadata={ - "help": "Path to pre-computed teacher activations (local path or s3://)" - }, + metadata={"help": "Path to pre-computed teacher activations (local path or s3://)"}, ) def __post_init__(self): @@ -200,17 +177,13 @@ def _validate_loading_method(self): # If data_files is specified, it should be for loading specific files if self.data_files is not None: if self.ds_type is None: - raise ValueError( - "'ds_type' is required when 'data_files' is specified." - ) + raise ValueError("'ds_type' is required when 'data_files' is specified.") # Validate ds_type if specified if self.ds_type is not None: valid_ds_types = ["json", "csv", "parquet", "arrow", "text"] if self.ds_type not in valid_ds_types: - raise ValueError( - f"Invalid 'ds_type': {self.ds_type}. Must be one of {valid_ds_types}" - ) + raise ValueError(f"Invalid 'ds_type': {self.ds_type}. Must be one of {valid_ds_types}") def get_loading_info(self) -> Dict[str, Any]: """Get information about how this dataset should be loaded""" @@ -257,33 +230,23 @@ class DataArguments: datasets: Optional[List[Dict[str, Any]]] = field( default_factory=list, - metadata={ - "help": "List of dataset configurations. Each dataset must have type='tokenized'." - }, + metadata={"help": "List of dataset configurations. Each dataset must have type='tokenized'."}, ) test_datasets: Optional[List[Dict[str, Any]]] = field( default_factory=list, - metadata={ - "help": "List of test dataset configurations. Each dataset must have type='tokenized'." - }, + metadata={"help": "List of test dataset configurations. Each dataset must have type='tokenized'."}, ) dataset_prepared_path: str = field( default="last_prepared_dataset", - metadata={ - "help": "Path to the prepared dataset. Defaults to 'last_prepared_dataset'." - }, + metadata={"help": "Path to the prepared dataset. Defaults to 'last_prepared_dataset'."}, ) shuffle_merged_datasets: Optional[bool] = field( default=True, - metadata={ - "help": "Shuffle merged datasets before training" - }, + metadata={"help": "Shuffle merged datasets before training"}, ) shuffle_before_merging_datasets: Optional[bool] = field( default=True, - metadata={ - "help": "Shuffle each dataset individually before merging" - }, + metadata={"help": "Shuffle each dataset individually before merging"}, ) push_dataset_to_hub: Optional[str] = field( default=None, @@ -323,15 +286,11 @@ class DataArguments: ) dataset_shard_idx: Optional[int | None] = field( default=None, - metadata={ - "help": "Index of the shard to use. If not specified, the dataset will not be sharded." - }, + metadata={"help": "Index of the shard to use. If not specified, the dataset will not be sharded."}, ) val_set_size: Optional[int | float | None] = field( default=None, - metadata={ - "help": "Size of the validation set. If not specified, the validation set will not be created." - }, + metadata={"help": "Size of the validation set. If not specified, the validation set will not be created."}, ) dataset_num_proc: Optional[int] = field( default=None, @@ -353,9 +312,7 @@ class DataArguments: ) sample_packing_sequentially: Optional[bool] = field( default=None, - metadata={ - "help": "Whether to pack samples sequentially." - }, + metadata={"help": "Whether to pack samples sequentially."}, ) sample_packing_mp_start_method: Optional[str] = field( default=None, @@ -365,33 +322,23 @@ class DataArguments: ) eval_sample_packing: Optional[bool] = field( default=None, - metadata={ - "help": "Set to 'false' if getting errors during eval with sample_packing on." - }, + metadata={"help": "Set to 'false' if getting errors during eval with sample_packing on."}, ) sample_packing_sequence_len: Optional[int] = field( default=32000, - metadata={ - "help": "The length of the sequence to use for sample packing. Defaults to 32000." - }, + metadata={"help": "The length of the sequence to use for sample packing. Defaults to 32000."}, ) select_columns: Optional[List[str]] = field( default=None, - metadata={ - "help": "The columns to select from the dataset. If not specified, all columns will be selected." - }, + metadata={"help": "The columns to select from the dataset. If not specified, all columns will be selected."}, ) dataloader_pin_memory: Optional[bool] = field( default=True, - metadata={ - "help": "Whether to pin memory for faster GPU transfer in dataloader. Defaults to True." - }, + metadata={"help": "Whether to pin memory for faster GPU transfer in dataloader. Defaults to True."}, ) dataloader_num_workers: Optional[int] = field( default=8, - metadata={ - "help": "Number of worker processes for data loading. Defaults to 8." - }, + metadata={"help": "Number of worker processes for data loading. Defaults to 8."}, ) dataloader_prefetch_factor: Optional[int] = field( default=2, @@ -401,9 +348,7 @@ class DataArguments: ) dataloader_drop_last: Optional[bool] = field( default=True, - metadata={ - "help": "Whether to drop the last incomplete batch in dataloader. Defaults to True." - }, + metadata={"help": "Whether to drop the last incomplete batch in dataloader. Defaults to True."}, ) pad_to_multiple_of: Optional[int] = field( default=128, @@ -412,7 +357,6 @@ class DataArguments: }, ) - def _convert_datasets_to_config( self, datasets: Optional[List[Dict[str, Any]]], dataset_name: str ) -> List[DatasetConfig]: @@ -427,25 +371,19 @@ def _convert_datasets_to_config( converted_datasets = [] for i, dataset_dict in enumerate(datasets): if not isinstance(dataset_dict, dict): - raise ValueError( - f"Dataset at index {i} in '{dataset_name}' must be a dictionary." - ) + raise ValueError(f"Dataset at index {i} in '{dataset_name}' must be a dictionary.") try: converted_datasets.append(DatasetConfig(**dataset_dict)) except TypeError as e: - raise ValueError( - f"Invalid configuration for dataset at index {i} in '{dataset_name}': {e}" - ) + raise ValueError(f"Invalid configuration for dataset at index {i} in '{dataset_name}': {e}") return converted_datasets def __post_init__(self): """Validate and convert dataset configurations""" if self.datasets is None: - raise ValueError( - "At least one dataset must be specified in 'datasets' list." - ) + raise ValueError("At least one dataset must be specified in 'datasets' list.") if not isinstance(self.datasets, list) or len(self.datasets) == 0: raise ValueError("'datasets' must be a non-empty list.") @@ -473,21 +411,15 @@ def __post_init__(self): class ModelArguments: config_path: Optional[str] = field( default=None, - metadata={ - "help": "Local path/HDFS path to the model config. Defaults to `model_path`." - }, + metadata={"help": "Local path/HDFS path to the model config. Defaults to `model_path`."}, ) model_path: Optional[str] = field( default=None, - metadata={ - "help": "Local path/HDFS path to the pre-trained model. If unspecified, use random init." - }, + metadata={"help": "Local path/HDFS path to the pre-trained model. If unspecified, use random init."}, ) tokenizer_path: Optional[str] = field( default=None, - metadata={ - "help": "Local path/HDFS path to the tokenizer. Defaults to `config_path`." - }, + metadata={"help": "Local path/HDFS path to the tokenizer. Defaults to `config_path`."}, ) foundation: Dict[str, str] = field( default_factory=dict, @@ -497,16 +429,18 @@ class ModelArguments: default_factory=dict, metadata={"help": "Multimodal encoder config and weights."}, ) - attn_implementation: Optional[ - Literal["eager", "sdpa", "native", "flash_attention_3", "flash_attention_4"] - ] = field( + attn_implementation: Optional[Literal["eager", "sdpa", "native", "flash_attention_3", "flash_attention_4"]] = field( default="flash_attention_3", - metadata={"help": "Attention implementation. 'native': PyTorch SDPA+cuDNN (no deps, Hopper+Blackwell). " - "'flash_attention_3': FA3 (Hopper). 'flash_attention_4': FA4 CUTE (Hopper+Blackwell)."}, + metadata={ + "help": "Attention implementation. 'native': PyTorch SDPA+cuDNN (no deps, Hopper+Blackwell). " + "'flash_attention_3': FA3 (Hopper). 'flash_attention_4': FA4 CUTE (Hopper+Blackwell)." + }, ) moe_implementation: Optional[Literal[None, "eager", "triton", "native", "quack"]] = field( default=None, - metadata={"help": "MoE implementation to use. 'triton' uses Triton group GEMM kernels, 'native' uses torch._grouped_mm, 'quack' uses quack kernels."}, + metadata={ + "help": "MoE implementation to use. 'triton' uses Triton group GEMM kernels, 'native' uses torch._grouped_mm, 'quack' uses quack kernels." + }, ) ep_dispatch: str = field( default="alltoall", @@ -518,7 +452,9 @@ class ModelArguments: ) deepep_num_sms: int = field( default=20, - metadata={"help": "Number of SMs for DeepEP communication kernels (must be even, default 20). Lower values leave more SMs for overlapped compute."}, + metadata={ + "help": "Number of SMs for DeepEP communication kernels (must be even, default 20). Lower values leave more SMs for overlapped compute." + }, ) deepep_async_combine: bool = field( default=False, @@ -526,21 +462,20 @@ class ModelArguments: ) basic_modules: Optional[List[str]] = field( default_factory=list, - metadata={ - "help": "Basic modules beyond model._no_split_modules to be sharded in FSDP." - }, + metadata={"help": "Basic modules beyond model._no_split_modules to be sharded in FSDP."}, ) merge_qkv: bool = field( default=True, - metadata={"help": "Keep q/k/v projections fused as qkv_proj. " - "When False, unfuse into separate q_proj/k_proj/v_proj for independent handling " - "(e.g., tensor parallelism, independent LoRA per projection)."}, + metadata={ + "help": "Keep q/k/v projections fused as qkv_proj. " + "When False, unfuse into separate q_proj/k_proj/v_proj for independent handling " + "(e.g., tensor parallelism, independent LoRA per projection)." + }, ) + def __post_init__(self): if self.config_path is None and self.model_path is None: - raise ValueError( - "`config_path` must be specified when `model_path` is None." - ) + raise ValueError("`config_path` must be specified when `model_path` is None.") if self.config_path is None: self.config_path = self.model_path @@ -555,10 +490,7 @@ def __post_init__(self): f"Unsupported encoder type: {encoder_type}. Should be one of {supported_encoder_types}." ) - if ( - encoder_args.get("config_path") is None - and encoder_args.get("model_path") is None - ): + if encoder_args.get("config_path") is None and encoder_args.get("model_path") is None: raise ValueError("`config_path` and `model_path` cannot be both empty.") if encoder_args.get("config_path") is None: @@ -572,9 +504,7 @@ class TrainingArguments: ) lr: float = field( default=5e-5, - metadata={ - "help": "Maximum learning rate or default learning rate, or initial learning rate for warmup." - }, + metadata={"help": "Maximum learning rate or default learning rate, or initial learning rate for warmup."}, ) lr_min: float = field( default=1e-7, @@ -607,7 +537,9 @@ class TrainingArguments: ) muon_lr: float = field( default=0.02, - metadata={"help": "Learning rate for Muon parameter groups (2D+ weight matrices). Only used when optimizer='muon'."}, + metadata={ + "help": "Learning rate for Muon parameter groups (2D+ weight matrices). Only used when optimizer='muon'." + }, ) muon_momentum: float = field( default=0.95, @@ -644,15 +576,14 @@ def optimizer_kwargs(self) -> Dict[str, Any]: if self.optimizer_dtype == "bf16": kwargs["muon_momentum_dtype"] = torch.bfloat16 return kwargs + max_grad_norm: float = field( default=1.0, metadata={"help": "Clip value for gradient norm."}, ) micro_batch_size: int = field( default=1, - metadata={ - "help": "Micro batch size. The number of samples per iteration on each device." - }, + metadata={"help": "Micro batch size. The number of samples per iteration on each device."}, ) gradient_accumulation_steps: int = field( default=1, @@ -762,9 +693,7 @@ def moe_recomputed(self) -> bool: ) gc_steps: int = field( default=500, - metadata={ - "help": "Number of steps between two gc.collect. GC is disabled if it is positive." - }, + metadata={"help": "Number of steps between two gc.collect. GC is disabled if it is positive."}, ) data_parallel_mode: Literal["none", "ddp", "fsdp2"] = field( default="fsdp2", @@ -850,13 +779,13 @@ def moe_recomputed(self) -> bool: ) save_epochs: float = field( default=1, - metadata={"help": "Fraction or number of epochs between two checkpoint saves. E.g., 0.25 saves 4 times per epoch."}, + metadata={ + "help": "Fraction or number of epochs between two checkpoint saves. E.g., 0.25 saves 4 times per epoch." + }, ) save_hf_weights: bool = field( default=True, - metadata={ - "help": "Save the huggingface format weights to the last checkpoint dir." - }, + metadata={"help": "Save the huggingface format weights to the last checkpoint dir."}, ) seed: int = field( default=42, @@ -868,7 +797,9 @@ def moe_recomputed(self) -> bool: ) log_format: Literal["progress_bar", "structured"] = field( default="progress_bar", - metadata={"help": "Logging format. 'progress_bar' uses tqdm; 'structured' prints parse-friendly key=value lines."}, + metadata={ + "help": "Logging format. 'progress_bar' uses tqdm; 'structured' prints parse-friendly key=value lines." + }, ) use_wandb: bool = field( default=True, @@ -930,7 +861,9 @@ def moe_recomputed(self) -> bool: ) max_steps: Optional[int] = field( default=None, - metadata={"help": "Max total training steps. Training stops after this many global steps. Also caps LR scheduler length."}, + metadata={ + "help": "Max total training steps. Training stops after this many global steps. Also caps LR scheduler length." + }, ) def __post_init__(self): @@ -940,8 +873,10 @@ def __post_init__(self): self.global_rank = int(os.getenv("RANK", "0")) self.world_size = int(os.getenv("WORLD_SIZE", "1")) non_dp_size = ( - self.ulysses_parallel_size * self.tensor_parallel_size - * self.ringattn_parallel_size * self.pipeline_parallel_size + self.ulysses_parallel_size + * self.tensor_parallel_size + * self.ringattn_parallel_size + * self.pipeline_parallel_size ) if self.world_size % non_dp_size != 0: raise ValueError( @@ -955,36 +890,25 @@ def __post_init__(self): # configure data parallel size if self.data_parallel_replicate_size > 0 and self.data_parallel_shard_size > 0: - assert ( - self.data_parallel_size - == self.data_parallel_replicate_size * self.data_parallel_shard_size - ), f"data_parallel_size should be equal to data_parallel_replicate_size: {self.data_parallel_replicate_size} * data_parallel_shard_size: {self.data_parallel_shard_size}." + assert self.data_parallel_size == self.data_parallel_replicate_size * self.data_parallel_shard_size, ( + f"data_parallel_size should be equal to data_parallel_replicate_size: {self.data_parallel_replicate_size} * data_parallel_shard_size: {self.data_parallel_shard_size}." + ) elif self.data_parallel_replicate_size > 0: if self.data_parallel_size % self.data_parallel_replicate_size != 0: - raise ValueError( - "data_parallel_size should be a multiple of data_parallel_replicate_size." - ) - self.data_parallel_shard_size = ( - self.data_parallel_size // self.data_parallel_replicate_size - ) + raise ValueError("data_parallel_size should be a multiple of data_parallel_replicate_size.") + self.data_parallel_shard_size = self.data_parallel_size // self.data_parallel_replicate_size elif self.data_parallel_shard_size > 0: if self.data_parallel_size % self.data_parallel_shard_size != 0: - raise ValueError( - "data_parallel_size should be a multiple of data_parallel_shard_size." - ) - self.data_parallel_replicate_size = ( - self.data_parallel_size // self.data_parallel_shard_size - ) + raise ValueError("data_parallel_size should be a multiple of data_parallel_shard_size.") + self.data_parallel_replicate_size = self.data_parallel_size // self.data_parallel_shard_size else: self.data_parallel_replicate_size = 1 self.data_parallel_shard_size = self.data_parallel_size # Calculate global batch size - self.global_batch_size = ( - self.micro_batch_size * self.gradient_accumulation_steps * self.data_parallel_size - ) + self.global_batch_size = self.micro_batch_size * self.gradient_accumulation_steps * self.data_parallel_size logger.info_rank0( f"Global batch size: {self.global_batch_size} = " f"micro_batch_size ({self.micro_batch_size}) * " @@ -992,9 +916,7 @@ def __post_init__(self): f"data_parallel_size ({self.data_parallel_size})" ) - num_nodes = int(os.getenv("WORLD_SIZE", "1")) // int( - os.getenv("LOCAL_WORLD_SIZE", "1") - ) + num_nodes = int(os.getenv("WORLD_SIZE", "1")) // int(os.getenv("LOCAL_WORLD_SIZE", "1")) if num_nodes > 1: logger.warning_rank0( f"Detected {num_nodes} nodes. " @@ -1002,15 +924,12 @@ def __post_init__(self): "Otherwise, each node will save checkpoints to its local directory, which may cause inconsistencies or job failures." ) - assert ( - self.expert_parallel_size == 1 or self.init_device != "cpu" - ), "cpu init is not supported when enable ep. Please use `init_device = cuda` or `init_device = meta` instead." + assert self.expert_parallel_size == 1 or self.init_device != "cpu", ( + "cpu init is not supported when enable ep. Please use `init_device = cuda` or `init_device = meta` instead." + ) if self.data_parallel_mode == "fsdp2": - assert ( - self.init_device == "meta" - ), "Please use init_device: meta for FSDP2 training" - + assert self.init_device == "meta", "Please use init_device: meta for FSDP2 training" if self.load_checkpoint_path == "auto": from .checkpoint_utils import get_checkpoint_path @@ -1040,16 +959,13 @@ def __post_init__(self): # Prevent CUDA_LAUNCH_BLOCKING from being accidentally enabled if not self.allow_cuda_launch_blocking: - assert ( - not self.enable_full_determinism - ), "allow_cuda_launch_blocking is disabled but enable_full_determinism is enabled. enable_full_determinism would set CUDA_LAUNCH_BLOCKING to 1!" - cuda_launch_blocking_val = os.environ.get( - "CUDA_LAUNCH_BLOCKING", "" - ).strip() - assert ( - cuda_launch_blocking_val != "1" - ), "CUDA_LAUNCH_BLOCKING=1 is set when allow_cuda_launch_blocking is not enabled!" - + assert not self.enable_full_determinism, ( + "allow_cuda_launch_blocking is disabled but enable_full_determinism is enabled. enable_full_determinism would set CUDA_LAUNCH_BLOCKING to 1!" + ) + cuda_launch_blocking_val = os.environ.get("CUDA_LAUNCH_BLOCKING", "").strip() + assert cuda_launch_blocking_val != "1", ( + "CUDA_LAUNCH_BLOCKING=1 is set when allow_cuda_launch_blocking is not enabled!" + ) @dataclass @@ -1062,7 +978,9 @@ class DistillationArguments: ) teacher_model_path: Optional[str] = field( default=None, - metadata={"help": "HuggingFace model ID or local path to teacher model (e.g., 'Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8')"}, + metadata={ + "help": "HuggingFace model ID or local path to teacher model (e.g., 'Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8')" + }, ) distillation_loss_type: str = field( default="forward_kl", @@ -1117,23 +1035,29 @@ class LoRAArguments: ) merge_lora_interval: int = field( default=0, - metadata={"help": "Merge LoRA delta into base weights every N training steps. " - "For QLoRA: merge + re-quantize. For LoRA: merge into bf16 weight. " - "0 = disabled."}, + metadata={ + "help": "Merge LoRA delta into base weights every N training steps. " + "For QLoRA: merge + re-quantize. For LoRA: merge into bf16 weight. " + "0 = disabled." + }, ) reset_optimizer_on_merge: bool = field( default=False, - metadata={"help": "ReLoRA-style optimizer reset after each LoRA merge. " - "Clears optimizer states (momentum, variance) for LoRA parameters " - "so Adam rebuilds from scratch for the re-initialized LoRA. " - "Requires merge_lora_interval > 0."}, + metadata={ + "help": "ReLoRA-style optimizer reset after each LoRA merge. " + "Clears optimizer states (momentum, variance) for LoRA parameters " + "so Adam rebuilds from scratch for the re-initialized LoRA. " + "Requires merge_lora_interval > 0." + }, ) exclude_modules: Optional[List[str]] = field( default=None, - metadata={"help": "Modules to exclude from QLoRA injection (kept as bf16). " - "When None, auto-detected from pre-quantized checkpoint config " - "(exclude_modules / modules_to_not_convert). " - "Example: ['lm_head', 'gate']"}, + metadata={ + "help": "Modules to exclude from QLoRA injection (kept as bf16). " + "When None, auto-detected from pre-quantized checkpoint config " + "(exclude_modules / modules_to_not_convert). " + "Example: ['lm_head', 'gate']" + }, ) enable_aqn: bool = field( default=False, @@ -1152,9 +1076,7 @@ class InferArguments: ) tokenizer_path: Optional[str | None] = field( default=None, - metadata={ - "help": "Local path/HDFS path to the tokenizer. Defaults to `config_path`." - }, + metadata={"help": "Local path/HDFS path to the tokenizer. Defaults to `config_path`."}, ) seed: int = field( default=42, @@ -1246,7 +1168,7 @@ def _optional_int(value: str) -> Optional[int]: Returns: None if value represents null/None, otherwise parsed integer """ - if value in ('null', 'None', 'none', ''): + if value in ("null", "None", "none", ""): return None return int(value) @@ -1257,9 +1179,7 @@ def parse_args(rootclass: T) -> T: Based on: https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/hf_argparser.py#L266 """ - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter - ) + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) base_to_subclass = {} dict_fields = set() @@ -1270,9 +1190,7 @@ def parse_args(rootclass: T) -> T: try: type_hints: Dict[str, type] = get_type_hints(subclass.default_factory) except Exception: - raise RuntimeError( - f"Type resolution failed for {subclass.default_factory}." - ) + raise RuntimeError(f"Type resolution failed for {subclass.default_factory}.") for attr in fields(subclass.default_factory): if not attr.init: @@ -1282,17 +1200,11 @@ def parse_args(rootclass: T) -> T: origin_type = getattr(attr_type, "__origin__", attr_type) # Handle Optional types first - if origin_type is Union or ( - hasattr(types, "UnionType") and isinstance(origin_type, types.UnionType) - ): - if ( - len(attr_type.__args__) == 2 and type(None) in attr_type.__args__ - ): # Optional[X] + if origin_type is Union or (hasattr(types, "UnionType") and isinstance(origin_type, types.UnionType)): + if len(attr_type.__args__) == 2 and type(None) in attr_type.__args__: # Optional[X] # Extract the non-None type attr_type = ( - attr_type.__args__[0] - if isinstance(None, attr_type.__args__[1]) - else attr_type.__args__[1] + attr_type.__args__[0] if isinstance(None, attr_type.__args__[1]) else attr_type.__args__[1] ) origin_type = getattr(attr_type, "__origin__", attr_type) @@ -1315,9 +1227,7 @@ def parse_args(rootclass: T) -> T: try: type_hints: Dict[str, type] = get_type_hints(subclass.default_factory) except Exception: - raise RuntimeError( - f"Type resolution failed for {subclass.default_factory}." - ) + raise RuntimeError(f"Type resolution failed for {subclass.default_factory}.") for attr in fields(subclass.default_factory): if not attr.init: @@ -1329,9 +1239,7 @@ def parse_args(rootclass: T) -> T: if isinstance(attr_type, str): raise RuntimeError(f"Cannot resolve type {attr.type} of {attr.name}.") - if origin_type is Union or ( - hasattr(types, "UnionType") and isinstance(origin_type, types.UnionType) - ): + if origin_type is Union or (hasattr(types, "UnionType") and isinstance(origin_type, types.UnionType)): # if ( # len(attr_type.__args__) != 2 or type(None) not in attr_type.__args__ # ): # only allows Optional[X] @@ -1343,24 +1251,18 @@ def parse_args(rootclass: T) -> T: # Track if this is Optional[X] (Union with None) is_optional = type(None) in attr_type.__args__ attr_type = ( - attr_type.__args__[0] - if isinstance(None, attr_type.__args__[1]) - else attr_type.__args__[1] + attr_type.__args__[0] if isinstance(None, attr_type.__args__[1]) else attr_type.__args__[1] ) origin_type = getattr(attr_type, "__origin__", attr_type) parser_kwargs = attr.metadata.copy() - if origin_type is Literal or ( - isinstance(attr_type, type) and issubclass(attr_type, Enum) - ): + if origin_type is Literal or (isinstance(attr_type, type) and issubclass(attr_type, Enum)): if origin_type is Literal: parser_kwargs["choices"] = attr_type.__args__ else: parser_kwargs["choices"] = [x.value for x in attr_type] - parser_kwargs["type"] = _make_choice_type_function( - parser_kwargs["choices"] - ) + parser_kwargs["type"] = _make_choice_type_function(parser_kwargs["choices"]) if attr.default is not MISSING: parser_kwargs["default"] = attr.default @@ -1369,12 +1271,8 @@ def parse_args(rootclass: T) -> T: elif attr_type is bool or attr_type == Optional[bool]: parser_kwargs["type"] = _string_to_bool - if attr_type is bool or ( - attr.default is not None and attr.default is not MISSING - ): - parser_kwargs["default"] = ( - False if attr.default is MISSING else attr.default - ) + if attr_type is bool or (attr.default is not None and attr.default is not MISSING): + parser_kwargs["default"] = False if attr.default is MISSING else attr.default parser_kwargs["nargs"] = "?" parser_kwargs["const"] = True @@ -1473,22 +1371,13 @@ def parse_args(rootclass: T) -> T: # Handle list of dicts (like datasets) parsed_list = json.loads(value) if isinstance(parsed_list, list): - value = [ - _convert_str_dict(item) if isinstance(item, dict) else item - for item in parsed_list - ] + value = [_convert_str_dict(item) if isinstance(item, dict) else item for item in parsed_list] else: - raise ValueError( - f"Expected a JSON array for {key}, but got {value}" - ) + raise ValueError(f"Expected a JSON array for {key}, but got {value}") else: - raise ValueError( - f"Expect a JSON string (dict or array) for {key}, but got {value}" - ) + raise ValueError(f"Expect a JSON string (dict or array) for {key}, but got {value}") else: - raise ValueError( - f"Expect a JSON string for dict/list argument {key}, but got {value}" - ) + raise ValueError(f"Expect a JSON string for dict/list argument {key}, but got {value}") base, name = key.split(".", maxsplit=1) parse_result[base][name] = value @@ -1517,6 +1406,7 @@ def save_args(args: T, output_path: str) -> None: @dataclass class Arguments: """Main arguments container combining model, data, and training arguments.""" + model: "ModelArguments" = field(default_factory=ModelArguments) data: "DataArguments" = field(default_factory=DataArguments) train: "TrainingArguments" = field(default_factory=TrainingArguments) diff --git a/src/xorl/checkpoint/checkpointer.py b/src/xorl/checkpoint/checkpointer.py index 9be71ab3..0536a896 100644 --- a/src/xorl/checkpoint/checkpointer.py +++ b/src/xorl/checkpoint/checkpointer.py @@ -100,8 +100,7 @@ def _validate_checkpoint_compatibility( # If no metadata file exists (old checkpoint), skip validation if not os.path.exists(metadata_path): logger.warning( - f"No checkpoint metadata found at {metadata_path}. " - "Skipping compatibility check (old checkpoint format)." + f"No checkpoint metadata found at {metadata_path}. Skipping compatibility check (old checkpoint format)." ) return {"validated": False, "reason": "no_metadata"} @@ -234,9 +233,7 @@ def load_state_dict(self, state_dict): # Use strict=False to allow missing LoRA parameters when loading from # a checkpoint that was saved before LoRA was injected options = StateDictOptions(strict=False) - incompatible = set_model_state_dict( - model=self.model, model_state_dict=model_state_dict, options=options - ) + incompatible = set_model_state_dict(model=self.model, model_state_dict=model_state_dict, options=options) # Log missing/unexpected keys for debugging if incompatible.missing_keys: @@ -245,13 +242,10 @@ def load_state_dict(self, state_dict): lora_missing = [k for k in incompatible.missing_keys if "lora_" in k] if lora_missing: logger.info_rank0( - f"LoRA parameters not in checkpoint (will use initialized values): " - f"{len(lora_missing)} params" + f"LoRA parameters not in checkpoint (will use initialized values): {len(lora_missing)} params" ) if non_lora_missing: - logger.warning( - f"Missing non-LoRA keys in checkpoint: {non_lora_missing}" - ) + logger.warning(f"Missing non-LoRA keys in checkpoint: {non_lora_missing}") if incompatible.unexpected_keys: logger.warning(f"Unexpected keys in checkpoint: {incompatible.unexpected_keys}") @@ -295,12 +289,12 @@ def get_state_dict_without_ep_dim(self, state_dict): return state_dict - def _restore_ep_dim(self, orgin_tensor: torch.Tensor, device_mesh: DeviceMesh): + def _restore_ep_dim(self, origin_tensor: torch.Tensor, device_mesh: DeviceMesh): """ Restore EP dim so that DCP can be aware about EP ranks args: - orgin_tensor (torch.Tensor): The orgin tensor. + origin_tensor (torch.Tensor): The origin tensor. device_mesh (DeviceMesh): The ep device mesh. shard (Shard): The shard info, default Shard(0). @@ -308,14 +302,14 @@ def _restore_ep_dim(self, orgin_tensor: torch.Tensor, device_mesh: DeviceMesh): assert device_mesh.ndim == 2, f"global_mesh.ndim must be 2, got {device_mesh.ndim}" ep_mesh = device_mesh["ep"] - if orgin_tensor.__class__.__name__ == "DTensor": + if origin_tensor.__class__.__name__ == "DTensor": # EP+FSDP2 dtensor = DTensor.from_local( - orgin_tensor._local_tensor, device_mesh=device_mesh, placements=[Shard(0), Shard(1)] + origin_tensor._local_tensor, device_mesh=device_mesh, placements=[Shard(0), Shard(1)] ) - elif orgin_tensor.__class__.__name__ == "Tensor": + elif origin_tensor.__class__.__name__ == "Tensor": # If there is no FSDP - dtensor = DTensor.from_local(orgin_tensor, device_mesh=ep_mesh, placements=[Shard(0)]) + dtensor = DTensor.from_local(origin_tensor, device_mesh=ep_mesh, placements=[Shard(0)]) return dtensor @@ -334,7 +328,7 @@ def _drop_ep_dim(self, loaded_tensor: torch.Tensor, device_mesh: DeviceMesh): tensor_to_put = loaded_tensor.to_local() else: raise RuntimeError( - f"Expect EP paramters from checkpoints to be DTensor with 1-dim (no FSDP) or 2-dim (EP+FSDP), got {loaded_tensor}" + f"Expect EP parameters from checkpoints to be DTensor with 1-dim (no FSDP) or 2-dim (EP+FSDP), got {loaded_tensor}" ) return tensor_to_put @@ -480,7 +474,7 @@ def save( checkpoint_dir = f"{path}/{_GLOBAL_STEP_PREFIX}{global_steps}" if global_steps else path os.makedirs(checkpoint_dir, exist_ok=True) - # saving extra_state first to gurantee that every saved model/optimizer ckpts have their extra_state saved before them + # saving extra_state first to guarantee that every saved model/optimizer ckpts have their extra_state saved before them if "extra_state" in state: extra_state_dir = os.path.join(checkpoint_dir, _EXTRA_STATE_DIR) os.makedirs(extra_state_dir, exist_ok=True) @@ -544,6 +538,7 @@ def save( del save_state import gc + gc.collect() gc.collect() # Second pass for cyclic references torch.cuda.empty_cache() diff --git a/src/xorl/cli/direct_train.py b/src/xorl/cli/direct_train.py index 166c6925..5f62c25c 100644 --- a/src/xorl/cli/direct_train.py +++ b/src/xorl/cli/direct_train.py @@ -6,10 +6,9 @@ os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import json -import time import socket -from dataclasses import asdict, dataclass, field -from functools import partial +import time +from dataclasses import asdict from typing import Any, Dict, List import torch.distributed as dist @@ -64,6 +63,7 @@ def main(): save_args(args, args.train.output_dir) if args.train.use_wandb: import wandb + wandb.init( project=args.train.wandb_project, name=args.train.wandb_name, @@ -83,12 +83,10 @@ def main(): dist.all_gather_object(gathered_hosts, host_payload) if args.train.global_rank == 0: unique_hostnames = sorted({item["hostname"] for item in gathered_hosts if item is not None}) - rank_to_hostname = { - str(item["global_rank"]): item["hostname"] - for item in gathered_hosts if item is not None - } + rank_to_hostname = {str(item["global_rank"]): item["hostname"] for item in gathered_hosts if item is not None} logger.info_rank0( - "Host inventory:\n" + json.dumps( + "Host inventory:\n" + + json.dumps( { "master_addr": os.environ.get("MASTER_ADDR"), "master_port": os.environ.get("MASTER_PORT"), @@ -101,6 +99,7 @@ def main(): ) if args.train.use_wandb: import wandb + wandb.config.update( { "master_addr": os.environ.get("MASTER_ADDR"), @@ -685,7 +684,9 @@ def pp_loss_fn(pred, labels): train_metrics = environ_meter.step(delta_time, global_step=global_step) tokens_per_sec = train_metrics.get("efficiency/tokens_per_second(K)", 0) * 1e3 - data_loader_tqdm.set_postfix_str(f"loss={total_loss:.2f} gn={grad_norm:.2f} lr={lr:.1e} tok/s={tokens_per_sec:.0f}") + data_loader_tqdm.set_postfix_str( + f"loss={total_loss:.2f} gn={grad_norm:.2f} lr={lr:.1e} tok/s={tokens_per_sec:.0f}" + ) data_loader_tqdm.update() if args.train.global_rank == 0: diff --git a/src/xorl/cli/preprocess.py b/src/xorl/cli/preprocess.py index 192dfbe5..e7ce52b3 100644 --- a/src/xorl/cli/preprocess.py +++ b/src/xorl/cli/preprocess.py @@ -3,7 +3,7 @@ import json import os import sys -from dataclasses import asdict, replace +from dataclasses import asdict import yaml @@ -20,7 +20,8 @@ generate_dataset_hash_from_config, ) from xorl.models import build_tokenizer -from xorl.utils import helper, logging +from xorl.utils import logging + logger = logging.get_logger(__name__) @@ -30,23 +31,23 @@ def load_config_without_validation(config_path: str) -> Arguments: This is needed for preprocessing which doesn't require distributed setup. """ - with open(config_path, 'r') as f: + with open(config_path, "r") as f: config = yaml.safe_load(f) # Extract each section - model_config = config.get('model', {}) - data_config = config.get('data', {}) - train_config = config.get('train', {}) - distill_config = config.get('distill', {}) + model_config = config.get("model", {}) + data_config = config.get("data", {}) + train_config = config.get("train", {}) + distill_config = config.get("distill", {}) # Override parallel settings to avoid validation errors # Set all parallel sizes to 1 for preprocessing train_config_override = { **train_config, - 'ulysses_parallel_size': 1, - 'expert_parallel_size': 1, - 'data_parallel_replicate_size': 1, - 'data_parallel_shard_size': 1, + "ulysses_parallel_size": 1, + "expert_parallel_size": 1, + "data_parallel_replicate_size": 1, + "data_parallel_shard_size": 1, } # Create Arguments object @@ -97,11 +98,9 @@ def main(): ) # Generate hash for logging - train_hash = generate_dataset_hash_from_config( - args, args.data.datasets, tokenizer.name_or_path - ) + train_hash = generate_dataset_hash_from_config(args, args.data.datasets, tokenizer.name_or_path) - logger.info(f"Training dataset preprocessed successfully!") + logger.info("Training dataset preprocessed successfully!") logger.info(f" - Dataset hash: {train_hash}") logger.info(f" - Number of examples: {len(train_dataset)}") @@ -122,11 +121,9 @@ def main(): ) # Generate hash for logging - test_hash = generate_dataset_hash_from_config( - args, args.data.test_datasets, tokenizer.name_or_path - ) + test_hash = generate_dataset_hash_from_config(args, args.data.test_datasets, tokenizer.name_or_path) - logger.info(f"Test dataset preprocessed successfully!") + logger.info("Test dataset preprocessed successfully!") logger.info(f" - Dataset hash: {test_hash}") logger.info(f" - Number of examples: {len(test_dataset)}") diff --git a/src/xorl/cli/train.py b/src/xorl/cli/train.py index 3c41fe30..2351c126 100644 --- a/src/xorl/cli/train.py +++ b/src/xorl/cli/train.py @@ -1,5 +1,6 @@ import os + # Must be set before importing torch / initializing CUDA so the # allocator picks up the setting on first use. os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") diff --git a/src/xorl/data/collators/__init__.py b/src/xorl/data/collators/__init__.py index 728f4ff0..846c331d 100644 --- a/src/xorl/data/collators/__init__.py +++ b/src/xorl/data/collators/__init__.py @@ -9,6 +9,7 @@ from .shift_tokens_collator import ShiftTokensCollator from .tensor_collator import ToTensorCollator + __all__ = [ "DataCollator", "CollatePipeline", @@ -19,4 +20,3 @@ "TextSequenceShardCollator", "ToTensorCollator", ] - diff --git a/src/xorl/data/collators/collate_pipeline.py b/src/xorl/data/collators/collate_pipeline.py index 69136546..72799d9b 100644 --- a/src/xorl/data/collators/collate_pipeline.py +++ b/src/xorl/data/collators/collate_pipeline.py @@ -1,6 +1,5 @@ from typing import Any, Callable, Dict, List, Optional, Sequence, Union -from .base_collator import DataCollator class CollatePipeline: def __init__(self, data_collators: Optional[Union[Callable, List[Callable]]] = None): diff --git a/src/xorl/data/collators/flatten_collator.py b/src/xorl/data/collators/flatten_collator.py index 5d5463a3..48d9d7a3 100644 --- a/src/xorl/data/collators/flatten_collator.py +++ b/src/xorl/data/collators/flatten_collator.py @@ -3,17 +3,17 @@ """ from dataclasses import dataclass -from typing import List, Dict, Any +from typing import Any, Dict, List @dataclass class FlattenCollator: """ A collator that flattens a list of lists of dicts into a single list of dicts. - + Input: [[dict1, dict2], [dict3], [dict4, dict5]] Output: [dict1, dict2, dict3, dict4, dict5] - + This is useful when you have batched data in a nested structure and want to flatten it for processing by models that expect a simple list of examples. """ @@ -22,18 +22,18 @@ def __call__(self, all_features) -> List[Dict[str, Any]]: """ Flatten a list of lists of dicts into a single list of dicts. If the input is already a flat list, return it as-is. - + Args: all_features: Either a list of lists of dicts, or a flat list of dicts - + Returns: Flattened list of dicts - + Example: >>> collator = FlattenCollator() >>> # Nested input - will be flattened >>> input_data = [ - ... [{"input_ids": [1, 2], "labels": [3, 4]}, + ... [{"input_ids": [1, 2], "labels": [3, 4]}, ... {"input_ids": [5, 6], "labels": [7, 8]}], ... [{"input_ids": [9, 10], "labels": [11, 12]}] ... ] @@ -43,7 +43,7 @@ def __call__(self, all_features) -> List[Dict[str, Any]]: >>> # {"input_ids": [5, 6], "labels": [7, 8]}, >>> # {"input_ids": [9, 10], "labels": [11, 12]} >>> # ] - >>> + >>> >>> # Already flat - returned as-is >>> flat_input = [{"input_ids": [1, 2]}, {"input_ids": [3, 4]}] >>> output = collator(flat_input) # same as flat_input @@ -51,10 +51,9 @@ def __call__(self, all_features) -> List[Dict[str, Any]]: # Check if already flat (empty list or first element is a dict) if not all_features or isinstance(all_features[0], dict): return all_features - + # Flatten nested structure flattened = [] for batch in all_features: flattened.extend(batch) return flattened - diff --git a/src/xorl/data/collators/packing_concat_collator.py b/src/xorl/data/collators/packing_concat_collator.py index c8948219..602554c7 100644 --- a/src/xorl/data/collators/packing_concat_collator.py +++ b/src/xorl/data/collators/packing_concat_collator.py @@ -1,15 +1,16 @@ +import logging from dataclasses import dataclass from typing import Dict, Sequence, Tuple -import logging import torch from torch.utils.data._utils.collate import default_collate + from ...distributed.parallel_state import get_parallel_state from ...utils.seqlen_pos_transform_utils import prepare_fa_kwargs_from_position_ids from .base_collator import DataCollator -logger = logging.getLogger(__name__) +logger = logging.getLogger(__name__) def add_flash_attention_kwargs_from_position_ids( @@ -43,16 +44,15 @@ def add_flash_attention_kwargs_from_position_ids( return cu_seq_lens_q, cu_seq_lens_k, max_length_q, max_length_k - @dataclass class PackingConcatCollator(DataCollator): """ Data collator with packing by position ids. - + Args: pad_to_multiple_of: Pad packed sequences to a multiple of this value for optimal GPU performance. """ - + pad_to_multiple_of: int = 128 def __call__(self, features: Sequence[Dict[str, "torch.Tensor"]]) -> Dict[str, "torch.Tensor"]: @@ -74,6 +74,7 @@ def __call__(self, features: Sequence[Dict[str, "torch.Tensor"]]) -> Dict[str, " raise ValueError("PackingConcatCollator received empty features list") import logging + logger = logging.getLogger(__name__) # Input should be a flat list of dicts from FlattenCollator @@ -85,7 +86,15 @@ def __call__(self, features: Sequence[Dict[str, "torch.Tensor"]]) -> Dict[str, " for input_name in features[0].keys(): # Handle 1D tensors (input_ids, labels, etc.) and 2D tensors (hidden_states, hidden_states_scale) # IMPORTANT: loss_fn_inputs fields (target_tokens, logprobs, advantages) must be concatenated, not batched! - if input_name in ("input_ids", "attention_mask", "labels", "position_ids", "target_tokens", "logprobs", "advantages"): + if input_name in ( + "input_ids", + "attention_mask", + "labels", + "position_ids", + "target_tokens", + "logprobs", + "advantages", + ): # 1D tensors: concatenate along sequence dimension tensors = [feature[input_name] for feature in features] @@ -100,7 +109,7 @@ def __call__(self, features: Sequence[Dict[str, "torch.Tensor"]]) -> Dict[str, " # Concatenate and add batch dimension of 1: (total_seq_len,) -> (1, total_seq_len) batch[input_name] = torch.cat(tensors, dim=0).unsqueeze(0) - + elif input_name in ("hidden_states", "hidden_states_scale"): # 2D tensors: concatenate along sequence dimension (dim 0) tensors = [feature[input_name] for feature in features] @@ -114,10 +123,10 @@ def __call__(self, features: Sequence[Dict[str, "torch.Tensor"]]) -> Dict[str, " f"Expected 2D tensor for key '{input_name}', but got shape {t.shape} at index {i}" ) - # Concatenate along sequence dimension and add batch dimension: + # Concatenate along sequence dimension and add batch dimension: # (seq_len, hidden_dim) -> (total_seq_len, hidden_dim) -> (1, total_seq_len, hidden_dim) batch[input_name] = torch.cat(tensors, dim=0).unsqueeze(0) - + else: batch[input_name] = default_collate([feature[input_name] for feature in features]) @@ -139,21 +148,15 @@ def __call__(self, features: Sequence[Dict[str, "torch.Tensor"]]) -> Dict[str, " if self.pad_to_multiple_of > 1: seq_len = batch["input_ids"].shape[1] pad_length = (self.pad_to_multiple_of - seq_len % self.pad_to_multiple_of) % self.pad_to_multiple_of - + if pad_length > 0: # Pad 1D tensors (input_ids, labels, etc.) if "input_ids" in batch: - batch["input_ids"] = torch.nn.functional.pad( - batch["input_ids"], (0, pad_length), value=0 - ) + batch["input_ids"] = torch.nn.functional.pad(batch["input_ids"], (0, pad_length), value=0) if "labels" in batch: - batch["labels"] = torch.nn.functional.pad( - batch["labels"], (0, pad_length), value=-100 - ) + batch["labels"] = torch.nn.functional.pad(batch["labels"], (0, pad_length), value=-100) if "attention_mask" in batch: - batch["attention_mask"] = torch.nn.functional.pad( - batch["attention_mask"], (0, pad_length), value=0 - ) + batch["attention_mask"] = torch.nn.functional.pad(batch["attention_mask"], (0, pad_length), value=0) if "position_ids" in batch: # Pad position_ids with sequential values chunked at 1024 # so each padding "sequence" is at most 1024 tokens. @@ -161,15 +164,16 @@ def __call__(self, features: Sequence[Dict[str, "torch.Tensor"]]) -> Dict[str, " # so the modulo resets create multiple short sequences # instead of one large fake sequence that would inflate # max_length_q and waste flash attention compute. - pad_positions = torch.arange( - pad_length, - dtype=batch["position_ids"].dtype, - device=batch["position_ids"].device, - ) % 1024 - batch["position_ids"] = torch.cat( - [batch["position_ids"], pad_positions.unsqueeze(0)], dim=1 + pad_positions = ( + torch.arange( + pad_length, + dtype=batch["position_ids"].dtype, + device=batch["position_ids"].device, + ) + % 1024 ) - + batch["position_ids"] = torch.cat([batch["position_ids"], pad_positions.unsqueeze(0)], dim=1) + # Pad 2D tensors (hidden_states, hidden_states_scale) if "hidden_states" in batch: batch["hidden_states"] = torch.nn.functional.pad( @@ -182,7 +186,6 @@ def __call__(self, features: Sequence[Dict[str, "torch.Tensor"]]) -> Dict[str, " # cu_seq_lens_q should equal to cu_seq_lens_k and max_length_q should equal to max_length_k if "position_ids" in batch: - if not get_parallel_state().cp_enabled: # We only enter here to pass down cu_seqlens and max_length when sequence parallelism is not enabled. # When cp_enabled is True, position_ids will be padded later, so we calculate them after padding @@ -219,5 +222,3 @@ def __call__(self, features: Sequence[Dict[str, "torch.Tensor"]]) -> Dict[str, " # batch["labels"] = labels_flat.view_as(labels) return batch - - diff --git a/src/xorl/data/collators/sequence_shard_collator.py b/src/xorl/data/collators/sequence_shard_collator.py index f13df1d7..6493b059 100644 --- a/src/xorl/data/collators/sequence_shard_collator.py +++ b/src/xorl/data/collators/sequence_shard_collator.py @@ -1,14 +1,13 @@ from dataclasses import dataclass -from typing import Dict, List, Sequence +from typing import Dict import torch -import torch.nn.functional as F -from .base_collator import DataCollator -from .packing_concat_collator import add_flash_attention_kwargs_from_position_ids -from ...utils.seqlen_pos_transform_utils import prepare_fa_kwargs_from_position_ids from ...data.constants import IGNORE_INDEX from ...distributed.parallel_state import get_parallel_state +from ...utils.seqlen_pos_transform_utils import prepare_fa_kwargs_from_position_ids +from .base_collator import DataCollator +from .packing_concat_collator import add_flash_attention_kwargs_from_position_ids def zigzag_reorder_packed_sequence( @@ -70,8 +69,8 @@ def zigzag_reorder_packed_sequence( chunks = list(doc.chunk(n, dim=dim)) for r in range(ringattn_size): - rank_parts[r].append(chunks[r]) # early sub-chunk - rank_parts[r].append(chunks[n - 1 - r]) # late sub-chunk + rank_parts[r].append(chunks[r]) # early sub-chunk + rank_parts[r].append(chunks[n - 1 - r]) # late sub-chunk # Concatenate: rank 0's data first, then rank 1's, etc. all_parts = [] @@ -216,9 +215,7 @@ def __call__(self, batch: Dict[str, "torch.Tensor"]) -> Dict[str, "torch.Tensor" # 1. cu_seqlens is computed from FULL padded position_ids # 2. For Ulysses, all SP ranks use the SAME cu_seqlens for flash attention # 3. Each SP rank only processes a slice of the sequence but needs full cu_seqlens - position_ids = self.sp_padding( - position_ids, dim=-1, pad_value=0, pad_length=pad_length, sequential=True - ) + position_ids = self.sp_padding(position_ids, dim=-1, pad_value=0, pad_length=pad_length, sequential=True) # Zigzag reorder: rearrange each document's tokens so that contiguous # sp_slice gives each CP rank balanced [early, late] sub-chunks. @@ -233,7 +230,9 @@ def __call__(self, batch: Dict[str, "torch.Tensor"]) -> Dict[str, "torch.Tensor" batch["attention_mask"], original_position_ids, self.ringattn_size, dim=-1 ) # Reorder position_ids last (uses its own original values for boundaries) - position_ids = zigzag_reorder_packed_sequence(position_ids, original_position_ids, self.ringattn_size, dim=-1) + position_ids = zigzag_reorder_packed_sequence( + position_ids, original_position_ids, self.ringattn_size, dim=-1 + ) # sp slice - only slice input_ids and labels, NOT position_ids batch["input_ids"] = self.sp_slice(input_ids, dim=-1) diff --git a/src/xorl/data/collators/shift_tokens_collator.py b/src/xorl/data/collators/shift_tokens_collator.py index adf9c85a..254d71a3 100644 --- a/src/xorl/data/collators/shift_tokens_collator.py +++ b/src/xorl/data/collators/shift_tokens_collator.py @@ -23,6 +23,7 @@ from .base_collator import DataCollator + logger = logging.getLogger(__name__) diff --git a/src/xorl/data/collators/tensor_collator.py b/src/xorl/data/collators/tensor_collator.py index c6c87721..caff02ca 100644 --- a/src/xorl/data/collators/tensor_collator.py +++ b/src/xorl/data/collators/tensor_collator.py @@ -25,7 +25,7 @@ class ToTensorCollator(DataCollator): Returns list of dicts with converted tensors, preserving the structure. Downstream collators (like PackingConcatCollator) handle batching/concatenation. """ - + def __call__(self, features: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]: """ Convert lists/arrays to tensors without batching. @@ -78,7 +78,7 @@ def __call__(self, features: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]: # Return list of dicts (not batched) # PackingConcatCollator will handle the concatenation/batching return converted_features - + def _convert_value(self, value: Any, key: str) -> Any: """ Recursively convert a value to tensor if appropriate. @@ -154,21 +154,20 @@ def _is_numeric_list(self, lst: list) -> bool: # Flat list - check if all are numeric return all(isinstance(item, (int, float, bool, np.number)) for item in lst) - + def _infer_dtype(self, key: str) -> torch.dtype: """ Infer appropriate dtype based on field name. - + Args: key: Field name - + Returns: PyTorch dtype """ # These fields should be long (int64) for model compatibility - if key in ['input_ids', 'labels', 'attention_mask', 'position_ids']: + if key in ["input_ids", "labels", "attention_mask", "position_ids"]: return torch.long - + # Let PyTorch infer for other fields return None # Will use default inference - diff --git a/src/xorl/data/constants.py b/src/xorl/data/constants.py index 07b3de5d..f9e3d8ad 100644 --- a/src/xorl/data/constants.py +++ b/src/xorl/data/constants.py @@ -1,4 +1,4 @@ IGNORE_INDEX = -100 IMAGE_INPUT_INDEX = -200 VIDEO_INPUT_INDEX = -300 -AUDIO_INPUT_INDEX = -400 \ No newline at end of file +AUDIO_INPUT_INDEX = -400 diff --git a/src/xorl/data/data_loader.py b/src/xorl/data/data_loader.py index 1d8db5d1..00c1e090 100644 --- a/src/xorl/data/data_loader.py +++ b/src/xorl/data/data_loader.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence from torchdata.stateful_dataloader import StatefulDataLoader from torchdata.stateful_dataloader.sampler import StatefulDistributedSampler @@ -25,17 +25,17 @@ class MicroBatchCollator(DataCollator): """ Collator that splits a large batch into multiple micro-batches for gradient accumulation. - + This collator receives a batch of size (micro_batch_size * gradient_accumulation_steps) and splits it into gradient_accumulation_steps separate micro-batches, each of size micro_batch_size. Each micro-batch is then processed by the internal_collator. - + Args: micro_batch_size: Size of each micro-batch gradient_accumulation_steps: Number of micro-batches to create internal_collator: The collator to apply to each micro-batch """ - + def __init__( self, micro_batch_size: int, @@ -45,16 +45,14 @@ def __init__( self.micro_batch_size = micro_batch_size self.gradient_accumulation_steps = gradient_accumulation_steps self.internal_collator = internal_collator - - def __call__( - self, features: Sequence[Dict[str, Any]] - ) -> List[Dict[str, Any]]: + + def __call__(self, features: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]: """ Split features into micro-batches and collate each one. - + Args: features: List of samples (length = micro_batch_size * gradient_accumulation_steps) - + Returns: List of collated micro-batches (length = gradient_accumulation_steps) """ @@ -67,18 +65,17 @@ def __call__( f"This usually means the dataset length is not divisible by the dataloader batch size. " f"Consider setting drop_last=True in build_dataloader." ) - + micro_batches = [] for i in range(0, len(features), self.micro_batch_size): micro_batch_features = features[i : i + self.micro_batch_size] collated_micro_batch = self.internal_collator(micro_batch_features) micro_batches.append(collated_micro_batch) - + assert len(micro_batches) == self.gradient_accumulation_steps, ( - f"Internal error: Expected {self.gradient_accumulation_steps} micro-batches, " - f"but got {len(micro_batches)}" + f"Internal error: Expected {self.gradient_accumulation_steps} micro-batches, but got {len(micro_batches)}" ) - + return micro_batches @@ -96,36 +93,36 @@ def set_epoch(self, epoch: int) -> None: class DataLoaderBuilder: """ Builder class for constructing distributed dataloaders with gradient accumulation support. - + This class provides a flexible interface for building dataloaders, allowing customization of collators, samplers, and other components through method chaining or inheritance. - + Example: # Basic usage builder = DataLoaderBuilder(dataset, micro_batch_size=2, gradient_accumulation_steps=4) dataloader = builder.build() - + # Add custom collators to the pipeline builder = DataLoaderBuilder(dataset, micro_batch_size=2, gradient_accumulation_steps=4) builder.add_collator(MyCustomCollator(), position="end") builder.print_pipeline() # Visualize the pipeline dataloader = builder.build() - + # Method chaining dataloader = (DataLoaderBuilder(dataset, micro_batch_size=2, gradient_accumulation_steps=4) .add_collator(MyPreprocessor(), position="start") .add_collator(MyPostprocessor(), position="end") .build()) - + # Custom builder subclass class CustomBuilder(DataLoaderBuilder): def _build_default_collator_list(self): return [MyCustomCollator()] - + builder = CustomBuilder(dataset, micro_batch_size=2, gradient_accumulation_steps=4) dataloader = builder.build() """ - + def __init__( self, dataset: "Dataset", @@ -166,41 +163,41 @@ def __init__( self.seed = seed self.pad_to_multiple_of = pad_to_multiple_of self.parallel_state = get_parallel_state() - + # Initialize collator pipeline if use_default_collators: self._collator_list: List[DataCollator] = self._build_default_collator_list() else: self._collator_list: List[DataCollator] = [] - + # Final collator and sampler will be set when build() is called self.collate_fn = None self.sampler = None - + def _build_default_collator_list(self) -> List[DataCollator]: """Build the default collator pipeline.""" from .collators import ToTensorCollator collators = [ - ToTensorCollator(), # 1. Convert to tensors - FlattenCollator(), # 2. Flatten list of lists to flat list - ShiftTokensCollator(auto_detect=True), # 3. Shift tokens for causal LM (auto-detects if needed) - PackingConcatCollator(pad_to_multiple_of=self.pad_to_multiple_of) # 4. Concatenate sequences for packing + ToTensorCollator(), # 1. Convert to tensors + FlattenCollator(), # 2. Flatten list of lists to flat list + ShiftTokensCollator(auto_detect=True), # 3. Shift tokens for causal LM (auto-detects if needed) + PackingConcatCollator(pad_to_multiple_of=self.pad_to_multiple_of), # 4. Concatenate sequences for packing ] if self.parallel_state.cp_enabled: collators.append(TextSequenceShardCollator()) # 5. Shard sequences (if SP enabled) return collators - + def add_collator(self, collator: DataCollator, position: str = "end") -> "DataLoaderBuilder": """ Add a collator to the pipeline. - + Args: collator: The collator to add position: Where to add ('start' or 'end') - + Returns: Self for method chaining """ @@ -211,80 +208,80 @@ def add_collator(self, collator: DataCollator, position: str = "end") -> "DataLo else: raise ValueError(f"Invalid position: {position}. Must be 'start' or 'end'") return self - + def insert_collator(self, collator: DataCollator, index: int) -> "DataLoaderBuilder": """ Insert a collator at a specific position in the pipeline. - + Args: collator: The collator to insert index: The position to insert at - + Returns: Self for method chaining """ self._collator_list.insert(index, collator) return self - + def remove_collator(self, index: int) -> "DataLoaderBuilder": """ Remove a collator from the pipeline. - + Args: index: The index of the collator to remove - + Returns: Self for method chaining """ self._collator_list.pop(index) return self - + def get_collator_pipeline(self) -> List[DataCollator]: """ Get the current list of collators in the pipeline. - + Returns: List of collators """ return self._collator_list.copy() - + def print_pipeline(self) -> None: """ Print the current collator pipeline in a readable format. - + This shows the order and types of all collators that will be applied to each micro-batch during data loading. """ logger.info_rank0("=" * 60) logger.info_rank0("Collator Pipeline:") logger.info_rank0("=" * 60) - + if not self._collator_list: logger.info_rank0(" (empty pipeline)") else: for i, collator in enumerate(self._collator_list, 1): collator_name = collator.__class__.__name__ logger.info_rank0(f" {i}. {collator_name}") - + # If it's a CollatePipeline, show its internal collators if isinstance(collator, CollatePipeline): for j, sub_collator in enumerate(collator.data_collators, 1): sub_collator_name = sub_collator.__class__.__name__ logger.info_rank0(f" {i}.{j} {sub_collator_name}") - + logger.info_rank0("-" * 60) - logger.info_rank0(f"Micro-batch splitting:") + logger.info_rank0("Micro-batch splitting:") logger.info_rank0(f" micro_batch_size: {self.micro_batch_size}") logger.info_rank0(f" gradient_accumulation_steps: {self.gradient_accumulation_steps}") logger.info_rank0(f" dataloader_batch_size: {self.micro_batch_size * self.gradient_accumulation_steps}") logger.info_rank0("=" * 60) - + def build_collator(self) -> MicroBatchCollator: """ Build the final collator that wraps the pipeline with micro-batch splitting. - + Override this method to customize the entire collator construction logic. - + Returns: A MicroBatchCollator that wraps the collator pipeline """ @@ -295,20 +292,20 @@ def build_collator(self) -> MicroBatchCollator: internal_collator = self._collator_list[0] else: internal_collator = CollatePipeline(self._collator_list) - + # Wrap with MicroBatchCollator for gradient accumulation return MicroBatchCollator( micro_batch_size=self.micro_batch_size, gradient_accumulation_steps=self.gradient_accumulation_steps, internal_collator=internal_collator, ) - + def build_sampler(self) -> StatefulDistributedSampler: """ Build the distributed sampler. - + Override this method to customize the sampler (e.g., change shuffle behavior). - + Returns: A StatefulDistributedSampler instance """ @@ -319,25 +316,24 @@ def build_sampler(self) -> StatefulDistributedSampler: shuffle=True, seed=self.seed, ) - + def build(self, verbose: bool = True) -> "DistributedDataloader": """ Build and return the final dataloader. - + Args: verbose: If True, prints the collator pipeline configuration - + Returns: A DistributedDataloader configured with all components """ if verbose: self.print_pipeline() - + dataloader_batch_size = self.micro_batch_size * self.gradient_accumulation_steps - + logger.info_rank0( - f"Building DataLoader with dp_size={self.parallel_state.dp_size}, " - f"cp_size={self.parallel_state.cp_size}" + f"Building DataLoader with dp_size={self.parallel_state.dp_size}, cp_size={self.parallel_state.cp_size}" ) # Validate dataset structure - each item can be either: @@ -376,5 +372,5 @@ def build(self, verbose: bool = True) -> "DistributedDataloader": drop_last=self.drop_last, prefetch_factor=self.prefetch_factor, ) - + return dataloader diff --git a/src/xorl/data/prepare/__init__.py b/src/xorl/data/prepare/__init__.py index e00eb420..0d94736f 100644 --- a/src/xorl/data/prepare/__init__.py +++ b/src/xorl/data/prepare/__init__.py @@ -1 +1 @@ -DEFAULT_DATASET_PREPARED_PATH = "last_prepared_dataset" \ No newline at end of file +DEFAULT_DATASET_PREPARED_PATH = "last_prepared_dataset" diff --git a/src/xorl/data/prepare/constants.py b/src/xorl/data/prepare/constants.py index e00eb420..0d94736f 100644 --- a/src/xorl/data/prepare/constants.py +++ b/src/xorl/data/prepare/constants.py @@ -1 +1 @@ -DEFAULT_DATASET_PREPARED_PATH = "last_prepared_dataset" \ No newline at end of file +DEFAULT_DATASET_PREPARED_PATH = "last_prepared_dataset" diff --git a/src/xorl/data/prepare/file_lock_loader.py b/src/xorl/data/prepare/file_lock_loader.py index a193e01d..bec17974 100644 --- a/src/xorl/data/prepare/file_lock_loader.py +++ b/src/xorl/data/prepare/file_lock_loader.py @@ -9,6 +9,7 @@ from ...arguments import Arguments from .constants import DEFAULT_DATASET_PREPARED_PATH + LOCK_FILE_NAME = "datasets_prep.lock" READY_FILE_NAME = "datasets_ready.flag" PROCESS_COUNTER_FILE_NAME = "process_counter.txt" @@ -23,9 +24,7 @@ class FileLockLoader: def __init__(self, args: Arguments): self.args = args - self.dataset_prepared_path = ( - args.data.dataset_prepared_path or DEFAULT_DATASET_PREPARED_PATH - ) + self.dataset_prepared_path = args.data.dataset_prepared_path or DEFAULT_DATASET_PREPARED_PATH self.lock_file_path = Path(self.dataset_prepared_path) / LOCK_FILE_NAME self.ready_flag_path = Path(self.dataset_prepared_path) / READY_FILE_NAME self.counter_path = Path(self.dataset_prepared_path) / PROCESS_COUNTER_FILE_NAME @@ -33,7 +32,7 @@ def __init__(self, args: Arguments): def load(self, load_fn: Callable[[], Any]) -> Any: # Ensure directory exists Path(self.dataset_prepared_path).mkdir(parents=True, exist_ok=True) - + with FileLock(str(self.lock_file_path)): self._increment_counter() @@ -58,7 +57,7 @@ def _increment_counter(self): else: count = 0 self.counter_path.write_text(str(count + 1)) - except (ValueError, OSError) as e: + except (ValueError, OSError): # Handle corrupted counter file or I/O errors # Reset to 1 for this process self.counter_path.write_text("1") @@ -78,7 +77,7 @@ def cleanup(self): else: # Still have active processes self.counter_path.write_text(str(count)) - except (ValueError, OSError) as e: + except (ValueError, OSError): # Handle corrupted counter file or I/O errors # Force cleanup since we can't determine the count self.ready_flag_path.unlink(missing_ok=True) diff --git a/src/xorl/data/prepare/hash.py b/src/xorl/data/prepare/hash.py index 1c37c0c7..33e20293 100644 --- a/src/xorl/data/prepare/hash.py +++ b/src/xorl/data/prepare/hash.py @@ -1,13 +1,12 @@ from typing import List +from datasets import Dataset + from ...arguments import Arguments, DatasetConfig from .utils import md5 -from datasets import Dataset -def generate_split_fingerprints( - dataset: Dataset, val_set_size: int | float, seed: int -) -> tuple[str, str]: +def generate_split_fingerprints(dataset: Dataset, val_set_size: int | float, seed: int) -> tuple[str, str]: """Generate consistent fingerprints for train/test splits.""" fingerprint = dataset._fingerprint @@ -40,9 +39,8 @@ def generate_packing_hash( packing_str += f"_align{doc_align}" return packing_str -def generate_dataset_hash_from_config( - args: Arguments, args_datasets: List[DatasetConfig], tokenizer_name: str -) -> str: + +def generate_dataset_hash_from_config(args: Arguments, args_datasets: List[DatasetConfig], tokenizer_name: str) -> str: """Generate a hash to uniquely identify a dataset configuration for SFT. Args: diff --git a/src/xorl/data/prepare/packing.py b/src/xorl/data/prepare/packing.py index 2b924ad1..ac780339 100644 --- a/src/xorl/data/prepare/packing.py +++ b/src/xorl/data/prepare/packing.py @@ -3,29 +3,24 @@ into fixed-capacity batches to optimize memory usage and training throughput. """ -from typing import Any, Dict, List, Optional, Union, Callable -import gc -import hashlib -import json -import math import os -import pickle import time from concurrent.futures import ProcessPoolExecutor from multiprocessing import cpu_count, get_context from pathlib import Path -from typing import Iterable, Iterator +from typing import Any, Callable, Dict, List, Optional -from torch.utils.data import Dataset -from datasets import Dataset as HFDataset -from transformers import PreTrainedTokenizer import numba import numpy as np +from datasets import Dataset as HFDataset +from torch.utils.data import Dataset +from transformers import PreTrainedTokenizer -from ...utils import logging from ...arguments import Arguments -from .shared import get_prepared_dataset_path +from ...utils import logging from .hash import generate_dataset_hash_from_config, generate_packing_hash +from .shared import get_prepared_dataset_path + LOG = logging.get_logger(__name__) @@ -97,10 +92,7 @@ def pack_group( # Try to place sequence in existing bins add_new_bin = True for bin_idx, _ in enumerate(bins_remaining_space): - if ( - bins_remaining_space[bin_idx] >= size - and len(bins_assigned_sequences[bin_idx]) < bin_size - ): + if bins_remaining_space[bin_idx] >= size and len(bins_assigned_sequences[bin_idx]) < bin_size: bins_remaining_space[bin_idx] -= size bins_assigned_sequences[bin_idx].append(global_idx) add_new_bin = False @@ -126,9 +118,7 @@ def _process_group( ) -> list[list[int]]: """Standalone function for multiprocessing.""" group_lengths, start_idx, bin_capacity, max_bins, bin_size, safe_mode = args - return pack_group( - group_lengths, start_idx, bin_capacity, max_bins, bin_size, safe_mode - ) + return pack_group(group_lengths, start_idx, bin_capacity, max_bins, bin_size, safe_mode) def pack_parallel( @@ -178,9 +168,7 @@ def pack_parallel( f"Failed to get multiprocessing context '{mp_start_method}'. " f"Falling back to default. Available: {get_context().get_all_start_methods()}" ) - mp_ctx = ( - None # Fallback to default context if specified one is not available - ) + mp_ctx = None # Fallback to default context if specified one is not available if num_processes == 1: LOG.debug("Using single process for pack_parallel, running sequentially.") @@ -190,9 +178,7 @@ def pack_parallel( else: # Use ProcessPoolExecutor only if num_processes > 1 # Pass mp_context if available - with ProcessPoolExecutor( - max_workers=num_processes, mp_context=mp_ctx - ) as executor: + with ProcessPoolExecutor(max_workers=num_processes, mp_context=mp_ctx) as executor: for group_bins in executor.map(_process_group, tasks): all_bins.extend(group_bins) @@ -312,9 +298,7 @@ def drop_no_trainable_tokens(sample: Dict[str, Any]): # drop samples with no trainable tokens -def filter_dataset_with_logging( - dataset: HFDataset, filter_func: Callable, dataset_name: str, num_proc: int -): +def filter_dataset_with_logging(dataset: HFDataset, filter_func: Callable, dataset_name: str, num_proc: int): """Filter dataset and log dropped samples.""" try: prior_len = len(dataset) @@ -331,16 +315,12 @@ def filter_dataset_with_logging( if prior_len: dropped = prior_len - len(filtered_dataset) if dropped: - LOG.warning( - f"Dropped {dropped} samples with no trainable tokens from {dataset_name} dataset" - ) + LOG.warning(f"Dropped {dropped} samples with no trainable tokens from {dataset_name} dataset") return filtered_dataset -def process_datasets_for_packing( - args: Arguments, train_dataset: HFDataset, eval_dataset: HFDataset | None = None -): +def process_datasets_for_packing(args: Arguments, train_dataset: HFDataset, eval_dataset: HFDataset | None = None): # drop samples with no trainable tokens train_dataset = filter_dataset_with_logging( train_dataset, drop_no_trainable_tokens, "train", args.data.dataset_num_proc @@ -370,16 +350,16 @@ def process_datasets_for_packing( class PackingDataset(Dataset): - def __init__(self, args: Arguments, tokenizer: PreTrainedTokenizer, dataset: HFDataset, split: str = "train") -> None: + def __init__( + self, args: Arguments, tokenizer: PreTrainedTokenizer, dataset: HFDataset, split: str = "train" + ) -> None: self.dataset: HFDataset = dataset self.args: Arguments = args datasets_configs = args.data.datasets if split == "train" else args.data.test_datasets - self.dataset_hash: str = generate_dataset_hash_from_config( - args, datasets_configs, tokenizer.name_or_path - ) + self.dataset_hash: str = generate_dataset_hash_from_config(args, datasets_configs, tokenizer.name_or_path) self.prepared_dataset_path: str | None = get_prepared_dataset_path(args, self.dataset_hash) - + # Get sequence lengths from the dataset assert "length" in self.dataset.features, "Length column not found in dataset" self.sequence_lengths: np.ndarray = np.array(dataset["length"]) @@ -414,7 +394,7 @@ def _get_bins_cache_path(self) -> Path | None: """Get the path where bins should be cached as HF dataset.""" if self.prepared_dataset_path is None: return None - + # Include packing args hash in the cache path packing_hash = generate_packing_hash( self.args.data.sample_packing_method, @@ -440,7 +420,7 @@ def _load_cached_bins(self) -> Optional[List[List[int]]]: # So if the file exists and loads successfully, it's valid # Extract bins from the dataset - bins = bins_dataset['bins'] + bins = bins_dataset["bins"] LOG.info(f"Loaded {len(bins)} cached bins from {cache_path}") return bins @@ -453,28 +433,25 @@ def _save_bins_cache(self, bins: List[List[int]]) -> None: cache_path = self._get_bins_cache_path() if cache_path is None: return - + # Ensure the directory exists cache_path.parent.mkdir(parents=True, exist_ok=True) - + # Create a HF dataset from the bins - bins_dataset = HFDataset.from_dict({ - 'bins': bins - }) - + bins_dataset = HFDataset.from_dict({"bins": bins}) + # Add metadata to the dataset info bins_dataset.info.description = ( f"Packing bins for dataset fingerprint: {self.dataset_hash}" f" | packing method: {self.args.data.sample_packing_method}" f" | bin capacity: {self.bin_capacity}" ) - + # Save the dataset num_workers = max(self.args.data.dataset_num_proc // 8, 1) bins_dataset.save_to_disk(str(cache_path), num_proc=num_workers) - + LOG.info(f"Saved bins cache to {cache_path}") - def _load_or_compute_bins(self) -> List[List[int]]: """Load cached bins or compute new ones with distributed coordination.""" @@ -549,27 +526,22 @@ def _compute_bins(self) -> List[List[int]]: return self._compute_sequential_bins() else: # default to multipack return self._compute_multipack_bins() - + def _compute_sequential_bins(self) -> List[List[int]]: """Compute bins using sequential packing.""" # Use the sequential allocator from the existing code rank: int = 0 # For single-process case num_ranks: int = 1 # For single-process case - - bins, _, _ = allocate_sequentially( - self.sequence_lengths, - rank, - self.bin_capacity, - num_ranks - ) + + bins, _, _ = allocate_sequentially(self.sequence_lengths, rank, self.bin_capacity, num_ranks) return bins - + def _compute_multipack_bins(self) -> List[List[int]]: """Compute bins using multipack algorithm.""" group_size: int = self.args.data.sample_packing_group_size or 100000 bin_size: int = self.bin_capacity # Maximum sequences per bin mp_start_method: str = self.args.data.sample_packing_mp_start_method or "fork" - + # Use the parallel packing function bins: List[List[int]] = pack_parallel( sequence_lengths=self.sequence_lengths, @@ -578,7 +550,7 @@ def _compute_multipack_bins(self) -> List[List[int]]: bin_size=bin_size, num_processes=1, # Use single process for simplicity safe_mode=True, - mp_start_method=mp_start_method + mp_start_method=mp_start_method, ) return bins @@ -597,4 +569,3 @@ def __getitem__(self, index: int) -> Dict[str, Any]: samples: List[Dict[str, Any]] = [self.dataset[idx] for idx in bin_indices] return samples - \ No newline at end of file diff --git a/src/xorl/data/prepare/prepare_datasets.py b/src/xorl/data/prepare/prepare_datasets.py index 910a516c..ecce8a2e 100644 --- a/src/xorl/data/prepare/prepare_datasets.py +++ b/src/xorl/data/prepare/prepare_datasets.py @@ -1,36 +1,38 @@ """Data handling specific to SFT.""" -import functools -import os -import tempfile from typing import Literal +import torch.distributed as dist from datasets import ( Dataset as HFDataset, +) +from datasets import ( DatasetDict as HFDatasetDict, +) +from datasets import ( IterableDataset as HFIterableDataset, +) +from datasets import ( IterableDatasetDict as HFIterableDatasetDict, - load_dataset, ) -from transformers import PreTrainedTokenizer, ProcessorMixin from torch.utils.data import Dataset -import torch.distributed as dist +from transformers import PreTrainedTokenizer, ProcessorMixin from ...arguments import Arguments, DatasetConfig -from ...utils import logging -from ...data.prepare.file_lock_loader import FileLockLoader from ...data.prepare.utils import retry_on_request_exceptions +from ...utils import logging +from .hash import generate_dataset_hash_from_config +from .packing import PackingDataset, process_datasets_for_packing from .shared import ( - load_dataset_with_config, + create_train_validation_split, datasets_with_name_generator, - merge_datasets, - try_load_from_hub, + load_dataset_with_config, load_preprocessed_dataset, + merge_datasets, save_preprocessed_dataset, - create_train_validation_split, + try_load_from_hub, ) -from .hash import generate_dataset_hash_from_config -from .packing import PackingDataset, process_datasets_for_packing + logger = logging.get_logger(__name__) @@ -91,12 +93,14 @@ def _create_dummy_dataset(seq_len: int, num_samples: int = 4096, seed: int = 42, pa_lengths = pa.array(lengths.tolist(), type=pa.int64()) - table = pa.table({ - "input_ids": tokens, - "labels": tokens, - "position_ids": position_ids, - "length": pa_lengths, - }) + table = pa.table( + { + "input_ids": tokens, + "labels": tokens, + "position_ids": position_ids, + "length": pa_lengths, + } + ) return HFDataset(table) @@ -124,28 +128,17 @@ def prepare_datasets( def _load_datasets(): # Load training dataset - train_dataset, eval_dataset = _load_and_prepare_datasets( - tokenizer, args, split="train", processor=processor - ) + train_dataset, eval_dataset = _load_and_prepare_datasets(tokenizer, args, split="train", processor=processor) # Override with test dataset if available if args.data.test_datasets: - _, eval_dataset = _load_and_prepare_datasets( - tokenizer, args, split="test", processor=processor - ) + _, eval_dataset = _load_and_prepare_datasets(tokenizer, args, split="test", processor=processor) # Apply sample packing if configured - if ( - args.data.sample_packing_method - and args.data.sample_packing_method != "none" - ): - train_dataset = PackingDataset( - args, tokenizer, train_dataset, split="train" - ) + if args.data.sample_packing_method and args.data.sample_packing_method != "none": + train_dataset = PackingDataset(args, tokenizer, train_dataset, split="train") if eval_dataset: - eval_dataset = PackingDataset( - args, tokenizer, eval_dataset, split="test" - ) + eval_dataset = PackingDataset(args, tokenizer, eval_dataset, split="test") else: eval_dataset = None @@ -196,14 +189,10 @@ def _load_tokenized_prepared_datasets( is_rank_zero = not is_distributed or dist.get_rank() == 0 # Select correct dataset configuration based on split - datasets_configs = ( - args.data.datasets if split == "train" else args.data.test_datasets - ) + datasets_configs = args.data.datasets if split == "train" else args.data.test_datasets # Generate dataset hash for caching - dataset_hash = generate_dataset_hash_from_config( - args, datasets_configs, tokenizer.name_or_path - ) + dataset_hash = generate_dataset_hash_from_config(args, datasets_configs, tokenizer.name_or_path) # Try loading from hub if push_dataset_to_hub is configured dataset = None @@ -279,9 +268,7 @@ def _load_raw_datasets( # Save the prepared dataset if not args.data.skip_prepare_dataset: - dataset_hash = generate_dataset_hash_from_config( - args, datasets_configs, tokenizer.name_or_path - ) + dataset_hash = generate_dataset_hash_from_config(args, datasets_configs, tokenizer.name_or_path) save_preprocessed_dataset(args, dataset, dataset_hash, split) return dataset @@ -298,7 +285,9 @@ def _load_and_process_single_dataset( """Load and process a single dataset based on the passed config.""" # Load the dataset dataset = load_dataset_with_config( - dataset_config, args.data.hf_use_auth_token, streaming=False, + dataset_config, + args.data.hf_use_auth_token, + streaming=False, num_proc=args.data.dataset_num_proc, ) @@ -313,16 +302,13 @@ def _load_and_process_single_dataset( dataset = dataset[split] else: raise ValueError( - f"no {split} split found for dataset {dataset_config.path}, you may " - "specify a split with 'split: ...'" + f"no {split} split found for dataset {dataset_config.path}, you may specify a split with 'split: ...'" ) # Apply sharding if configured if dataset_config.shards: shards_idx = dataset_config.shards_idx or 0 - dataset = dataset.shuffle(seed=seed).shard( - num_shards=dataset_config.shards, index=shards_idx - ) + dataset = dataset.shuffle(seed=seed).shard(num_shards=dataset_config.shards, index=shards_idx) # Select columns if configured if args.data.select_columns: @@ -350,11 +336,11 @@ def _load_and_process_single_dataset( # Add activations_path field if configured (for distillation training) if dataset_config.activations_path: logger.info_rank0(f"Adding activations_path '{dataset_config.activations_path}' to dataset") - + def add_activations_path(example): example["activations_path"] = dataset_config.activations_path return example - + dataset = dataset.map( add_activations_path, desc=f"Adding activations_path to {dataset_config.name or 'dataset'}", @@ -364,22 +350,19 @@ def add_activations_path(example): return dataset -def _handle_train_dataset_split( - dataset: HFDataset, args: Arguments -) -> tuple[HFDataset, HFDataset | None]: +def _handle_train_dataset_split(dataset: HFDataset, args: Arguments) -> tuple[HFDataset, HFDataset | None]: """Handle processing for train split, including validation set creation.""" val_set_size = ( int(args.data.val_set_size) if args.data.val_set_size is not None and args.data.val_set_size > 1 - else float(args.data.val_set_size) if args.data.val_set_size is not None + else float(args.data.val_set_size) + if args.data.val_set_size is not None else 0 ) if val_set_size: # Create train/validation split - train_dataset, eval_dataset = create_train_validation_split( - dataset, args, val_set_size - ) + train_dataset, eval_dataset = create_train_validation_split(dataset, args, val_set_size) return train_dataset, eval_dataset else: return dataset, None @@ -396,9 +379,7 @@ def _apply_dataset_sharding(dataset: HFDataset, args: Arguments) -> HFDataset: Sharded dataset or original dataset if no sharding configured. """ if args.data.dataset_shard_num and args.data.dataset_shard_idx is not None: - logger.info_rank0( - f"Using index #{args.data.dataset_shard_idx} of {args.data.dataset_shard_num} shards" - ) + logger.info_rank0(f"Using index #{args.data.dataset_shard_idx} of {args.data.dataset_shard_num} shards") dataset = dataset.shard( num_shards=args.data.dataset_shard_num, index=args.data.dataset_shard_idx, diff --git a/src/xorl/data/prepare/shared.py b/src/xorl/data/prepare/shared.py index d6d98467..5a641bff 100644 --- a/src/xorl/data/prepare/shared.py +++ b/src/xorl/data/prepare/shared.py @@ -3,17 +3,15 @@ from __future__ import annotations import functools -import hashlib import os from pathlib import Path -from typing import List, Dict, TYPE_CHECKING, Any, Generator +from typing import TYPE_CHECKING, Any, Generator from datasets import ( Dataset, DatasetDict, IterableDataset, IterableDatasetDict, - interleave_datasets, concatenate_datasets, load_dataset, load_dataset_builder, @@ -27,9 +25,9 @@ ) from ...arguments import Arguments, DatasetConfig -from ...data.prepare.hash import generate_split_fingerprints from ...data.prepare.constants import DEFAULT_DATASET_PREPARED_PATH -from ...data.prepare.utils import md5 +from ...data.prepare.hash import generate_split_fingerprints + if TYPE_CHECKING: try: @@ -41,6 +39,8 @@ pass from ...utils import logging + + LOG = logging.get_logger(__name__) EXTENSIONS_TO_DATASET_TYPES = { @@ -51,8 +51,6 @@ } - - def get_dataset_type(dataset_config: DatasetConfig) -> str: """Get the dataset type from the path if it's not specified.""" if dataset_config.ds_type: @@ -148,9 +146,7 @@ def load_dataset_with_config( if is_hub_dataset: return _load_from_hub(dataset_config, use_auth_token, load_dataset_kwargs) if is_cloud_dataset: - return _load_from_cloud( - dataset_config, remote_fs, storage_options, load_dataset_kwargs - ) + return _load_from_cloud(dataset_config, remote_fs, storage_options, load_dataset_kwargs) if dataset_config.path.startswith("https://"): return _load_from_url(dataset_config, load_dataset_kwargs) if dataset_config.data_files: @@ -187,9 +183,7 @@ def _check_if_hub_dataset(dataset_config: DatasetConfig, use_auth_token: bool) - def _get_remote_filesystem( path: str, -) -> tuple[ - S3FileSystem | GCSFileSystem | AzureBlobFileSystem | OCIFileSystem | None, dict -]: +) -> tuple[S3FileSystem | GCSFileSystem | AzureBlobFileSystem | OCIFileSystem | None, dict]: """Get the appropriate filesystem for a remote path.""" if path.startswith("s3://"): try: @@ -207,9 +201,7 @@ def _get_remote_filesystem( storage_options = {"token": None} # type: ignore return gcsfs.GCSFileSystem(**storage_options), storage_options except ImportError as exc: - raise ImportError( - "gs:// or gcs:// paths require gcsfs to be installed" - ) from exc + raise ImportError("gs:// or gcs:// paths require gcsfs to be installed") from exc elif path.startswith(("adl://", "abfs://", "az://")): try: @@ -218,9 +210,7 @@ def _get_remote_filesystem( storage_options = {"anon": False} return adlfs.AzureBlobFileSystem(**storage_options), storage_options except ImportError as exc: - raise ImportError( - "adl:// or abfs:// paths require adlfs to be installed" - ) from exc + raise ImportError("adl:// or abfs:// paths require adlfs to be installed") from exc elif path.startswith("oci://"): try: @@ -260,9 +250,7 @@ def _load_from_local_path( **load_dataset_kwargs, ) else: - raise ValueError( - "Unhandled dataset load: local path exists, but is neither a directory or a file" - ) + raise ValueError("Unhandled dataset load: local path exists, but is neither a directory or a file") def _load_from_hub( @@ -287,11 +275,7 @@ def _load_from_hub( trust_remote_code=dataset_config.trust_remote_code, ) data_files = getattr(getattr(builder, "config", None), "data_files", None) - if ( - data_files - and requested_split in data_files - and len(data_files) > 1 - ): + if data_files and requested_split in data_files and len(data_files) > 1: # Load only the parquet files for the requested split, skipping others split_files = data_files[requested_split] kwargs = {k: v for k, v in load_dataset_kwargs.items() if k != "split"} @@ -314,6 +298,7 @@ def _load_from_hub( **load_dataset_kwargs, ) + def _load_from_cloud( dataset_config: DatasetConfig, remote_fs: S3FileSystem | GCSFileSystem | AzureBlobFileSystem | OCIFileSystem, @@ -336,9 +321,7 @@ def _load_from_cloud( **load_dataset_kwargs, ) - raise ValueError( - f"Cloud path {dataset_config.path} is neither a directory nor a file" - ) + raise ValueError(f"Cloud path {dataset_config.path} is neither a directory nor a file") def _load_from_url( @@ -382,8 +365,6 @@ def _load_from_data_files( return load_dataset("json", data_files=file_path, **load_dataset_kwargs) - - def get_prepared_dataset_path(args: Arguments, dataset_hash: str) -> Path: """Get standardized path for prepared datasets. @@ -411,9 +392,7 @@ def create_train_validation_split( Returns: Tuple of (train_dataset, eval_dataset). """ - train_fingerprint, test_fingerprint = generate_split_fingerprints( - dataset, val_set_size, args.train.seed - ) + train_fingerprint, test_fingerprint = generate_split_fingerprints(dataset, val_set_size, args.train.seed) split_dataset = dataset.train_test_split( test_size=val_set_size, @@ -511,9 +490,7 @@ def load_preprocessed_dataset(args: Arguments, dataset_hash: str) -> Dataset | N return None -def try_load_from_hub( - args: Arguments, dataset_hash: str, split: str -) -> Dataset | None: +def try_load_from_hub(args: Arguments, dataset_hash: str, split: str) -> Dataset | None: """Try to load the prepared dataset from HuggingFace Hub.""" try: LOG.info_rank0( @@ -530,6 +507,7 @@ def try_load_from_hub( LOG.info_rank0("Unable to find prepared dataset in HuggingFace Hub") return None + def merge_datasets(datasets: list[Dataset], args: Arguments) -> Dataset: """Merge multiple datasets into one with optional shuffling. @@ -557,7 +535,7 @@ def merge_datasets(datasets: list[Dataset], args: Arguments) -> Dataset: datasets = [ds.shuffle(seed=args.train.seed) for ds in datasets] LOG.info_rank0("Merging datasets...") - #TODO: add interleave_datasets support + # TODO: add interleave_datasets support merged_dataset = concatenate_datasets(datasets) if args.data.shuffle_merged_datasets: diff --git a/src/xorl/data/prepare/utils.py b/src/xorl/data/prepare/utils.py index 05668c67..6f0f9ea6 100644 --- a/src/xorl/data/prepare/utils.py +++ b/src/xorl/data/prepare/utils.py @@ -1,6 +1,5 @@ """Data handling helpers""" -import contextlib import functools import hashlib import time @@ -8,12 +7,11 @@ from typing import Callable import huggingface_hub -import numpy as np import requests -from datasets import Dataset, IterableDataset from ...utils import logging + logger = logging.get_logger(__name__) @@ -78,4 +76,3 @@ def md5(to_hash: str, encoding: str = "utf-8") -> str: def sha256(to_hash: str, encoding: str = "utf-8") -> str: """Generate SHA256 hash of a string.""" return hashlib.sha256(to_hash.encode(encoding)).hexdigest() - diff --git a/src/xorl/distributed/gradient_accumulate_loss.py b/src/xorl/distributed/gradient_accumulate_loss.py index 377bb45a..3ab0f4dc 100644 --- a/src/xorl/distributed/gradient_accumulate_loss.py +++ b/src/xorl/distributed/gradient_accumulate_loss.py @@ -57,9 +57,7 @@ def backward( # Gradient from second output (sum loss) grad_from_sum = ( - grad_loss_sum * local_valid_tokens - if grad_loss_sum is not None - else torch.zeros_like(grad_output) + grad_loss_sum * local_valid_tokens if grad_loss_sum is not None else torch.zeros_like(grad_output) ) return grad_from_normalized + grad_from_sum, None, None @@ -71,4 +69,4 @@ def gradient_accumulate_loss( global_valid_tokens: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: """User-facing helper function to apply GradientAccumulateLoss.""" - return GradientAccumulateLoss.apply(loss, local_valid_tokens, global_valid_tokens) \ No newline at end of file + return GradientAccumulateLoss.apply(loss, local_valid_tokens, global_valid_tokens) diff --git a/src/xorl/distributed/moe/__init__.py b/src/xorl/distributed/moe/__init__.py index 97407721..17c8ccda 100644 --- a/src/xorl/distributed/moe/__init__.py +++ b/src/xorl/distributed/moe/__init__.py @@ -10,12 +10,12 @@ def __getattr__(name): "alltoall_post_combine", ): from .alltoall import ( + AllToAllDispatchContext, + alltoall_post_combine, + alltoall_pre_dispatch, preprocess, token_pre_all2all, tokens_post_all2all, - AllToAllDispatchContext, - alltoall_pre_dispatch, - alltoall_post_combine, ) globals().update( @@ -29,14 +29,21 @@ def __getattr__(name): } ) return globals()[name] - if name in ("DeepEPBuffer", "DEEPEP_AVAILABLE", "token_pre_dispatch", "tokens_post_combine", "get_default_buffer", "destroy_default_buffer"): + if name in ( + "DeepEPBuffer", + "DEEPEP_AVAILABLE", + "token_pre_dispatch", + "tokens_post_combine", + "get_default_buffer", + "destroy_default_buffer", + ): from .deepep import ( - DeepEPBuffer, DEEPEP_AVAILABLE, + DeepEPBuffer, + destroy_default_buffer, + get_default_buffer, token_pre_dispatch, tokens_post_combine, - get_default_buffer, - destroy_default_buffer, ) globals().update( diff --git a/src/xorl/distributed/moe/alltoall.py b/src/xorl/distributed/moe/alltoall.py index 8117e960..baf871d2 100644 --- a/src/xorl/distributed/moe/alltoall.py +++ b/src/xorl/distributed/moe/alltoall.py @@ -144,6 +144,7 @@ class AllToAllDispatchContext: Carries all information between ``alltoall_pre_dispatch()`` and ``alltoall_post_combine()`` so the EP compute step is stateless. """ + input_splits: List[int] output_splits: List[int] num_tokens_per_expert: torch.Tensor @@ -180,9 +181,7 @@ def alltoall_pre_dispatch( - cumsum: Cumulative sum of tokens per local expert ``[num_local_experts]``. - ctx: :class:`AllToAllDispatchContext` for ``alltoall_post_combine()``. """ - expert_mask = F.one_hot( - selected_experts, num_classes=num_experts - ).permute(2, 1, 0) + expert_mask = F.one_hot(selected_experts, num_classes=num_experts).permute(2, 1, 0) input_splits, output_splits, num_tokens_per_expert, sum_tokens = preprocess( expert_mask=expert_mask, diff --git a/src/xorl/distributed/moe/deepep.py b/src/xorl/distributed/moe/deepep.py index fcb73333..88b9b3e9 100644 --- a/src/xorl/distributed/moe/deepep.py +++ b/src/xorl/distributed/moe/deepep.py @@ -20,7 +20,7 @@ try: import deep_ep - from deep_ep.utils import EventOverlap, EventHandle + from deep_ep.utils import EventHandle, EventOverlap DEEPEP_AVAILABLE = True except ImportError: @@ -32,10 +32,7 @@ def check_deepep_available(): if not DEEPEP_AVAILABLE: - raise ImportError( - "DeepEP is not installed. Please install it from " - "https://github.com/deepseek-ai/DeepEP" - ) + raise ImportError("DeepEP is not installed. Please install it from https://github.com/deepseek-ai/DeepEP") def get_hidden_bytes(x: torch.Tensor) -> int: @@ -209,8 +206,8 @@ def permute_for_experts( ) # Flat view of expert IDs and scores — no copy, no boolean indexing - flat_expert_ids = recv_topk_idx.reshape(-1) # [num_recv_tokens * topk] - flat_scores = recv_topk_weights.reshape(-1) # [num_recv_tokens * topk] + flat_expert_ids = recv_topk_idx.reshape(-1) # [num_recv_tokens * topk] + flat_scores = recv_topk_weights.reshape(-1) # [num_recv_tokens * topk] # Invalid entries (-1) → max int64 so they sort to the end sort_keys = torch.where( @@ -269,20 +266,18 @@ def forward( buffer.buffer.get_dispatch_layout(topk_idx_deepep, num_experts) ) previous_event = EventOverlap(EventHandle()) - recv_x, recv_topk_idx, recv_topk_weights, recv_counts, handle, event = ( - buffer.buffer.dispatch( - x=x.contiguous(), - num_tokens_per_rank=num_tokens_per_rank, - num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, - is_token_in_rank=is_token_in_rank, - num_tokens_per_expert=num_tokens_per_expert, - topk_idx=topk_idx_deepep, - topk_weights=topk_weights_f32, - config=buffer.dispatch_config, - previous_event=previous_event, - async_finish=True, - allocate_on_comm_stream=True, - ) + recv_x, recv_topk_idx, recv_topk_weights, recv_counts, handle, event = buffer.buffer.dispatch( + x=x.contiguous(), + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + is_token_in_rank=is_token_in_rank, + num_tokens_per_expert=num_tokens_per_expert, + topk_idx=topk_idx_deepep, + topk_weights=topk_weights_f32, + config=buffer.dispatch_config, + previous_event=previous_event, + async_finish=True, + allocate_on_comm_stream=True, ) event.current_stream_wait() @@ -294,7 +289,10 @@ def forward( # --- permute --- expert_input, permuted_scores, permuted_indices, num_valid = permute_for_experts( - recv_x, recv_topk_idx, recv_topk_weights, num_valid=total_valid_count, + recv_x, + recv_topk_idx, + recv_topk_weights, + num_valid=total_valid_count, ) # Save for backward @@ -329,8 +327,10 @@ def backward(ctx, grad_expert_input, grad_cumsum, grad_dispatch_ctx): # Step 1: Scatter gradient from expert order → recv order grad_recv_x = torch.zeros( - ctx.num_recv_tokens, ctx.hidden_dim, - dtype=grad_expert_input.dtype, device=grad_expert_input.device, + ctx.num_recv_tokens, + ctx.hidden_dim, + dtype=grad_expert_input.dtype, + device=grad_expert_input.device, ) grad_recv_x.index_add_(0, permuted_indices, grad_expert_input) @@ -379,12 +379,18 @@ def forward( # Step 1: Weighted scatter-add (unpermute) if expert_output.shape[0] == 0: gather_output = torch.zeros( - dispatch_ctx.num_recv_tokens, hidden_dim, dtype=dtype, device=device, + dispatch_ctx.num_recv_tokens, + hidden_dim, + dtype=dtype, + device=device, ) else: weighted = expert_output * dispatch_ctx.permuted_scores.unsqueeze(1).to(expert_output.dtype) gather_output = torch.zeros( - dispatch_ctx.num_recv_tokens, hidden_dim, dtype=expert_output.dtype, device=device, + dispatch_ctx.num_recv_tokens, + hidden_dim, + dtype=expert_output.dtype, + device=device, ) idx_2d = dispatch_ctx.permuted_indices.unsqueeze(1).expand(-1, hidden_dim) gather_output.scatter_add_(0, idx_2d, weighted) @@ -501,9 +507,7 @@ def combine_no_grad( # --------------------------------------------------------------------------- # Profiling flag # --------------------------------------------------------------------------- -_DEEPEP_PROFILE = _os.environ.get("XORL_DEEPEP_PROFILE", "0").strip().lower() not in { - "0", "false", "no", "off", "" -} +_DEEPEP_PROFILE = _os.environ.get("XORL_DEEPEP_PROFILE", "0").strip().lower() not in {"0", "false", "no", "off", ""} # --------------------------------------------------------------------------- @@ -526,11 +530,19 @@ def token_pre_dispatch( if _DEEPEP_PROFILE: return _token_pre_dispatch_profiled( - buffer, hidden_states, routing_weights, selected_experts, num_experts, + buffer, + hidden_states, + routing_weights, + selected_experts, + num_experts, ) expert_input, cumsum, ctx = _FusedDispatchAndPermute.apply( - hidden_states, selected_experts, routing_weights, buffer, num_experts, + hidden_states, + selected_experts, + routing_weights, + buffer, + num_experts, ) return expert_input, cumsum, ctx @@ -559,7 +571,11 @@ def tokens_post_combine( # Profiled versions # --------------------------------------------------------------------------- def _token_pre_dispatch_profiled( - buffer, hidden_states, routing_weights, selected_experts, num_experts, + buffer, + hidden_states, + routing_weights, + selected_experts, + num_experts, ): """Profiled version of token_pre_dispatch — enabled by XORL_DEEPEP_PROFILE=1.""" rank = dist.get_rank() if dist.is_initialized() else 0 @@ -567,7 +583,11 @@ def _token_pre_dispatch_profiled( ev[0].record() expert_input, cumsum, ctx = _FusedDispatchAndPermute.apply( - hidden_states, selected_experts, routing_weights, buffer, num_experts, + hidden_states, + selected_experts, + routing_weights, + buffer, + num_experts, ) ev[1].record() @@ -596,8 +616,7 @@ def _tokens_post_combine_profiled(buffer, expert_output, ctx, async_combine): t_total = ev[0].elapsed_time(ev[1]) if rank == 0: print( - f"[DEEPEP POST r{rank}] unpermute+combine={t_total:.1f}ms " - f"expert_out={expert_output.shape}", + f"[DEEPEP POST r{rank}] unpermute+combine={t_total:.1f}ms expert_out={expert_output.shape}", flush=True, ) return result diff --git a/src/xorl/distributed/parallel_plan.py b/src/xorl/distributed/parallel_plan.py index b72bc3e9..596d6f97 100644 --- a/src/xorl/distributed/parallel_plan.py +++ b/src/xorl/distributed/parallel_plan.py @@ -112,7 +112,7 @@ def get_fsdp_no_shard_info(self, model: nn.Module): def update_prefix(self, prefix: str): """ - Update ep_plan when model is wrappered. + Update ep_plan when model is wrapped. """ self.ep_plan = {prefix + "." + k: v for k, v in self.ep_plan.items()} self.ep_param_suffix = {k.split(".")[-1] for k in self.ep_plan.keys()} diff --git a/src/xorl/distributed/parallel_state.py b/src/xorl/distributed/parallel_state.py index 733141fd..aa07e4b8 100644 --- a/src/xorl/distributed/parallel_state.py +++ b/src/xorl/distributed/parallel_state.py @@ -83,7 +83,9 @@ def _resolve_mesh_name(self, name: str) -> str: def __post_init__(self): if self.cp_fsdp_mode not in ("all", "ulysses_only", "ring_only", "none"): - raise ValueError(f"Invalid cp_fsdp_mode: {self.cp_fsdp_mode}. Must be 'all', 'ulysses_only', 'ring_only', or 'none'.") + raise ValueError( + f"Invalid cp_fsdp_mode: {self.cp_fsdp_mode}. Must be 'all', 'ulysses_only', 'ring_only', or 'none'." + ) if self.pp_size * self.dp_size * self.ringattn_size * self.ulysses_size * self.tp_size != self.world_size: raise ValueError("The product of parallel sizes should be equal to the world size.") @@ -96,8 +98,8 @@ def __post_init__(self): if self.cp_enabled: from ..distributed.sequence_parallel import ( init_sequence_parallel, - set_ringattn_group, set_data_parallel_group, + set_ringattn_group, set_ulysses_sequence_parallel_group, set_unified_sequence_parallel_group, ) diff --git a/src/xorl/distributed/pipeline_parallel.py b/src/xorl/distributed/pipeline_parallel.py index 518d6585..ea96cd11 100644 --- a/src/xorl/distributed/pipeline_parallel.py +++ b/src/xorl/distributed/pipeline_parallel.py @@ -35,6 +35,7 @@ from ..utils import logging + logger = logging.get_logger(__name__) __all__ = [ @@ -103,8 +104,7 @@ def generate_llm_fqn_per_model_part( if num_stages > num_effective_layers: raise ValueError( - f"Number of stages ({num_stages}) cannot be greater than " - f"effective layers ({num_effective_layers})" + f"Number of stages ({num_stages}) cannot be greater than effective layers ({num_effective_layers})" ) # Calculate layers per stage (distribute evenly) @@ -121,13 +121,9 @@ def generate_llm_fqn_per_model_part( ) if input_weight > layers_per_stage: - raise ValueError( - f"input_weight ({input_weight}) exceeds minimum layers per stage ({layers_per_stage})." - ) + raise ValueError(f"input_weight ({input_weight}) exceeds minimum layers per stage ({layers_per_stage}).") if output_weight > layers_per_stage: - raise ValueError( - f"output_weight ({output_weight}) exceeds minimum layers per stage ({layers_per_stage})." - ) + raise ValueError(f"output_weight ({output_weight}) exceeds minimum layers per stage ({layers_per_stage}).") module_names_per_stage = [] current_layer = 0 @@ -191,8 +187,9 @@ def _pp_forward(self, x): When the queue is not set (e.g. unit tests), falls back to generating sequential position_ids scaled by cp_size for RoPE cache correctness. """ + from xorl.models.layers.moe.routing_replay import get_replay_stage, is_r3_mode, set_replay_stage + from .parallel_state import get_parallel_state - from xorl.models.layers.moe.routing_replay import get_replay_stage, set_replay_stage, is_r3_mode ps = get_parallel_state() @@ -231,8 +228,7 @@ def _pp_forward(self, x): position_ids = metadata.pop("position_ids", None) if position_ids is not None: position_ids = position_ids.to(x.device) - extra_kwargs = {k: v.to(x.device) if isinstance(v, torch.Tensor) else v - for k, v in metadata.items()} + extra_kwargs = {k: v.to(x.device) if isinstance(v, torch.Tensor) else v for k, v in metadata.items()} # Fallback: generate sequential position_ids covering the full SP range # so that RoPE embeddings have a large enough cache. @@ -244,15 +240,19 @@ def _pp_forward(self, x): if self._pp_is_first: # x is input_ids outputs = self._pp_original_forward( - input_ids=x, position_ids=position_ids, - use_cache=False, output_hidden_states=False, + input_ids=x, + position_ids=position_ids, + use_cache=False, + output_hidden_states=False, **extra_kwargs, ) else: # x is hidden_states from previous stage outputs = self._pp_original_forward( - inputs_embeds=x, position_ids=position_ids, - use_cache=False, output_hidden_states=False, + inputs_embeds=x, + position_ids=position_ids, + use_cache=False, + output_hidden_states=False, **extra_kwargs, ) @@ -282,7 +282,7 @@ def _recursive_prune(module: nn.Module, prefix: str, fqns_to_keep: set): # Same as torchtitan: filter children of container modules by FQN children_prefix = fqn + "." layers_to_keep = { - fqn_to_keep[len(children_prefix):].split(".")[0] + fqn_to_keep[len(children_prefix) :].split(".")[0] for fqn_to_keep in fqns_to_keep if fqn_to_keep.startswith(children_prefix) } @@ -292,9 +292,7 @@ def _recursive_prune(module: nn.Module, prefix: str, fqns_to_keep: set): if layer_name not in layers_to_keep: del child[layer_name] elif isinstance(child, nn.ModuleList): - indices_to_keep = { - int(idx) for idx in layers_to_keep if idx.isdigit() - } + indices_to_keep = {int(idx) for idx in layers_to_keep if idx.isdigit()} # Preserve original indices so checkpoint keys match. # Set unwanted entries to None (named_parameters/modules skip None). for i in range(len(child)): @@ -436,8 +434,7 @@ def build_pipeline_schedule( """ if schedule_name.lower() not in _SUPPORTED_PP_SCHEDULES: raise ValueError( - f"Unsupported PP schedule '{schedule_name}'. " - f"Supported schedules: {sorted(_SUPPORTED_PP_SCHEDULES)}" + f"Unsupported PP schedule '{schedule_name}'. Supported schedules: {sorted(_SUPPORTED_PP_SCHEDULES)}" ) schedule_class = get_schedule_class(schedule_name) diff --git a/src/xorl/distributed/sequence_parallel/__init__.py b/src/xorl/distributed/sequence_parallel/__init__.py index c886e398..7688782e 100644 --- a/src/xorl/distributed/sequence_parallel/__init__.py +++ b/src/xorl/distributed/sequence_parallel/__init__.py @@ -5,11 +5,11 @@ divide_qkv_linear_weight, ) from .comm import ( + get_data_parallel_group, + get_data_parallel_rank, get_ringattn_group, get_ringattn_rank, get_ringattn_world_size, - get_data_parallel_group, - get_data_parallel_rank, get_ulysses_sequence_parallel_group, get_ulysses_sequence_parallel_rank, get_ulysses_sequence_parallel_world_size, @@ -17,8 +17,8 @@ get_unified_sequence_parallel_rank, get_unified_sequence_parallel_world_size, init_sequence_parallel, - set_ringattn_group, set_data_parallel_group, + set_ringattn_group, set_ulysses_sequence_parallel_group, set_unified_sequence_parallel_group, ) @@ -29,20 +29,20 @@ slice_input_tensor_scale_grad, slice_position_embedding, ) -from .ulysses import ( - all_to_all_images, - gather_heads_scatter_seq, - gather_seq_scatter_heads, -) from .strategy import ( + CPStrategy, HybridUlyssesRingStrategy, NoopStrategy, RingAttentionStrategy, - CPStrategy, UlyssesAsyncStrategy, UlyssesSyncStrategy, get_cp_strategy, ) +from .ulysses import ( + all_to_all_images, + gather_heads_scatter_seq, + gather_seq_scatter_heads, +) from .utils import pad_tensor, unpad_tensor, vlm_images_a2a_meta diff --git a/src/xorl/distributed/sequence_parallel/comm.py b/src/xorl/distributed/sequence_parallel/comm.py index b3b05be2..cea1c3c0 100644 --- a/src/xorl/distributed/sequence_parallel/comm.py +++ b/src/xorl/distributed/sequence_parallel/comm.py @@ -256,11 +256,9 @@ def init_sequence_parallel( assert (ulysses_group_key == "default" and _ULYSSES_SEQUENCE_PARALLEL_GROUP[ulysses_group_key] is None) or ( ulysses_group_key != "default" and ulysses_group_key not in _ULYSSES_SEQUENCE_PARALLEL_GROUP ), f"Ulysses sequence parallel group ({ulysses_group_key}) has already been initialized!" - assert ( - ulysses_group_key == "default" and _ULYSSES_SEQUENCE_PARALLEL_CPU_GROUP[ulysses_group_key] is None - ) or (ulysses_group_key != "default" and ulysses_group_key not in _ULYSSES_SEQUENCE_PARALLEL_CPU_GROUP), ( - f"Ulysses sequence parallel ({ulysses_group_key}) group has already been initialized!" - ) + assert (ulysses_group_key == "default" and _ULYSSES_SEQUENCE_PARALLEL_CPU_GROUP[ulysses_group_key] is None) or ( + ulysses_group_key != "default" and ulysses_group_key not in _ULYSSES_SEQUENCE_PARALLEL_CPU_GROUP + ), f"Ulysses sequence parallel ({ulysses_group_key}) group has already been initialized!" for i in range(data_parallel_size): # build ulysses group diff --git a/src/xorl/distributed/sequence_parallel/ring_attention.py b/src/xorl/distributed/sequence_parallel/ring_attention.py index 7ece20c4..adceb0ce 100644 --- a/src/xorl/distributed/sequence_parallel/ring_attention.py +++ b/src/xorl/distributed/sequence_parallel/ring_attention.py @@ -16,9 +16,8 @@ import torch import torch.distributed as dist import torch.nn.functional as F -from torch import Tensor - from flash_attn_interface import _flash_attn_backward, _flash_attn_forward +from torch import Tensor # --------------------------------------------------------------------------- @@ -26,7 +25,6 @@ # --------------------------------------------------------------------------- - def _merge_attn_outputs( out_acc: Tensor, lse_acc: Tensor, @@ -186,9 +184,7 @@ def _zigzag_min_chunk_id(rank: int, ringattn_size: int) -> int: return min(rank, 2 * ringattn_size - 1 - rank) -def _get_zigzag_step_section( - ringattn_rank: int, ringattn_size: int, step: int -) -> str: +def _get_zigzag_step_section(ringattn_rank: int, ringattn_size: int, step: int) -> str: """Determine the section type for a load-balanced ring step. After zigzag, each rank has chunks at positions [min_r, max_r]. @@ -203,9 +199,9 @@ def _get_zigzag_step_section( min_q = _zigzag_min_chunk_id(ringattn_rank, ringattn_size) min_kv = _zigzag_min_chunk_id(source, ringattn_size) if min_q > min_kv: - return "lower" # Q's early chunk is later → all Q attends to KV first half + return "lower" # Q's early chunk is later → all Q attends to KV first half else: - return "upper" # Q's early chunk is earlier → Q second half attends to all KV + return "upper" # Q's early chunk is earlier → Q second half attends to all KV def _compute_half_indices(cu_seqlens: Tensor) -> Tuple[Tensor, Tensor]: @@ -238,25 +234,54 @@ def _compute_half_indices(cu_seqlens: Tensor) -> Tuple[Tensor, Tensor]: # --------------------------------------------------------------------------- -def _flash_fwd(q, k, v, softmax_scale, dropout_p, causal, is_varlen, - cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=None, max_seqlen_k=None): +def _flash_fwd( + q, + k, + v, + softmax_scale, + dropout_p, + causal, + is_varlen, + cu_seqlens_q=None, + cu_seqlens_k=None, + max_seqlen_q=None, + max_seqlen_k=None, +): """Flash attention forward (FA3). Handles both batched and varlen via cu_seqlens kwargs.""" kwargs = dict(softmax_scale=softmax_scale, causal=causal) if is_varlen: - kwargs.update(cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, - max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k) + kwargs.update( + cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k + ) out, lse, _, _ = _flash_attn_forward(q, k, v, **kwargs) return out, lse -def _flash_bwd(grad_out, q, k, v, out, lse, dq, dk, dv, softmax_scale, dropout_p, - causal, is_varlen, cu_seqlens_q=None, cu_seqlens_k=None, - max_seqlen_q=None, max_seqlen_k=None): +def _flash_bwd( + grad_out, + q, + k, + v, + out, + lse, + dq, + dk, + dv, + softmax_scale, + dropout_p, + causal, + is_varlen, + cu_seqlens_q=None, + cu_seqlens_k=None, + max_seqlen_q=None, + max_seqlen_k=None, +): """Flash attention backward (FA3). Handles both batched and varlen via cu_seqlens kwargs.""" kwargs = dict(softmax_scale=softmax_scale, is_causal=causal, dq=dq, dk=dk, dv=dv) if is_varlen: - kwargs.update(cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, - max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k) + kwargs.update( + cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k + ) _flash_attn_backward(grad_out, q, k, v, out, lse, **kwargs) @@ -371,40 +396,74 @@ def forward( k_half = k_step.index_select(0, early_idx).contiguous() v_half = v_step.index_select(0, early_idx).contiguous() out_step, lse_step = _flash_fwd( - q, k_half, v_half, softmax_scale, dropout_p, - False, True, - cu_seqlens_q, cu_half_k, max_seqlen_q, max_half_k, + q, + k_half, + v_half, + softmax_scale, + dropout_p, + False, + True, + cu_seqlens_q, + cu_half_k, + max_seqlen_q, + max_half_k, ) else: half = k_step.shape[seq_dim] // 2 k_half = k_step.narrow(seq_dim, 0, half).contiguous() v_half = v_step.narrow(seq_dim, 0, half).contiguous() out_step, lse_step = _flash_fwd( - q, k_half, v_half, softmax_scale, dropout_p, - False, False, + q, + k_half, + v_half, + softmax_scale, + dropout_p, + False, + False, ) elif section == "upper": # Late half of each doc in Q attends to all KV (no mask) if is_varlen: q_half = q.index_select(0, late_idx).contiguous() out_step, lse_step = _flash_fwd( - q_half, k_step, v_step, softmax_scale, dropout_p, - False, True, - cu_half_q, cu_seqlens_k, max_half_q, max_seqlen_k, + q_half, + k_step, + v_step, + softmax_scale, + dropout_p, + False, + True, + cu_half_q, + cu_seqlens_k, + max_half_q, + max_seqlen_k, ) else: half = q.shape[seq_dim] // 2 q_half = q.narrow(seq_dim, half, half).contiguous() out_step, lse_step = _flash_fwd( - q_half, k_step, v_step, softmax_scale, dropout_p, - False, False, + q_half, + k_step, + v_step, + softmax_scale, + dropout_p, + False, + False, ) else: # Diagonal: causal on full sequence out_step, lse_step = _flash_fwd( - q, k_step, v_step, softmax_scale, dropout_p, - True, is_varlen, - cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, + q, + k_step, + v_step, + softmax_scale, + dropout_p, + True, + is_varlen, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, ) computed_steps.append((i, section)) @@ -412,9 +471,17 @@ def forward( # Non-causal or ringattn_size==1: no zigzag needed step_causal = False out_step, lse_step = _flash_fwd( - q, k_step, v_step, softmax_scale, dropout_p, - step_causal, is_varlen, - cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, + q, + k_step, + v_step, + softmax_scale, + dropout_p, + step_causal, + is_varlen, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, ) computed_steps.append((i, "full")) @@ -426,15 +493,19 @@ def forward( if is_varlen: out[late_idx] = out_step lse = torch.full( - (q.shape[1], q.shape[0]), float("-inf"), - dtype=torch.float32, device=q.device, + (q.shape[1], q.shape[0]), + float("-inf"), + dtype=torch.float32, + device=q.device, ) lse[:, late_idx] = lse_step else: out.narrow(seq_dim, half, half).copy_(out_step) lse = torch.full( - (q.shape[0], q.shape[2], q.shape[1]), float("-inf"), - dtype=torch.float32, device=q.device, + (q.shape[0], q.shape[2], q.shape[1]), + float("-inf"), + dtype=torch.float32, + device=q.device, ) lse[:, :, half:] = lse_step else: @@ -446,7 +517,11 @@ def forward( out_late = out[late_idx] lse_late = lse[:, late_idx] merged_out, merged_lse = _merge_attn_outputs( - out_late, lse_late, out_step, lse_step, True, + out_late, + lse_late, + out_step, + lse_step, + True, ) out[late_idx] = merged_out.to(out.dtype) lse[:, late_idx] = merged_lse @@ -455,13 +530,21 @@ def forward( out_second = out.narrow(seq_dim, half, half) lse_second = lse[:, :, half:] merged_out, merged_lse = _merge_attn_outputs( - out_second, lse_second, out_step, lse_step, False, + out_second, + lse_second, + out_step, + lse_step, + False, ) out.narrow(seq_dim, half, half).copy_(merged_out) lse[:, :, half:] = merged_lse else: out, lse = _merge_attn_outputs( - out, lse, out_step, lse_step, is_varlen, + out, + lse, + out_step, + lse_step, + is_varlen, ) # Cast output back to input dtype @@ -473,13 +556,17 @@ def forward( out = torch.zeros_like(q) if is_varlen: lse = torch.full( - (q.shape[1], q.shape[0]), float("-inf"), - dtype=torch.float32, device=q.device, + (q.shape[1], q.shape[0]), + float("-inf"), + dtype=torch.float32, + device=q.device, ) else: lse = torch.full( - (q.shape[0], q.shape[2], q.shape[1]), float("-inf"), - dtype=torch.float32, device=q.device, + (q.shape[0], q.shape[2], q.shape[1]), + float("-inf"), + dtype=torch.float32, + device=q.device, ) # Save for backward @@ -562,10 +649,23 @@ def backward(ctx, grad_output: Tensor): dk_step = torch.empty_like(k_half) dv_step = torch.empty_like(v_half) _flash_bwd( - grad_output, q, k_half, v_half, - out, lse, dq_step, dk_step, dv_step, - ctx.softmax_scale, ctx.dropout_p, False, True, - cu_seqlens_q, cu_half_k, ctx.max_seqlen_q, max_half_k, + grad_output, + q, + k_half, + v_half, + out, + lse, + dq_step, + dk_step, + dv_step, + ctx.softmax_scale, + ctx.dropout_p, + False, + True, + cu_seqlens_q, + cu_half_k, + ctx.max_seqlen_q, + max_half_k, ) dq.add_(dq_step) # Scatter dk/dv back to early positions @@ -579,9 +679,19 @@ def backward(ctx, grad_output: Tensor): dk_step = torch.empty_like(k_half) dv_step = torch.empty_like(v_half) _flash_bwd( - grad_output, q, k_half, v_half, - out, lse, dq_step, dk_step, dv_step, - ctx.softmax_scale, ctx.dropout_p, False, False, + grad_output, + q, + k_half, + v_half, + out, + lse, + dq_step, + dk_step, + dv_step, + ctx.softmax_scale, + ctx.dropout_p, + False, + False, ) dq.add_(dq_step) dk_all[kv_source].narrow(seq_dim, 0, half).add_(dk_step) @@ -598,10 +708,23 @@ def backward(ctx, grad_output: Tensor): dk_step = torch.empty_like(k_step) dv_step = torch.empty_like(v_step) _flash_bwd( - grad_half, q_half, k_step, v_step, - out_half, lse_half, dq_step, dk_step, dv_step, - ctx.softmax_scale, ctx.dropout_p, False, True, - cu_half_q, cu_seqlens_k, max_half_q, ctx.max_seqlen_k, + grad_half, + q_half, + k_step, + v_step, + out_half, + lse_half, + dq_step, + dk_step, + dv_step, + ctx.softmax_scale, + ctx.dropout_p, + False, + True, + cu_half_q, + cu_seqlens_k, + max_half_q, + ctx.max_seqlen_k, ) dq[late_idx] += dq_step dk_all[kv_source].add_(dk_step) @@ -616,9 +739,19 @@ def backward(ctx, grad_output: Tensor): dk_step = torch.empty_like(k_step) dv_step = torch.empty_like(v_step) _flash_bwd( - grad_half, q_half, k_step, v_step, - out_half, lse_half, dq_step, dk_step, dv_step, - ctx.softmax_scale, ctx.dropout_p, False, False, + grad_half, + q_half, + k_step, + v_step, + out_half, + lse_half, + dq_step, + dk_step, + dv_step, + ctx.softmax_scale, + ctx.dropout_p, + False, + False, ) dq.narrow(seq_dim, half, half).add_(dq_step) dk_all[kv_source].add_(dk_step) @@ -630,10 +763,23 @@ def backward(ctx, grad_output: Tensor): dk_step = torch.empty_like(k_step) dv_step = torch.empty_like(v_step) _flash_bwd( - grad_output, q, k_step, v_step, - out, lse, dq_step, dk_step, dv_step, - ctx.softmax_scale, ctx.dropout_p, True, is_varlen, - cu_seqlens_q, cu_seqlens_k, ctx.max_seqlen_q, ctx.max_seqlen_k, + grad_output, + q, + k_step, + v_step, + out, + lse, + dq_step, + dk_step, + dv_step, + ctx.softmax_scale, + ctx.dropout_p, + True, + is_varlen, + cu_seqlens_q, + cu_seqlens_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, ) dq.add_(dq_step) dk_all[kv_source].add_(dk_step) @@ -644,10 +790,23 @@ def backward(ctx, grad_output: Tensor): dk_step = torch.empty_like(k_step) dv_step = torch.empty_like(v_step) _flash_bwd( - grad_output, q, k_step, v_step, - out, lse, dq_step, dk_step, dv_step, - ctx.softmax_scale, ctx.dropout_p, step_causal, is_varlen, - cu_seqlens_q, cu_seqlens_k, ctx.max_seqlen_q, ctx.max_seqlen_k, + grad_output, + q, + k_step, + v_step, + out, + lse, + dq_step, + dk_step, + dv_step, + ctx.softmax_scale, + ctx.dropout_p, + step_causal, + is_varlen, + cu_seqlens_q, + cu_seqlens_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, ) dq.add_(dq_step) dk_all[kv_source].add_(dk_step) diff --git a/src/xorl/distributed/sequence_parallel/strategy.py b/src/xorl/distributed/sequence_parallel/strategy.py index 66281b7e..360e1552 100644 --- a/src/xorl/distributed/sequence_parallel/strategy.py +++ b/src/xorl/distributed/sequence_parallel/strategy.py @@ -14,7 +14,7 @@ """ from abc import ABC, abstractmethod -from typing import Optional, Tuple +from typing import Optional import torch import torch.distributed as dist @@ -62,6 +62,7 @@ def _scale_cu_seqlens_for_ringattn(kwargs, ringattn_group): # Base class # ------------------------------------------------------------------ # + class CPStrategy(ABC): """Base class for sequence parallel strategies. @@ -110,6 +111,7 @@ def prepare_position_embeddings(self, position_embeddings, dim, sp_group, num_kv # NoopStrategy — no sequence parallelism # ------------------------------------------------------------------ # + class NoopStrategy(CPStrategy): """No sequence parallelism — delegates to module's own methods.""" @@ -118,9 +120,7 @@ def project_qkv(self, module, hidden_states, position_embeddings): def compute_attention(self, module, q, k, v, attention_mask, **kwargs): attn_fn = module._get_attention_fn() - attn_output, _ = attn_fn( - module, q, k, v, attention_mask, **module._attention_kwargs(), **kwargs - ) + attn_output, _ = attn_fn(module, q, k, v, attention_mask, **module._attention_kwargs(), **kwargs) return attn_output def project_output(self, module, attn_output): @@ -134,6 +134,7 @@ def prepare_position_embeddings(self, position_embeddings, **kwargs): # Ulysses sync — all-to-all around attention # ------------------------------------------------------------------ # + class UlyssesSyncStrategy(CPStrategy): """Ulysses SP (sync): fused KV all-to-all with GQA head expansion. @@ -155,8 +156,7 @@ def project_qkv(self, module, hidden_states, position_embeddings): kv_head_num = k.shape[2] if self.ulysses_size > kv_head_num: assert self.ulysses_size % kv_head_num == 0, ( - f"ulysses_size ({self.ulysses_size}) must be divisible by " - f"num_key_value_heads ({kv_head_num})" + f"ulysses_size ({self.ulysses_size}) must be divisible by num_key_value_heads ({kv_head_num})" ) n_repeat = self.ulysses_size // kv_head_num # repeat_kv expects [batch, num_heads, seq, head_dim] @@ -190,13 +190,9 @@ def project_qkv(self, module, hidden_states, position_embeddings): q = gather_seq_scatter_heads(q, seq_dim=1, head_dim=2, group=self.group) # Fused KV a2a for 4D tensors kv = torch.stack([k, v], dim=3) # [B, S, H_kv, 2, D] - kv = kv.reshape( - k.size(0), k.size(1), 2 * k.size(2), k.size(3) - ) # [B, S, 2*H_kv, D] + kv = kv.reshape(k.size(0), k.size(1), 2 * k.size(2), k.size(3)) # [B, S, 2*H_kv, D] kv = gather_seq_scatter_heads(kv, seq_dim=1, head_dim=2, group=self.group) - kv = kv.reshape( - kv.size(0), kv.size(1), kv.size(2) // 2, 2, kv.size(3) - ) # [B, S_full, H_kv/SP, 2, D] + kv = kv.reshape(kv.size(0), kv.size(1), kv.size(2) // 2, 2, kv.size(3)) # [B, S_full, H_kv/SP, 2, D] k = kv[:, :, :, 0, :].contiguous() v = kv[:, :, :, 1, :].contiguous() @@ -204,23 +200,17 @@ def project_qkv(self, module, hidden_states, position_embeddings): def compute_attention(self, module, q, k, v, attention_mask, **kwargs): attn_fn = module._get_attention_fn() - attn_output, _ = attn_fn( - module, q, k, v, attention_mask, **module._attention_kwargs(), **kwargs - ) + attn_output, _ = attn_fn(module, q, k, v, attention_mask, **module._attention_kwargs(), **kwargs) return attn_output def project_output(self, module, attn_output): # Post-attention a2a: gather heads, scatter seq if attn_output.ndim == 4 and attn_output.size(0) == 1: attn_output = attn_output.squeeze(0) - attn_output = gather_heads_scatter_seq( - attn_output, seq_dim=0, head_dim=1, group=self.group - ) + attn_output = gather_heads_scatter_seq(attn_output, seq_dim=0, head_dim=1, group=self.group) attn_output = attn_output.unsqueeze(0) else: - attn_output = gather_heads_scatter_seq( - attn_output, seq_dim=1, head_dim=2, group=self.group - ) + attn_output = gather_heads_scatter_seq(attn_output, seq_dim=1, head_dim=2, group=self.group) return module._project_output(attn_output) @@ -232,6 +222,7 @@ def prepare_position_embeddings(self, position_embeddings, dim, sp_group, **kwar # Ulysses async — overlapped linear + a2a communication # ------------------------------------------------------------------ # + class UlyssesAsyncStrategy(CPStrategy): """Ulysses SP (async): overlaps QKV linear projections with a2a. @@ -323,9 +314,7 @@ def project_qkv(self, module, hidden_states, position_embeddings): def compute_attention(self, module, q, k, v, attention_mask, **kwargs): attn_fn = module._get_attention_fn() - attn_output, _ = attn_fn( - module, q, k, v, attention_mask, **module._attention_kwargs(), **kwargs - ) + attn_output, _ = attn_fn(module, q, k, v, attention_mask, **module._attention_kwargs(), **kwargs) return attn_output def project_output(self, module, attn_output): @@ -334,14 +323,10 @@ def project_output(self, module, attn_output): if module.o_proj.weight is None: if attn_output.ndim == 4 and attn_output.size(0) == 1: attn_output = attn_output.squeeze(0) - attn_output = gather_heads_scatter_seq( - attn_output, seq_dim=0, head_dim=1, group=self.group - ) + attn_output = gather_heads_scatter_seq(attn_output, seq_dim=0, head_dim=1, group=self.group) attn_output = attn_output.unsqueeze(0) else: - attn_output = gather_heads_scatter_seq( - attn_output, seq_dim=1, head_dim=2, group=self.group - ) + attn_output = gather_heads_scatter_seq(attn_output, seq_dim=1, head_dim=2, group=self.group) return module._project_output(attn_output) # Async output projection: a2a + o_proj (backward overlaps a2a with weight grad) @@ -367,6 +352,7 @@ def prepare_position_embeddings(self, position_embeddings, **kwargs): # Ring attention # ------------------------------------------------------------------ # + class RingAttentionStrategy(CPStrategy): """Ring attention without Ulysses. @@ -385,8 +371,8 @@ def compute_attention(self, module, q, k, v, attention_mask, **kwargs): from .ring_attention import ring_flash_attention_forward attn_kwargs = module._attention_kwargs() - cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k = ( - _scale_cu_seqlens_for_ringattn(kwargs, self.ringattn_group) + cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k = _scale_cu_seqlens_for_ringattn( + kwargs, self.ringattn_group ) return ring_flash_attention_forward( @@ -414,6 +400,7 @@ def prepare_position_embeddings(self, position_embeddings, dim, sp_group, **kwar # Hybrid Ulysses + Ring # ------------------------------------------------------------------ # + class HybridUlyssesRingStrategy(CPStrategy): """Hybrid Ulysses + Ring attention. @@ -441,8 +428,8 @@ def compute_attention(self, module, q, k, v, attention_mask, **kwargs): from .ring_attention import ring_flash_attention_forward attn_kwargs = module._attention_kwargs() - cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k = ( - _scale_cu_seqlens_for_ringattn(kwargs, self.ringattn_group) + cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k = _scale_cu_seqlens_for_ringattn( + kwargs, self.ringattn_group ) return ring_flash_attention_forward( diff --git a/src/xorl/distributed/sequence_parallel/ulysses.py b/src/xorl/distributed/sequence_parallel/ulysses.py index 074b76c8..d1cc48a4 100644 --- a/src/xorl/distributed/sequence_parallel/ulysses.py +++ b/src/xorl/distributed/sequence_parallel/ulysses.py @@ -256,10 +256,10 @@ def gather_seq_scatter_heads_qkv( group: ProcessGroup = None, ) -> Tensor: """ - A func to sync splited qkv tensor + A func to sync split qkv tensor qkv_tensor: the tensor we want to do alltoall with. The last dim must be the projection_idx, which we will split into 3 part. After - spliting, the gather idx will be projecttion_idx + 1 + splitting, the gather idx will be projecttion_idx + 1 seq_dim: gather_dim for all2all comm restore_shape: if True, output will has the same shape length as input """ diff --git a/src/xorl/distributed/sequence_parallel/utils.py b/src/xorl/distributed/sequence_parallel/utils.py index de2c19d9..255fb9a0 100644 --- a/src/xorl/distributed/sequence_parallel/utils.py +++ b/src/xorl/distributed/sequence_parallel/utils.py @@ -78,7 +78,9 @@ def has_overlap(x1, x2, y1, y2) -> Tuple[bool, int]: return max_value < min_value, min_value - max_value -def all2all_splits(image_lens: List, image_lens_per_rank: List, ulysses_size: int, ulysses_rank: int) -> Tuple[List, List]: +def all2all_splits( + image_lens: List, image_lens_per_rank: List, ulysses_size: int, ulysses_rank: int +) -> Tuple[List, List]: """ A func to generate splits for all2all communication """ @@ -113,7 +115,7 @@ def vlm_images_a2a_meta( ulysses_rank: int, ulysses_size: int, image_lens: List, image_masks: torch.Tensor ) -> Tuple[List, List, torch.Tensor]: """ - A func to generate metadata for all2all communication after we balance the computaion in vision encoder + A func to generate metadata for all2all communication after we balance the computation in vision encoder Usually we will split the batches of images for vision encoder in sp group. However, before we feed images tokens into language model, we need to use all2all communication to gather necessary tokens into the current rank. """ diff --git a/src/xorl/distributed/sync_padding.py b/src/xorl/distributed/sync_padding.py index 9237297b..4953df13 100644 --- a/src/xorl/distributed/sync_padding.py +++ b/src/xorl/distributed/sync_padding.py @@ -1,9 +1,10 @@ """Synchronize padding across distributed ranks to prevent load imbalance.""" +from typing import Any, Dict, List, Optional + import torch import torch.distributed as dist import torch.nn.functional as F -from typing import Any, Dict, List, Optional from xorl.data.constants import IGNORE_INDEX from xorl.distributed.parallel_state import get_parallel_state @@ -69,8 +70,12 @@ def synchronize_micro_batch_padding( # Identify which paddable keys are present sample = micro_batches[0] - keys_2d = [k for k in _2D_PAD_VALUES if k in sample and isinstance(sample[k], torch.Tensor) and sample[k].dim() == 2] - keys_3d = [k for k in _3D_PAD_VALUES if k in sample and isinstance(sample[k], torch.Tensor) and sample[k].dim() == 3] + keys_2d = [ + k for k in _2D_PAD_VALUES if k in sample and isinstance(sample[k], torch.Tensor) and sample[k].dim() == 2 + ] + keys_3d = [ + k for k in _3D_PAD_VALUES if k in sample and isinstance(sample[k], torch.Tensor) and sample[k].dim() == 3 + ] if not keys_2d and not keys_3d: return diff --git a/src/xorl/distributed/torch_parallelize.py b/src/xorl/distributed/torch_parallelize.py index 7b5d6b51..b7c434db 100644 --- a/src/xorl/distributed/torch_parallelize.py +++ b/src/xorl/distributed/torch_parallelize.py @@ -231,7 +231,7 @@ def _experts_shard_placement_fn(param): for sub_mod in layer_mod.modules(): if isinstance(sub_mod, mp_ignored_classes) and sub_mod is not layer_mod: # this will also create a AllGather communication group - # when modules here are small (like gating), this would slightly impacts the peformance + # when modules here are small (like gating), this would slightly impacts the performance # a better method might be adding them to ignored_params of fully_shard # but then they will need to be initialized separately fully_shard(sub_mod, **fsdp_kwargs_without_mp) @@ -300,6 +300,7 @@ def _experts_shard_placement_fn(param): # Skip EP-sharded expert modules: they manage their own divide factor (ep_size) # for averaging expert gradients across EP ranks. from torch.distributed._composable.fsdp.fully_shard import FSDPModule + for module in model.modules(): if isinstance(module, FSDPModule) and not getattr(module, "_is_ep_fsdp", False): module.set_gradient_divide_factor(1.0) @@ -465,9 +466,13 @@ def build_parallelize_model( logger.info_rank0("TP enabled: loading weights before FSDP wrapping...") load_weights_mode = kwargs.get("load_weights_mode", "broadcast") if load_weights_mode == "broadcast": - rank0_load_and_broadcast_weights(model_part, weights_path, get_device_type(), dtensor_factory=distribute_tensor) + rank0_load_and_broadcast_weights( + model_part, weights_path, get_device_type(), dtensor_factory=distribute_tensor + ) else: - all_ranks_load_weights(model_part, weights_path, get_device_type(), dtensor_factory=distribute_tensor) + all_ranks_load_weights( + model_part, weights_path, get_device_type(), dtensor_factory=distribute_tensor + ) kwargs["skip_weight_loading"] = True # torch.compile (if enabled) @@ -575,6 +580,7 @@ def _reentrant_ckpt_with_kwargs(fn, *args, **kw): if kw: fn = functools.partial(fn, **kw) return orig_fn(fn, *args) + return _reentrant_ckpt_with_kwargs module._gradient_checkpointing_func = _make_wrapper(_orig) @@ -608,7 +614,9 @@ def _reentrant_ckpt_with_kwargs(fn, *args, **kw): logger.info_rank0("TP enabled: loading weights before FSDP wrapping...") load_weights_mode = kwargs.get("load_weights_mode", "broadcast") if load_weights_mode == "broadcast": - rank0_load_and_broadcast_weights(model, weights_path, get_device_type(), dtensor_factory=distribute_tensor) + rank0_load_and_broadcast_weights( + model, weights_path, get_device_type(), dtensor_factory=distribute_tensor + ) else: all_ranks_load_weights(model, weights_path, get_device_type(), dtensor_factory=distribute_tensor) # Mark weights as already loaded so FSDP path skips loading diff --git a/src/xorl/lora/__init__.py b/src/xorl/lora/__init__.py index 14ea01c8..32d54919 100644 --- a/src/xorl/lora/__init__.py +++ b/src/xorl/lora/__init__.py @@ -9,8 +9,6 @@ """ # Base class and implementations -from xorl.lora.modules import LoraModule, LoraLinear - # Mapping from xorl.lora.mapping import ( LORA_MAPPING, @@ -18,35 +16,36 @@ get_lora_class_for_module, register_lora_mapping, ) +from xorl.lora.modules import LoraLinear, LoraModule # Utility functions from xorl.lora.utils import ( + count_lora_parameters, + freeze_base_parameters, + get_all_lora_state_dict, + get_lora_parameters, + get_lora_state_dict, + get_moe_lora_state_dict, inject_lora_into_model, inject_lora_into_model_with_moe, inject_lora_into_moe_blocks, - get_lora_state_dict, - get_moe_lora_state_dict, - get_all_lora_state_dict, + load_lora_checkpoint, load_lora_state_dict, - freeze_base_parameters, - get_lora_parameters, save_lora_checkpoint, - load_lora_checkpoint, - count_lora_parameters, ) # MoE LoRA — deferred to avoid circular import with xorl.models.layers.moe. # Access via xorl.lora.MoEExpertsLoRA etc. triggers __getattr__ below. - # Group GEMM ops for direct access from xorl.ops.group_gemm.kernel import ( - init_lora_weights_stacked, compute_lora_scaling, + get_lora_delta_weight_stacked, + init_lora_weights_stacked, merge_lora_weights_stacked, unmerge_lora_weights_stacked, - get_lora_delta_weight_stacked, ) + __all__ = [ # Base class "LoraModule", @@ -91,20 +90,24 @@ # Lazy imports for MoE LoRA to break circular import: # moe/lora.py → xorl.lora.modules.base → xorl.lora → xorl.models.layers.moe (cycle) _MOE_LAZY_ATTRS = { - "MoELoRAConfig", "MoEExpertsLoRA", - "copy_weights_to_lora_experts", "mark_only_lora_as_trainable", "lora_state_dict", + "MoELoRAConfig", + "MoEExpertsLoRA", + "copy_weights_to_lora_experts", + "mark_only_lora_as_trainable", + "lora_state_dict", } def __getattr__(name): if name in _MOE_LAZY_ATTRS: from xorl.models.layers.moe.lora import ( - MoELoRAConfig, MoEExpertsLoRA, + MoELoRAConfig, copy_weights_to_lora_experts, - mark_only_lora_as_trainable, lora_state_dict, + mark_only_lora_as_trainable, ) + # Cache all at once to avoid repeated imports g = globals() g["MoELoRAConfig"] = MoELoRAConfig diff --git a/src/xorl/lora/mapping.py b/src/xorl/lora/mapping.py index 422e8c52..301172e8 100644 --- a/src/xorl/lora/mapping.py +++ b/src/xorl/lora/mapping.py @@ -29,6 +29,7 @@ def _ensure_moe_mapping(): if not _moe_registered: _moe_registered = True from xorl.models.layers.moe import MoEExperts, MoEExpertsLoRA + LORA_MAPPING[MoEExperts] = MoEExpertsLoRA @@ -55,7 +56,7 @@ def get_lora_class_for_module(module: nn.Module) -> Optional[Type[Union[LoraModu return None # Check if module has 'lora_config' (already a LoRA module, e.g., MoE LoRA) - if hasattr(module, 'lora_config'): + if hasattr(module, "lora_config"): return None for source_cls, lora_cls in LORA_MAPPING.items(): diff --git a/src/xorl/lora/modules/__init__.py b/src/xorl/lora/modules/__init__.py index 74904dc0..b78e2b3a 100644 --- a/src/xorl/lora/modules/__init__.py +++ b/src/xorl/lora/modules/__init__.py @@ -8,6 +8,7 @@ from .base import LoraModule from .linear import LoraLinear + __all__ = [ "LoraModule", "LoraLinear", diff --git a/src/xorl/lora/modules/base.py b/src/xorl/lora/modules/base.py index 196265f7..d666765c 100644 --- a/src/xorl/lora/modules/base.py +++ b/src/xorl/lora/modules/base.py @@ -153,8 +153,4 @@ def get_lora_state_dict(self) -> dict[str, torch.Tensor]: if not isinstance(self, nn.Module): raise TypeError("LoraModule must be mixed with nn.Module") - return { - name: param.detach() - for name, param in self.named_parameters() - if "lora_" in name - } + return {name: param.detach() for name, param in self.named_parameters() if "lora_" in name} diff --git a/src/xorl/lora/modules/linear.py b/src/xorl/lora/modules/linear.py index 3230effb..c6c4772f 100644 --- a/src/xorl/lora/modules/linear.py +++ b/src/xorl/lora/modules/linear.py @@ -73,12 +73,8 @@ def __init__( # LoRA weights (trainable, float32 for numerical stability) # lora_A: down-projection [r, in_features] # lora_B: up-projection [out_features, r] - self.lora_A = nn.Parameter( - torch.empty(r, in_features, device=device, dtype=torch.float32) - ) - self.lora_B = nn.Parameter( - torch.empty(out_features, r, device=device, dtype=torch.float32) - ) + self.lora_A = nn.Parameter(torch.empty(r, in_features, device=device, dtype=torch.float32)) + self.lora_B = nn.Parameter(torch.empty(out_features, r, device=device, dtype=torch.float32)) # Initialize LoRA parameters self.reset_lora_parameters() diff --git a/src/xorl/lora/utils.py b/src/xorl/lora/utils.py index e71c2afd..ecc1ca42 100644 --- a/src/xorl/lora/utils.py +++ b/src/xorl/lora/utils.py @@ -8,15 +8,16 @@ import logging import os import re -from typing import Dict, Iterator, List, Optional, Set, Tuple, Union +from typing import Dict, Iterator, List, Optional, Tuple import torch import torch.nn as nn -from safetensors.torch import save_file, load_file +from safetensors.torch import load_file, save_file +from xorl.lora.mapping import can_apply_lora, get_lora_class_for_module from xorl.lora.modules import LoraLinear from xorl.lora.modules.base import LoraModule -from xorl.lora.mapping import can_apply_lora, get_lora_class_for_module + logger = logging.getLogger(__name__) @@ -244,8 +245,6 @@ def get_lora_parameters(model: nn.Module) -> Iterator[nn.Parameter]: yield param - - def _gather_ep_tensor(tensor: torch.Tensor, spec_info) -> torch.Tensor: """Gather an EP-sharded tensor from all EP ranks. @@ -295,7 +294,7 @@ def get_lora_state_dict( """ from torch.distributed._tensor import DTensor - fqn2spec_info = getattr(model, '_fqn2spec_info', None) + fqn2spec_info = getattr(model, "_fqn2spec_info", None) lora_state_dict = {} for name, param in model.named_parameters(): @@ -354,9 +353,7 @@ def load_lora_state_dict( if strict and (missing_keys or unexpected_keys): raise RuntimeError( - f"Error loading LoRA state dict.\n" - f"Missing keys: {missing_keys}\n" - f"Unexpected keys: {unexpected_keys}" + f"Error loading LoRA state dict.\nMissing keys: {missing_keys}\nUnexpected keys: {unexpected_keys}" ) # Load matching keys @@ -405,7 +402,6 @@ def save_lora_checkpoint( Returns: Path to saved checkpoint directory """ - import re os.makedirs(save_path, exist_ok=True) @@ -465,7 +461,7 @@ def _unmerge_moe_lora_weights(name: str, stacked_tensor: torch.Tensor) -> Dict[s result = {} # Check if this is a shared weight (hybrid shared LoRA) - is_shared = (num_experts == 1) + is_shared = num_experts == 1 if is_shared: # Shared weight: use ".shared." in the key name @@ -612,9 +608,9 @@ def _convert_from_peft_key(name: str) -> str: return name.replace("lm_head.lora_embedding_B", "lm_head.lora_B") # Handle .weight suffix for Linear layers elif name.endswith(".lora_A.weight"): - return name[:-len(".weight")] + return name[: -len(".weight")] elif name.endswith(".lora_B.weight"): - return name[:-len(".weight")] + return name[: -len(".weight")] return name # Convert PEFT format keys to xorl format if needed @@ -622,7 +618,7 @@ def _convert_from_peft_key(name: str) -> str: for key, value in state_dict.items(): # Remove PEFT prefix: base_model.model.{key} -> {key} if key.startswith("base_model.model."): - new_key = key[len("base_model.model."):] + new_key = key[len("base_model.model.") :] else: new_key = key # Convert from PEFT naming to xorl naming @@ -713,7 +709,7 @@ def inject_lora_into_moe_blocks( for name, module in model.named_modules(): # Check if this is a MoE block that supports LoRA injection - if hasattr(module, 'inject_lora') and hasattr(module, 'lora_adapter'): + if hasattr(module, "inject_lora") and hasattr(module, "lora_adapter"): # This is a MoEBlock or similar module.inject_lora( r=r, @@ -732,8 +728,7 @@ def inject_lora_into_moe_blocks( ) else: logger.warning( - "No LoRA-compatible MoE blocks found. Make sure model uses " - "moe_implementation='triton' or 'native'" + "No LoRA-compatible MoE blocks found. Make sure model uses moe_implementation='triton' or 'native'" ) return injected_count @@ -830,10 +825,10 @@ def get_moe_lora_state_dict(model: nn.Module) -> Dict[str, torch.Tensor]: for name, module in model.named_modules(): # Check if this is an MoE LoRA expert module (has lora_config attribute) - if hasattr(module, 'lora_config') and module.lora_config is not None: + if hasattr(module, "lora_config") and module.lora_config is not None: # Get all lora_ parameters from this module for param_name, param in module.named_parameters(): - if 'lora_' in param_name: + if "lora_" in param_name: full_key = f"{name}.{param_name}" moe_lora_state_dict[full_key] = param.detach().cpu() diff --git a/src/xorl/models/__init__.py b/src/xorl/models/__init__.py index 827fac6f..9262632d 100644 --- a/src/xorl/models/__init__.py +++ b/src/xorl/models/__init__.py @@ -1,8 +1,8 @@ from . import transformers from .auto import build_foundation_model, build_processor, build_tokenizer from .module_utils import ( - init_empty_weights, all_ranks_load_weights, + init_empty_weights, rank0_load_and_broadcast_weights, save_model_assets, save_model_weights, @@ -19,5 +19,4 @@ "save_model_assets", "save_model_weights", "transformers", - ] diff --git a/src/xorl/models/base.py b/src/xorl/models/base.py index d03e11a2..bb266ea3 100644 --- a/src/xorl/models/base.py +++ b/src/xorl/models/base.py @@ -46,7 +46,6 @@ def post_init(self): """Initialize weights and apply weight tying.""" self.apply(self._init_weights) - def init_weights(self, buffer_device=None): """Initialize weights, matching torchtitan's ModelProtocol. @@ -97,11 +96,9 @@ def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): recompute_modules = gradient_checkpointing_kwargs.pop("recompute_modules", None) moe_checkpoint_method = gradient_checkpointing_kwargs.pop("moe_checkpoint_method", None) - gc_func = partial( - torch.utils.checkpoint.checkpoint, **gradient_checkpointing_kwargs - ) + gc_func = partial(torch.utils.checkpoint.checkpoint, **gradient_checkpointing_kwargs) - moe_act = (moe_checkpoint_method == "moe_act") + moe_act = moe_checkpoint_method == "moe_act" for module in self.modules(): if hasattr(module, "gradient_checkpointing"): module.gradient_checkpointing = True diff --git a/src/xorl/models/checkpoint_handlers/__init__.py b/src/xorl/models/checkpoint_handlers/__init__.py index dc44149f..9c5dda02 100644 --- a/src/xorl/models/checkpoint_handlers/__init__.py +++ b/src/xorl/models/checkpoint_handlers/__init__.py @@ -3,13 +3,14 @@ ExpertWeightBuffer, GateUpMergeBuffer, QKVMergeBuffer, - parse_expert_key, checkpoint_has_per_expert_weights, - model_needs_expert_merging, - model_has_gate_up_merged, checkpoint_has_separate_gate_up, + model_has_gate_up_merged, + model_needs_expert_merging, + parse_expert_key, ) + __all__ = [ "CheckpointHandler", "ExpertWeightBuffer", diff --git a/src/xorl/models/checkpoint_handlers/base.py b/src/xorl/models/checkpoint_handlers/base.py index b7a6bd42..52ded506 100644 --- a/src/xorl/models/checkpoint_handlers/base.py +++ b/src/xorl/models/checkpoint_handlers/base.py @@ -27,9 +27,7 @@ class CheckpointHandler: add to output state_dict """ - def on_load_weight( - self, key: str, tensor: torch.Tensor - ) -> List[Tuple[str, torch.Tensor]]: + def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: """Process one checkpoint key/tensor during loading. Returns: @@ -40,9 +38,7 @@ def on_load_weight( """ return [(key, tensor)] - def on_skip_weight( - self, key: str - ) -> List[Tuple[str, torch.Tensor]]: + def on_skip_weight(self, key: str) -> List[Tuple[str, torch.Tensor]]: """Notify the handler that a checkpoint key was skipped (not loaded from disk). Used by EP-aware filtered loading: out-of-range expert weight tensors @@ -77,9 +73,7 @@ def get_skip_key_fn(self) -> Optional[Callable[[str], bool]]: """ return None - def on_save_weight( - self, param_name: str, tensor: torch.Tensor - ) -> List[Tuple[str, torch.Tensor]]: + def on_save_weight(self, param_name: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: """Process one model parameter during saving. Returns: diff --git a/src/xorl/models/checkpoint_handlers/buffers.py b/src/xorl/models/checkpoint_handlers/buffers.py index 5639b7d8..76771c7b 100644 --- a/src/xorl/models/checkpoint_handlers/buffers.py +++ b/src/xorl/models/checkpoint_handlers/buffers.py @@ -20,9 +20,7 @@ # Per-expert HuggingFace weight keys # e.g., model.layers.0.mlp.experts.5.gate_proj.weight -EXPERT_KEY_PATTERN = re.compile( - r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.weight$" -) +EXPERT_KEY_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.weight$") # Per-expert quantized auxiliary keys (NVFP4 / modelopt format) # e.g., model.layers.0.mlp.experts.5.gate_proj.weight_scale @@ -34,45 +32,31 @@ ) # Generic expert key: matches any suffix (weight, weight_scale, weight_scale_inv, etc.) -_EXPERT_ANY_KEY_PATTERN = re.compile( - r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.(.+)$" -) +_EXPERT_ANY_KEY_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.(.+)$") # Model parameter names in fused expert format # e.g., model.layers.0.mlp.experts.gate_proj -FUSED_EXPERT_PATTERN = re.compile( - r"^model\.layers\.\d+\.mlp\.experts\.(gate|up|down)_proj$" -) +FUSED_EXPERT_PATTERN = re.compile(r"^model\.layers\.\d+\.mlp\.experts\.(gate|up|down)_proj$") # Dense MLP gate/up projection weight keys # e.g., model.layers.0.mlp.gate_proj.weight or model.layers.0.mlp.shared_expert.up_proj.weight -DENSE_GATE_UP_PATTERN = re.compile( - r"^(.*)\.(gate|up)_proj\.weight$" -) +DENSE_GATE_UP_PATTERN = re.compile(r"^(.*)\.(gate|up)_proj\.weight$") # Attention o_proj weight key # e.g., model.layers.0.self_attn.o_proj.weight -OPROJ_WEIGHT_PATTERN = re.compile( - r"^.*\.self_attn\.o_proj\.weight$" -) +OPROJ_WEIGHT_PATTERN = re.compile(r"^.*\.self_attn\.o_proj\.weight$") # Dense/shared-expert down_proj weight key (non-expert) # e.g., model.layers.0.mlp.down_proj.weight or model.layers.0.mlp.shared_expert.down_proj.weight -DENSE_DOWN_PROJ_PATTERN = re.compile( - r"^.*\.mlp\.(?:shared_expert\.)?down_proj\.weight$" -) +DENSE_DOWN_PROJ_PATTERN = re.compile(r"^.*\.mlp\.(?:shared_expert\.)?down_proj\.weight$") # Attention QKV projection weight/bias keys # e.g., model.layers.0.self_attn.q_proj.weight -QKV_PROJ_PATTERN = re.compile( - r"^(.*\.self_attn)\.(q|k|v)_proj\.(weight|bias)$" -) +QKV_PROJ_PATTERN = re.compile(r"^(.*\.self_attn)\.(q|k|v)_proj\.(weight|bias)$") # Quantized auxiliary suffixes used by modelopt NVFP4 checkpoints. # Matches any key ending in .weight_scale, .weight_scale_2, or .input_scale -QUANT_AUX_SUFFIX_PATTERN = re.compile( - r"\.(weight_scale|weight_scale_2|input_scale)$" -) +QUANT_AUX_SUFFIX_PATTERN = re.compile(r"\.(weight_scale|weight_scale_2|input_scale)$") # Block FP8 auxiliary suffix (HuggingFace FP8 checkpoint format) # Matches any key ending in .weight_scale_inv @@ -83,6 +67,7 @@ # Key parsing and detection helpers # ============================================================================= + def parse_expert_key(key: str) -> Optional[Tuple[int, int, str]]: """Parse a per-expert weight key to extract layer index, expert index, and projection name. @@ -136,6 +121,7 @@ def _resolve_weights_path(weights_path: Optional[str]) -> Optional[str]: # Try resolving as a HF hub model ID try: from transformers.utils import cached_file + config_path = cached_file(weights_path, "config.json", _raise_exceptions_for_missing_entries=False) if config_path and os.path.isfile(config_path): return os.path.dirname(config_path) @@ -169,10 +155,7 @@ def detect_prequantized_checkpoint(weights_path: Optional[str]) -> bool: try: with open(quant_config_path) as f: quant_config = json.load(f) - algo = ( - quant_config.get("quantization", {}).get("quant_algo") - or quant_config.get("quant_algo") - ) + algo = quant_config.get("quantization", {}).get("quant_algo") or quant_config.get("quant_algo") if algo == "NVFP4": return True except (json.JSONDecodeError, OSError): @@ -230,8 +213,7 @@ def detect_prequantized_block_fp8_checkpoint(weights_path: Optional[str]) -> boo with open(config_path) as f: config = json.load(f) qc = config.get("quantization_config", {}) - if (qc.get("quant_method") == "fp8" - and qc.get("weight_block_size") == [128, 128]): + if qc.get("quant_method") == "fp8" and qc.get("weight_block_size") == [128, 128]: return True except (json.JSONDecodeError, OSError): pass @@ -287,10 +269,7 @@ def get_prequantized_exclude_modules(weights_path: Optional[str]) -> Set[str]: try: with open(quant_config_path) as f: qc = json.load(f) - exclude = ( - qc.get("quantization", {}).get("exclude_modules") - or qc.get("exclude_modules") - ) + exclude = qc.get("quantization", {}).get("exclude_modules") or qc.get("exclude_modules") if exclude: return set(exclude) except (json.JSONDecodeError, OSError): @@ -303,11 +282,7 @@ def get_prequantized_exclude_modules(weights_path: Optional[str]) -> Set[str]: with open(config_path) as f: config = json.load(f) qc = config.get("quantization_config", {}) - exclude = ( - qc.get("exclude_modules") - or qc.get("modules_to_not_convert") - or qc.get("ignore") - ) + exclude = qc.get("exclude_modules") or qc.get("modules_to_not_convert") or qc.get("ignore") if exclude: # Normalize full FQNs (e.g. "model.layers.0.mlp.gate") # to short suffixes (e.g. "gate", "lm_head") @@ -344,6 +319,7 @@ def checkpoint_has_separate_gate_up(checkpoint_keys: Set[str]) -> bool: # ExpertWeightBuffer # ============================================================================= + class ExpertWeightBuffer: """Buffer for collecting per-expert weights and merging them into stacked (G,K,N) tensors. @@ -361,7 +337,10 @@ class ExpertWeightBuffer: """ def __init__( - self, num_experts: int, ep_rank: int = 0, ep_size: int = 1, + self, + num_experts: int, + ep_rank: int = 0, + ep_size: int = 1, device: Optional[torch.device] = None, ): self.num_experts = num_experts @@ -398,18 +377,14 @@ def add(self, layer_idx: int, expert_idx: int, proj: str, tensor: torch.Tensor) if key not in self._stacked_buffers: stacked_shape = (self.local_num_experts,) + tensor.shape - self._stacked_buffers[key] = torch.empty( - stacked_shape, dtype=tensor.dtype, device=self._device - ) + self._stacked_buffers[key] = torch.empty(stacked_shape, dtype=tensor.dtype, device=self._device) else: # CPU path: transpose on CPU tensor = tensor.t().contiguous() if key not in self._stacked_buffers: stacked_shape = (self.local_num_experts,) + tensor.shape - self._stacked_buffers[key] = torch.empty( - stacked_shape, dtype=tensor.dtype, device="cpu" - ) + self._stacked_buffers[key] = torch.empty(stacked_shape, dtype=tensor.dtype, device="cpu") local_idx = expert_idx - self.expert_start self._stacked_buffers[key][local_idx].copy_(tensor) @@ -464,6 +439,7 @@ def get_pending_counts(self) -> Dict[Tuple[int, str], int]: # GateUpMergeBuffer # ============================================================================= + class GateUpMergeBuffer: """Buffer for merging separate gate_proj and up_proj checkpoint weights into gate_up_proj. @@ -510,6 +486,7 @@ def get_pending(self) -> Dict[str, List[str]]: # QKVMergeBuffer # ============================================================================= + class QKVMergeBuffer: """Buffer for merging separate q_proj, k_proj, v_proj checkpoint weights into qkv_proj. @@ -529,9 +506,9 @@ def add(self, key: str, tensor: torch.Tensor) -> Optional[Tuple[str, torch.Tenso if match is None: return None - prefix = match.group(1) # e.g., "model.layers.0.self_attn" - proj_type = match.group(2) # "q", "k", or "v" - param_type = match.group(3) # "weight" or "bias" + prefix = match.group(1) # e.g., "model.layers.0.self_attn" + proj_type = match.group(2) # "q", "k", or "v" + param_type = match.group(3) # "weight" or "bias" buf_key = (prefix, param_type) if buf_key not in self._pending: @@ -566,6 +543,7 @@ def get_pending(self) -> Dict[Tuple[str, str], List[str]]: class _SourceInfo: """Mapping from a checkpoint key prefix to a target module.""" + __slots__ = ("module_fqn", "src_proj") def __init__(self, module_fqn: str, src_proj: Optional[str]): @@ -575,10 +553,17 @@ def __init__(self, module_fqn: str, src_proj: Optional[str]): class _ModuleAccum: """Accumulator for collecting quantized weight pieces for one QLoRALinear module.""" + __slots__ = ( - "module", "target_format", "target_group_size", - "expected_pieces", "merge_sources", "received", "received_count", - "ema_amax", "scale_dtypes", + "module", + "target_format", + "target_group_size", + "expected_pieces", + "merge_sources", + "received", + "received_count", + "ema_amax", + "scale_dtypes", ) def __init__( @@ -632,11 +617,13 @@ def __init__(self, model: "torch.nn.Module"): for src_proj in merge_sources: prefix = f"{module._source_fqn}.{src_proj}" self._prefix_map[prefix] = _SourceInfo( - module_fqn=fqn, src_proj=src_proj, + module_fqn=fqn, + src_proj=src_proj, ) else: self._prefix_map[module._source_fqn] = _SourceInfo( - module_fqn=fqn, src_proj=None, + module_fqn=fqn, + src_proj=None, ) self._accums[fqn] = _ModuleAccum( @@ -649,9 +636,7 @@ def __init__(self, model: "torch.nn.Module"): # ----- public API (shared) ----- - def try_consume( - self, key: str, tensor: "torch.Tensor" - ) -> Optional[List[Tuple[str, "torch.Tensor"]]]: + def try_consume(self, key: str, tensor: "torch.Tensor") -> Optional[List[Tuple[str, "torch.Tensor"]]]: """Try to consume a checkpoint key as a QLoRA quantized weight. Returns: @@ -662,7 +647,7 @@ def try_consume( dot_pos = key.rfind(".") if dot_pos < 0: return None - prefix, suffix = key[:dot_pos], key[dot_pos + 1:] + prefix, suffix = key[:dot_pos], key[dot_pos + 1 :] if prefix not in self._prefix_map: return None @@ -689,7 +674,7 @@ def is_qlora_key(self, key: str) -> bool: dot_pos = key.rfind(".") if dot_pos < 0: return False - prefix, suffix = key[:dot_pos], key[dot_pos + 1:] + prefix, suffix = key[:dot_pos], key[dot_pos + 1 :] return prefix in self._prefix_map and suffix in self._suffixes def set_inline_metadata(self) -> None: @@ -723,9 +708,7 @@ def set_inline_metadata(self) -> None: # Source format string ("nvfp4" or "block_fp8"). _source_format: str - def _pack( - self, fqn: str, accum: _ModuleAccum - ) -> List[Tuple[str, "torch.Tensor"]]: + def _pack(self, fqn: str, accum: _ModuleAccum) -> List[Tuple[str, "torch.Tensor"]]: raise NotImplementedError # ----- shared helpers ----- @@ -747,8 +730,10 @@ def _convert_format( ) -> List[Tuple[str, "torch.Tensor"]]: """Dequantize from source format, re-quantize in target format.""" from xorl.ops.quantize import ( - nvfp4_quantize, nvfp4_dequantize, - block_fp8_quantize_gkn, block_fp8_dequantize_gkn, + block_fp8_dequantize_gkn, + block_fp8_quantize_gkn, + nvfp4_dequantize, + nvfp4_quantize, ) from xorl.qlora.modules.linear import QLoRALinear @@ -791,7 +776,9 @@ def _convert_format( else: ema_amax = w.float().abs().max().reshape(1) packed, block_scales, global_scale = nvfp4_quantize( - w, target_gs, global_amax=ema_amax, + w, + target_gs, + global_amax=ema_amax, ) new_packed = QLoRALinear._to_uint8(packed).contiguous().view(torch.float32) new_packed = new_packed.reshape(module.packed_weight_f32.shape) @@ -817,9 +804,7 @@ class NvFP4WeightBuffer(_QLoRAWeightBufferBase): _suffixes = _QLORA_SUFFIXES_NVFP4 _source_format = "nvfp4" - def _pack( - self, fqn: str, accum: _ModuleAccum - ) -> List[Tuple[str, "torch.Tensor"]]: + def _pack(self, fqn: str, accum: _ModuleAccum) -> List[Tuple[str, "torch.Tensor"]]: from xorl.qlora.modules.linear import QLoRALinear module = accum.module @@ -872,9 +857,7 @@ class BlockFP8WeightBuffer(_QLoRAWeightBufferBase): _suffixes = _QLORA_SUFFIXES_FP8 _source_format = "block_fp8" - def _pack( - self, fqn: str, accum: _ModuleAccum - ) -> List[Tuple[str, "torch.Tensor"]]: + def _pack(self, fqn: str, accum: _ModuleAccum) -> List[Tuple[str, "torch.Tensor"]]: from xorl.qlora.modules.linear import QLoRALinear module = accum.module @@ -990,9 +973,7 @@ def expert_end(self) -> int: def has_modules(self) -> bool: return len(self._layer_modules) > 0 - def try_consume( - self, key: str, tensor: torch.Tensor - ) -> Optional[List[Tuple[str, torch.Tensor]]]: + def try_consume(self, key: str, tensor: torch.Tensor) -> Optional[List[Tuple[str, torch.Tensor]]]: """Try to consume an expert quantized key. Returns None if not an expert key. @@ -1079,6 +1060,7 @@ def get_pending(self) -> List[Tuple[int, str]]: def set_inline_metadata(self) -> None: """Finalize inline-loaded modules: materialize LoRA params from meta device.""" import math + from torch.distributed._tensor import DTensor for module in self._layer_modules.values(): @@ -1093,12 +1075,12 @@ def set_inline_metadata(self) -> None: placement = param.placements mesh = param.device_mesh local_data = torch.zeros( - local_shape, dtype=param.dtype, device="cuda", + local_shape, + dtype=param.dtype, + device="cuda", ) materialized = torch.nn.Parameter( - DTensor.from_local( - local_data, mesh, placement, run_check=False - ), + DTensor.from_local(local_data, mesh, placement, run_check=False), requires_grad=param.requires_grad, ) else: @@ -1109,14 +1091,10 @@ def set_inline_metadata(self) -> None: parts = name.split("_") if parts[-1] == "A": for i in range(materialized.shape[0]): - torch.nn.init.kaiming_uniform_( - materialized.data[i], a=math.sqrt(5) - ) + torch.nn.init.kaiming_uniform_(materialized.data[i], a=math.sqrt(5)) setattr(module, name, materialized) - def _process_expert( - self, layer: int, proj: str, local_idx: int, pieces: Dict[str, torch.Tensor] - ) -> None: + def _process_expert(self, layer: int, proj: str, local_idx: int, pieces: Dict[str, torch.Tensor]) -> None: """Process one expert's data (transpose, scale absorption). Subclass implements.""" raise NotImplementedError @@ -1135,7 +1113,7 @@ def _process_expert(self, layer, proj, local_idx, pieces): from xorl.qlora.modules.moe_experts import QLoRAMoeExperts lp = (layer, proj) - packed = pieces["weight"] # [N, K//2] uint8 + packed = pieces["weight"] # [N, K//2] uint8 block_scales = pieces["weight_scale"] # [N, K//bs] fp8 global_scale = pieces["weight_scale_2"] # [1] fp32 @@ -1168,18 +1146,14 @@ def _finalize_proj(self, layer, proj): setattr(module, f"{proj}_block_scales", stacked_scales) # Global scale absorbed into block_scales — store 1.0 per expert - gs = QLoRAMoeExperts._to_uint8( - torch.ones(self._local_num, 1, dtype=torch.float32) - ).to(device) + gs = QLoRAMoeExperts._to_uint8(torch.ones(self._local_num, 1, dtype=torch.float32)).to(device) setattr(module, f"{proj}_global_scale", gs) module._scale_dtypes[proj] = { "weight_block_scales": torch.float32, "weight_global_scale": torch.float32, } - module._ema_amax[proj] = torch.tensor( - amax_list, dtype=torch.float32, device=device - ) + module._ema_amax[proj] = torch.tensor(amax_list, dtype=torch.float32, device=device) class BlockFP8QLoRAExpertBuffer(_QLoRAExpertBufferBase): @@ -1191,8 +1165,8 @@ def _process_expert(self, layer, proj, local_idx, pieces): from xorl.qlora.modules.moe_experts import QLoRAMoeExperts lp = (layer, proj) - fp8_w = pieces["weight"] # [N, K] float8_e4m3fn - scales = pieces["weight_scale_inv"] # [N//128, K//128] f32 + fp8_w = pieces["weight"] # [N, K] float8_e4m3fn + scales = pieces["weight_scale_inv"] # [N//128, K//128] f32 # Transpose HF [N, K] -> GKN [K, N] fp8_w_gkn = fp8_w.T.contiguous() @@ -1206,8 +1180,6 @@ def _process_expert(self, layer, proj, local_idx, pieces): self._scales_lists[lp][local_idx] = QLoRAMoeExperts._to_uint8(scales_gkn) def _finalize_proj(self, layer, proj): - from xorl.qlora.modules.moe_experts import QLoRAMoeExperts - lp = (layer, proj) module = self._layer_modules[layer] device = torch.device("cuda") @@ -1229,16 +1201,13 @@ def QLoRAExpertBuffer( ) -> Optional[_QLoRAExpertBufferBase]: """Factory: create format-specific QLoRA expert buffer, or None if no MoE experts.""" from xorl.qlora.modules.moe_experts import ( - QLoRAMoeExperts, BlockFP8QLoRAMoeExperts, + BlockFP8QLoRAMoeExperts, + QLoRAMoeExperts, ) for _, m in model.named_modules(): if isinstance(m, QLoRAMoeExperts) and not m._weights_loaded: - cls = ( - BlockFP8QLoRAExpertBuffer - if isinstance(m, BlockFP8QLoRAMoeExperts) - else NvFP4QLoRAExpertBuffer - ) + cls = BlockFP8QLoRAExpertBuffer if isinstance(m, BlockFP8QLoRAMoeExperts) else NvFP4QLoRAExpertBuffer buf = cls(model, ep_rank, ep_size, num_experts) return buf if buf.has_modules() else None diff --git a/src/xorl/models/layers/__init__.py b/src/xorl/models/layers/__init__.py index 455d4517..1304bf36 100644 --- a/src/xorl/models/layers/__init__.py +++ b/src/xorl/models/layers/__init__.py @@ -1,8 +1,8 @@ from .activations import ACT2FN from .attention import ( ATTENTION_FUNCTIONS, - AttentionKwargs, FLASH_ATTENTION_IMPLEMENTATIONS, + AttentionKwargs, FlashAttentionKwargs, eager_attention_forward, is_flash_attention, @@ -27,6 +27,7 @@ rotate_half, ) + __all__ = [ "ACT2FN", "ATTENTION_FUNCTIONS", diff --git a/src/xorl/models/layers/activations.py b/src/xorl/models/layers/activations.py index 462a7da1..42a95ed5 100644 --- a/src/xorl/models/layers/activations.py +++ b/src/xorl/models/layers/activations.py @@ -225,9 +225,7 @@ def __init__( ): super().__init__() self.alpha_p = nn.Parameter(torch.log(torch.expm1(torch.tensor(alpha_p_init, dtype=dtype))).unsqueeze(0)) - self.alpha_n = nn.Parameter( - torch.log(torch.expm1(torch.tensor(alpha_n_init - beta, dtype=dtype))).unsqueeze(0) - ) + self.alpha_n = nn.Parameter(torch.log(torch.expm1(torch.tensor(alpha_n_init - beta, dtype=dtype))).unsqueeze(0)) self.register_buffer("beta", torch.tensor(beta, dtype=dtype)) self.register_buffer("eps", torch.tensor(eps, dtype=dtype)) self.with_vector_loads = with_vector_loads diff --git a/src/xorl/models/layers/attention/__init__.py b/src/xorl/models/layers/attention/__init__.py index a6c61c93..89163870 100644 --- a/src/xorl/models/layers/attention/__init__.py +++ b/src/xorl/models/layers/attention/__init__.py @@ -1,16 +1,17 @@ -from .multi_head_attention import MultiHeadAttention -from .utils import repeat_kv from .backend import ( ATTENTION_FUNCTIONS, - AttentionKwargs, CAUSAL_MASK_FUNCTIONS, FLASH_ATTENTION_IMPLEMENTATIONS, + AttentionKwargs, FlashAttentionKwargs, is_flash_attention, update_causal_mask, ) -from .backend.eager import eager_attention_forward from .backend._mask_utils import prepare_4d_causal_attention_mask_with_cache_position +from .backend.eager import eager_attention_forward +from .multi_head_attention import MultiHeadAttention +from .utils import repeat_kv + # Conditional import — flash may not be installed try: diff --git a/src/xorl/models/layers/attention/backend/__init__.py b/src/xorl/models/layers/attention/backend/__init__.py index 8254dcec..67875664 100644 --- a/src/xorl/models/layers/attention/backend/__init__.py +++ b/src/xorl/models/layers/attention/backend/__init__.py @@ -15,14 +15,17 @@ import torch from typing_extensions import TypedDict -from .eager import eager_attention_forward, prepare_causal_mask as eager_prepare_causal_mask -from .native import native_attention_forward, prepare_causal_mask as native_prepare_causal_mask +from .eager import eager_attention_forward +from .eager import prepare_causal_mask as eager_prepare_causal_mask +from .native import native_attention_forward +from .native import prepare_causal_mask as native_prepare_causal_mask # ------------------------------------------------------------------ # # Generic attention kwargs (renamed from FlashAttentionKwargs) # ------------------------------------------------------------------ # + class AttentionKwargs(TypedDict, total=False): """Kwargs for attention functions (packed/varlen sequences). @@ -108,6 +111,7 @@ def is_flash_attention(attn_implementation: str) -> bool: # Thin dispatcher # ------------------------------------------------------------------ # + def update_causal_mask( attn_implementation: str, attention_mask: Union[torch.Tensor, None], diff --git a/src/xorl/models/layers/attention/backend/_mask_utils.py b/src/xorl/models/layers/attention/backend/_mask_utils.py index 76e08db7..95578f24 100644 --- a/src/xorl/models/layers/attention/backend/_mask_utils.py +++ b/src/xorl/models/layers/attention/backend/_mask_utils.py @@ -24,9 +24,7 @@ def prepare_4d_causal_attention_mask_with_cache_position( else: device = cache_position.device min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device - ) + causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device) diagonal_attend_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) if sliding_window is not None: @@ -46,7 +44,5 @@ def prepare_4d_causal_attention_mask_with_cache_position( mask_length = attention_mask.shape[-1] padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(device) padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(padding_mask, min_dtype) return causal_mask diff --git a/src/xorl/models/layers/attention/backend/eager.py b/src/xorl/models/layers/attention/backend/eager.py index f05c3d9a..63e97ef6 100644 --- a/src/xorl/models/layers/attention/backend/eager.py +++ b/src/xorl/models/layers/attention/backend/eager.py @@ -59,7 +59,9 @@ def eager_attention_forward( elif k_len % mk == 0: causal_mask = causal_mask.repeat_interleave(k_len // mk, dim=-1) else: - causal_mask = torch.nn.functional.pad(causal_mask, (0, k_len - mk), value=torch.finfo(causal_mask.dtype).min) + causal_mask = torch.nn.functional.pad( + causal_mask, (0, k_len - mk), value=torch.finfo(causal_mask.dtype).min + ) # Match query axis mq = causal_mask.shape[-2] @@ -92,9 +94,7 @@ def prepare_causal_mask( past_seen_tokens = 0 sequence_length = input_tensor.shape[1] target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 + attention_mask.shape[-1] if isinstance(attention_mask, torch.Tensor) else past_seen_tokens + sequence_length + 1 ) return prepare_4d_causal_attention_mask_with_cache_position( diff --git a/src/xorl/models/layers/attention/backend/flex_attention.py b/src/xorl/models/layers/attention/backend/flex_attention.py index 699d6c6f..08d17986 100644 --- a/src/xorl/models/layers/attention/backend/flex_attention.py +++ b/src/xorl/models/layers/attention/backend/flex_attention.py @@ -4,6 +4,7 @@ import torch + try: from torch.nn.attention.flex_attention import BlockMask, create_block_mask diff --git a/src/xorl/models/layers/attention/backend/native.py b/src/xorl/models/layers/attention/backend/native.py index ab5c92e9..d0db8217 100644 --- a/src/xorl/models/layers/attention/backend/native.py +++ b/src/xorl/models/layers/attention/backend/native.py @@ -48,12 +48,18 @@ def native_attention_forward( # so SDPA applies causal masking per-document within the packed batch. max_length_q = kwargs.get("max_length_q", query.shape[2]) mask = _build_varlen_causal_mask( - cu_seq_lens_q, cu_seq_lens_k, query.shape[2], key.shape[2], - device=query.device, dtype=query.dtype, + cu_seq_lens_q, + cu_seq_lens_k, + query.shape[2], + key.shape[2], + device=query.device, + dtype=query.dtype, ) with sdpa_kernel(SDPBackend.CUDNN_ATTENTION): attn_output = torch.nn.functional.scaled_dot_product_attention( - query, key, value, + query, + key, + value, attn_mask=mask, dropout_p=dropout if module.training else 0.0, scale=scaling, @@ -62,7 +68,9 @@ def native_attention_forward( else: with sdpa_kernel(SDPBackend.CUDNN_ATTENTION): attn_output = torch.nn.functional.scaled_dot_product_attention( - query, key, value, + query, + key, + value, is_causal=causal, dropout_p=dropout if module.training else 0.0, scale=scaling, diff --git a/src/xorl/models/layers/attention/multi_head_attention.py b/src/xorl/models/layers/attention/multi_head_attention.py index 0fd7a608..765b3fa7 100644 --- a/src/xorl/models/layers/attention/multi_head_attention.py +++ b/src/xorl/models/layers/attention/multi_head_attention.py @@ -5,11 +5,12 @@ import torch from torch import nn -from xorl.models.layers.normalization import RMSNorm -from xorl.models.layers.rope import apply_rotary_pos_emb +from xorl.distributed.sequence_parallel.strategy import get_cp_strategy from xorl.models.layers.attention.backend import ATTENTION_FUNCTIONS, AttentionKwargs from xorl.models.layers.attention.backend.eager import eager_attention_forward -from xorl.distributed.sequence_parallel.strategy import get_cp_strategy +from xorl.models.layers.normalization import RMSNorm +from xorl.models.layers.rope import apply_rotary_pos_emb + class MultiHeadAttention(nn.Module): """Base multi-head attention shared across all decoder model variants. @@ -34,9 +35,7 @@ def __init__(self, config, layer_idx: int): self.q_dim = config.num_attention_heads * self.head_dim self.kv_dim = config.num_key_value_heads * self.head_dim - self.qkv_proj = nn.Linear( - config.hidden_size, self.q_dim + 2 * self.kv_dim, bias=config.attention_bias - ) + self.qkv_proj = nn.Linear(config.hidden_size, self.q_dim + 2 * self.kv_dim, bias=config.attention_bias) self.o_proj = nn.Linear( config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias ) @@ -56,9 +55,15 @@ def unfuse_for_tp(self): """Replace fused qkv_proj with separate q_proj, k_proj, v_proj for tensor parallelism.""" device = self.qkv_proj.weight.device dtype = self.qkv_proj.weight.dtype - self.q_proj = nn.Linear(self.config.hidden_size, self.q_dim, bias=self.config.attention_bias, device=device, dtype=dtype) - self.k_proj = nn.Linear(self.config.hidden_size, self.kv_dim, bias=self.config.attention_bias, device=device, dtype=dtype) - self.v_proj = nn.Linear(self.config.hidden_size, self.kv_dim, bias=self.config.attention_bias, device=device, dtype=dtype) + self.q_proj = nn.Linear( + self.config.hidden_size, self.q_dim, bias=self.config.attention_bias, device=device, dtype=dtype + ) + self.k_proj = nn.Linear( + self.config.hidden_size, self.kv_dim, bias=self.config.attention_bias, device=device, dtype=dtype + ) + self.v_proj = nn.Linear( + self.config.hidden_size, self.kv_dim, bias=self.config.attention_bias, device=device, dtype=dtype + ) del self.qkv_proj def _project_qkv( @@ -91,7 +96,7 @@ def _project_qkv( q, k = apply_rotary_pos_emb(q, k, cos, sin) # Optionally cast to bfloat16 after RoPE for SGLang numerical alignment - if getattr(self.config, '_attention_cast_bf16', False): + if getattr(self.config, "_attention_cast_bf16", False): q = q.to(torch.bfloat16) k = k.to(torch.bfloat16) @@ -132,8 +137,6 @@ def forward( attention_mask: Optional[torch.Tensor], **kwargs: Unpack[AttentionKwargs], ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - - attn_strategy = get_cp_strategy(num_kv_heads=self.config.num_key_value_heads) # Phase 1: QKV projection + norm + RoPE (+ pre-attention SP communication) diff --git a/src/xorl/models/layers/attention/utils.py b/src/xorl/models/layers/attention/utils.py index 76451ad3..8e45675b 100644 --- a/src/xorl/models/layers/attention/utils.py +++ b/src/xorl/models/layers/attention/utils.py @@ -13,7 +13,5 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand( - batch, num_key_value_heads, n_rep, slen, head_dim - ) + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) diff --git a/src/xorl/models/layers/moe/backend/__init__.py b/src/xorl/models/layers/moe/backend/__init__.py index d51c4c6b..30f2c3c2 100644 --- a/src/xorl/models/layers/moe/backend/__init__.py +++ b/src/xorl/models/layers/moe/backend/__init__.py @@ -7,11 +7,13 @@ from typing import Callable, Dict + MOE_EXPERT_BACKENDS: Dict[str, Callable] = {} # Eager is always available (no kernel deps) from .eager import eager_expert_forward + MOE_EXPERT_BACKENDS["eager"] = eager_expert_forward # Triton group GEMM kernels (custom Triton) @@ -83,7 +85,7 @@ # AllToAll dispatch (always available when EP ops are available) try: - from xorl.distributed.moe import alltoall_pre_dispatch, alltoall_post_combine + from xorl.distributed.moe import alltoall_post_combine, alltoall_pre_dispatch EP_DISPATCH["alltoall"] = alltoall_pre_dispatch EP_COMBINE["alltoall"] = alltoall_post_combine @@ -94,6 +96,8 @@ try: from xorl.distributed.moe.deepep import ( token_pre_dispatch as deepep_pre_dispatch, + ) + from xorl.distributed.moe.deepep import ( tokens_post_combine as deepep_post_combine, ) diff --git a/src/xorl/models/layers/moe/backend/eager.py b/src/xorl/models/layers/moe/backend/eager.py index 67243a55..6cf0a0cc 100644 --- a/src/xorl/models/layers/moe/backend/eager.py +++ b/src/xorl/models/layers/moe/backend/eager.py @@ -28,9 +28,7 @@ def eager_expert_forward( Returns: Output tensor of shape ``(num_tokens, hidden_dim)``. """ - assert not get_parallel_state().ep_enabled, ( - "_moe_implementation='eager' does not support EP" - ) + assert not get_parallel_state().ep_enabled, "_moe_implementation='eager' does not support EP" gate_proj_out = torch.matmul(hidden_states, gate_proj[expert_idx]) up_proj_out = torch.matmul(hidden_states, up_proj[expert_idx]) out = act_fn(gate_proj_out) * up_proj_out diff --git a/src/xorl/models/layers/moe/backend/native.py b/src/xorl/models/layers/moe/backend/native.py index f20c865a..664dadd8 100644 --- a/src/xorl/models/layers/moe/backend/native.py +++ b/src/xorl/models/layers/moe/backend/native.py @@ -17,6 +17,7 @@ import torch import torch.nn.functional as F + # Alignment for token groups: 8 for bf16/fp16. _TOKEN_GROUP_ALIGN = 8 @@ -30,6 +31,7 @@ # Padding helpers (vectorized — no Python for-loops, no .tolist()) # --------------------------------------------------------------------------- + @torch.compiler.disable def _compute_pad_indices( counts: torch.Tensor, @@ -85,9 +87,9 @@ def _pad_to_alignment( ``(x_padded, padded_counts)`` where ``padded_counts`` are aligned. """ padded_counts = torch.clamp_min(counts, _TOKEN_GROUP_ALIGN) - padded_counts = ( - (padded_counts + _TOKEN_GROUP_ALIGN - 1) // _TOKEN_GROUP_ALIGN * _TOKEN_GROUP_ALIGN - ).to(torch.int32) + padded_counts = ((padded_counts + _TOKEN_GROUP_ALIGN - 1) // _TOKEN_GROUP_ALIGN * _TOKEN_GROUP_ALIGN).to( + torch.int32 + ) total_sorted = x_sorted.shape[0] total_padded = int(padded_counts.sum().item()) @@ -118,6 +120,7 @@ def _unpad( # Core grouped-mm expert forward (compiled like torchtitan) # --------------------------------------------------------------------------- + def _run_experts_grouped_mm( gate_proj: torch.Tensor, up_proj: torch.Tensor, @@ -189,9 +192,7 @@ def _gate_up_swiglu( compute_dtype = torch.bfloat16 x_bf16 = x.to(compute_dtype) - gate_out = F.silu( - torch._grouped_mm(x_bf16, gate_proj.to(compute_dtype), offs=offsets) - ) + gate_out = F.silu(torch._grouped_mm(x_bf16, gate_proj.to(compute_dtype), offs=offsets)) up_out = torch._grouped_mm(x_bf16, up_proj.to(compute_dtype), offs=offsets) return gate_out * up_out @@ -223,12 +224,18 @@ def _run_experts_moe_act( # instead of saving gate_out + up_out tensors. The compiled function is # called twice (forward + recompute in backward) — this is intentional. h = torch.utils.checkpoint.checkpoint( - _gate_up_swiglu_compiled, x, gate_proj, up_proj, offsets, + _gate_up_swiglu_compiled, + x, + gate_proj, + up_proj, + offsets, use_reentrant=False, ) out = torch._grouped_mm( - h, down_proj.to(torch.bfloat16), offs=offsets, + h, + down_proj.to(torch.bfloat16), + offs=offsets, ).to(x.dtype) return out @@ -238,13 +245,12 @@ def _run_experts_moe_act( # Shared token preparation / scatter helper # --------------------------------------------------------------------------- + def _ensure_dynamo_configured(): global _native_dynamo_configured if not _native_dynamo_configured: torch._dynamo.config.capture_scalar_outputs = True - torch._dynamo.config.cache_size_limit = max( - torch._dynamo.config.cache_size_limit, 64 - ) + torch._dynamo.config.cache_size_limit = max(torch._dynamo.config.cache_size_limit, 64) _native_dynamo_configured = True @@ -270,32 +276,33 @@ def _native_expert_forward_impl( device = hidden_states.device # 1. Flatten top-k assignments - flat_experts = selected_experts.view(-1) # (N * top_k,) - flat_weights = routing_weights.view(-1) # (N * top_k,) + flat_experts = selected_experts.view(-1) # (N * top_k,) + flat_weights = routing_weights.view(-1) # (N * top_k,) # 2. Sort by expert (stable preserves token order within each expert) sorted_order = torch.argsort(flat_experts, stable=True) - token_ids = sorted_order // top_k # original token index + token_ids = sorted_order // top_k # original token index - sorted_hidden = hidden_states[token_ids] # (N * top_k, dim) - sorted_weights = flat_weights[sorted_order] # (N * top_k,) + sorted_hidden = hidden_states[token_ids] # (N * top_k, dim) + sorted_weights = flat_weights[sorted_order] # (N * top_k,) # 3. Histogram — how many tokens per expert num_tokens_per_expert = torch.histc( - flat_experts.float(), bins=num_experts, min=0, max=num_experts - 1, + flat_experts.float(), + bins=num_experts, + min=0, + max=num_experts - 1, ).to(torch.int32) # 4. Compute padded counts for grouped_mm alignment padded_counts = torch.clamp_min(num_tokens_per_expert, _TOKEN_GROUP_ALIGN) - padded_counts = ( - (padded_counts + _TOKEN_GROUP_ALIGN - 1) // _TOKEN_GROUP_ALIGN * _TOKEN_GROUP_ALIGN - ).to(torch.int32) + padded_counts = ((padded_counts + _TOKEN_GROUP_ALIGN - 1) // _TOKEN_GROUP_ALIGN * _TOKEN_GROUP_ALIGN).to( + torch.int32 + ) # 5. Compute pad index mapping (once — reused for pad + unpad) total_sorted = num_tokens * top_k - pad_dst = _compute_pad_indices( - num_tokens_per_expert, padded_counts, num_experts, total_sorted, device - ) + pad_dst = _compute_pad_indices(num_tokens_per_expert, padded_counts, num_experts, total_sorted, device) # 6. Scatter sorted tokens into padded layout total_padded = int(padded_counts.sum().item()) @@ -303,9 +310,7 @@ def _native_expert_forward_impl( sorted_hidden_padded[pad_dst] = sorted_hidden # 7. Expert compute (backend-specific) - expert_out_padded = compute_fn( - gate_proj, up_proj, down_proj, sorted_hidden_padded, padded_counts - ) + expert_out_padded = compute_fn(gate_proj, up_proj, down_proj, sorted_hidden_padded, padded_counts) # 8. Gather from padded layout (reuse pad_dst) + apply routing weights expert_out = expert_out_padded[pad_dst] * sorted_weights.unsqueeze(-1) @@ -321,6 +326,7 @@ def _native_expert_forward_impl( # Public entry point (same interface as triton/quack backends) # --------------------------------------------------------------------------- + def native_expert_forward( hidden_states: torch.Tensor, routing_weights: torch.Tensor, @@ -333,8 +339,13 @@ def native_expert_forward( ) -> torch.Tensor: """Forward pass using native PyTorch ``torch._grouped_mm``.""" return _native_expert_forward_impl( - hidden_states, routing_weights, selected_experts, - gate_proj, up_proj, down_proj, num_experts, + hidden_states, + routing_weights, + selected_experts, + gate_proj, + up_proj, + down_proj, + num_experts, _run_experts_compiled, ) @@ -353,9 +364,7 @@ def expand_shared_lora_weight(tensor: torch.Tensor, num_experts: int) -> torch.T return tensor -def align_lora_rank_for_grouped_mm( - lora_A: torch.Tensor, lora_B: torch.Tensor -) -> tuple: +def align_lora_rank_for_grouped_mm(lora_A: torch.Tensor, lora_B: torch.Tensor) -> tuple: """Pad LoRA rank dimension so ``torch._grouped_mm`` stride is 16-byte aligned. With (G, K, N) format: @@ -387,6 +396,7 @@ def _cumsum_to_counts(cumsum: torch.Tensor, num_experts: int) -> torch.Tensor: # Core grouped-mm with LoRA (NOT compiled — LoRA expand/align not compile-friendly) # --------------------------------------------------------------------------- + def _run_grouped_mm_with_lora( x: torch.Tensor, padded_counts: torch.Tensor, @@ -417,7 +427,8 @@ def _run_grouped_mm_with_lora( gate_lora_A, gate_lora_B = align_lora_rank_for_grouped_mm(gate_lora_A, gate_lora_B) gate_lora_out = torch._grouped_mm( torch._grouped_mm(x_bf16, gate_lora_A, offs=offsets), - gate_lora_B, offs=offsets, + gate_lora_B, + offs=offsets, ) gate_out = F.silu(gate_out + gate_lora_out * scaling) @@ -428,7 +439,8 @@ def _run_grouped_mm_with_lora( up_lora_A, up_lora_B = align_lora_rank_for_grouped_mm(up_lora_A, up_lora_B) up_lora_out = torch._grouped_mm( torch._grouped_mm(x_bf16, up_lora_A, offs=offsets), - up_lora_B, offs=offsets, + up_lora_B, + offs=offsets, ) up_out = up_out + up_lora_out * scaling @@ -442,7 +454,8 @@ def _run_grouped_mm_with_lora( down_lora_A, down_lora_B = align_lora_rank_for_grouped_mm(down_lora_A, down_lora_B) down_lora_out = torch._grouped_mm( torch._grouped_mm(h, down_lora_A, offs=offsets), - down_lora_B, offs=offsets, + down_lora_B, + offs=offsets, ) out = out + down_lora_out * scaling @@ -453,6 +466,7 @@ def _run_grouped_mm_with_lora( # Public entry point: local native LoRA forward # --------------------------------------------------------------------------- + def native_expert_lora_forward( hidden_states: torch.Tensor, routing_weights: torch.Tensor, @@ -491,29 +505,34 @@ def native_expert_lora_forward( # 3. Histogram num_tokens_per_expert = torch.histc( - flat_experts.float(), bins=num_experts, min=0, max=num_experts - 1, + flat_experts.float(), + bins=num_experts, + min=0, + max=num_experts - 1, ).to(torch.int32) # 4. Pad for grouped_mm alignment - sorted_hidden_padded, padded_counts = _pad_to_alignment( - sorted_hidden, num_tokens_per_expert, num_experts - ) + sorted_hidden_padded, padded_counts = _pad_to_alignment(sorted_hidden, num_tokens_per_expert, num_experts) # 5. Grouped GEMM with LoRA expert_out_padded = _run_grouped_mm_with_lora( - sorted_hidden_padded, padded_counts, - gate_proj, up_proj, down_proj, - gate_proj_lora_A, gate_proj_lora_B, - up_proj_lora_A, up_proj_lora_B, - down_proj_lora_A, down_proj_lora_B, + sorted_hidden_padded, + padded_counts, + gate_proj, + up_proj, + down_proj, + gate_proj_lora_A, + gate_proj_lora_B, + up_proj_lora_A, + up_proj_lora_B, + down_proj_lora_A, + down_proj_lora_B, scaling, ) # 6. Unpad + apply routing weights total_sorted = num_tokens * top_k - expert_out = _unpad( - expert_out_padded, num_tokens_per_expert, padded_counts, num_experts, total_sorted - ) + expert_out = _unpad(expert_out_padded, num_tokens_per_expert, padded_counts, num_experts, total_sorted) expert_out = expert_out * sorted_weights.unsqueeze(-1) # 7. Scatter-add back @@ -527,6 +546,7 @@ def native_expert_lora_forward( # EP compute function (same interface as TritonEPGroupGemm.apply / QuackEPGroupGemm.apply) # --------------------------------------------------------------------------- + def native_ep_compute( permute_tokens: torch.Tensor, cumsum: torch.Tensor, @@ -568,6 +588,7 @@ def native_ep_compute( # EP compute with LoRA (same interface as TritonEPGroupGemmWithLoRA.apply) # --------------------------------------------------------------------------- + def native_ep_compute_lora( permute_tokens: torch.Tensor, cumsum: torch.Tensor, @@ -594,11 +615,17 @@ def native_ep_compute_lora( padded_tokens, padded_counts = _pad_to_alignment(permute_tokens, counts, num_local_experts) out_padded = _run_grouped_mm_with_lora( - padded_tokens, padded_counts, - gate_proj, up_proj, down_proj, - gate_proj_lora_A, gate_proj_lora_B, - up_proj_lora_A, up_proj_lora_B, - down_proj_lora_A, down_proj_lora_B, + padded_tokens, + padded_counts, + gate_proj, + up_proj, + down_proj, + gate_proj_lora_A, + gate_proj_lora_B, + up_proj_lora_A, + up_proj_lora_B, + down_proj_lora_A, + down_proj_lora_B, scaling, ) @@ -609,6 +636,7 @@ def native_ep_compute_lora( # moe_act variants: gate+up checkpointed, recomputed in backward # --------------------------------------------------------------------------- + def native_ep_compute_moe_act( permute_tokens: torch.Tensor, cumsum: torch.Tensor, @@ -650,7 +678,12 @@ def native_expert_forward_moe_act( instead of the compiled grouped GEMM. """ return _native_expert_forward_impl( - hidden_states, routing_weights, selected_experts, - gate_proj, up_proj, down_proj, num_experts, + hidden_states, + routing_weights, + selected_experts, + gate_proj, + up_proj, + down_proj, + num_experts, _run_experts_moe_act, ) diff --git a/src/xorl/models/layers/moe/experts.py b/src/xorl/models/layers/moe/experts.py index e9a90f6c..417687d5 100644 --- a/src/xorl/models/layers/moe/experts.py +++ b/src/xorl/models/layers/moe/experts.py @@ -1,6 +1,7 @@ """MoE expert weight container with backend dispatch.""" import os + import torch import torch.nn as nn @@ -105,12 +106,11 @@ def forward( # Check EP — use unified dispatch/compute/combine path from xorl.distributed.parallel_state import get_parallel_state + parallel_state = get_parallel_state() if parallel_state.ep_enabled: - return self._ep_forward( - hidden_states, routing_weights, selected_experts, parallel_state - ) + return self._ep_forward(hidden_states, routing_weights, selected_experts, parallel_state) # Local single-GPU path — select moe_act variant when available if _moe_act and self.moe_implementation in MOE_EXPERT_BACKENDS_MOE_ACT: @@ -144,7 +144,7 @@ def _ep_forward( Dispatch strategy is selected by ``self.ep_dispatch`` (``"alltoall"`` or ``"deepep"``). Compute backend by ``self.moe_implementation``. """ - from .backend import EP_DISPATCH, EP_COMBINE, EP_EXPERT_COMPUTE, EP_EXPERT_COMPUTE_MOE_ACT + from .backend import EP_COMBINE, EP_DISPATCH, EP_EXPERT_COMPUTE, EP_EXPERT_COMPUTE_MOE_ACT if self.moe_implementation not in EP_EXPERT_COMPUTE: raise ValueError( @@ -153,8 +153,7 @@ def _ep_forward( ) if self.ep_dispatch not in EP_DISPATCH: raise ValueError( - f"ep_dispatch={self.ep_dispatch!r} is not available. " - f"Available: {list(EP_DISPATCH.keys())}" + f"ep_dispatch={self.ep_dispatch!r} is not available. Available: {list(EP_DISPATCH.keys())}" ) dispatch_fn = EP_DISPATCH[self.ep_dispatch] @@ -168,13 +167,15 @@ def _ep_forward( compute_fn = EP_EXPERT_COMPUTE[self.moe_implementation] # Step 1: Dispatch tokens to expert-owning ranks - dispatch_kwargs = self._build_dispatch_kwargs( - hidden_states, routing_weights, selected_experts, parallel_state - ) + dispatch_kwargs = self._build_dispatch_kwargs(hidden_states, routing_weights, selected_experts, parallel_state) if _DEBUG_EP: return self._ep_forward_debug( - dispatch_fn, combine_fn, compute_fn, dispatch_kwargs, parallel_state, + dispatch_fn, + combine_fn, + compute_fn, + dispatch_kwargs, + parallel_state, ) permute_tokens, cumsum, ctx = dispatch_fn(**dispatch_kwargs) @@ -184,14 +185,15 @@ def _ep_forward( # Step 2: Expert computation (backend-specific GEMM only) expert_output = compute_fn( - permute_tokens, cumsum, - self.gate_proj, self.up_proj, self.down_proj, + permute_tokens, + cumsum, + self.gate_proj, + self.up_proj, + self.down_proj, ) # Step 3: Combine expert outputs back to original ranks - combine_kwargs = self._build_combine_kwargs( - expert_output, ctx, dispatch_kwargs, parallel_state - ) + combine_kwargs = self._build_combine_kwargs(expert_output, ctx, dispatch_kwargs, parallel_state) return combine_fn(**combine_kwargs) def _ep_forward_debug(self, dispatch_fn, combine_fn, compute_fn, dispatch_kwargs, parallel_state): @@ -215,15 +217,16 @@ def _ep_forward_debug(self, dispatch_fn, combine_fn, compute_fn, dispatch_kwargs # --- compute --- ev[2].record() expert_output = compute_fn( - permute_tokens, cumsum, - self.gate_proj, self.up_proj, self.down_proj, + permute_tokens, + cumsum, + self.gate_proj, + self.up_proj, + self.down_proj, ) ev[3].record() # --- combine --- - combine_kwargs = self._build_combine_kwargs( - expert_output, ctx, dispatch_kwargs, parallel_state - ) + combine_kwargs = self._build_combine_kwargs(expert_output, ctx, dispatch_kwargs, parallel_state) ev[4].record() result = combine_fn(**combine_kwargs) ev[5].record() @@ -264,6 +267,7 @@ def _build_dispatch_kwargs(self, hidden_states, routing_weights, selected_expert kwargs["ep_group"] = parallel_state.ep_group elif self.ep_dispatch == "deepep": from xorl.distributed.moe.deepep import get_default_buffer + kwargs["buffer"] = get_default_buffer( ep_group=parallel_state.ep_group, buffer_size_gb=self.deepep_buffer_size_gb, diff --git a/src/xorl/models/layers/moe/lora.py b/src/xorl/models/layers/moe/lora.py index 73fa065f..d29fb86c 100644 --- a/src/xorl/models/layers/moe/lora.py +++ b/src/xorl/models/layers/moe/lora.py @@ -15,18 +15,19 @@ down_proj_lora_A: [E, inter, r] down_proj_lora_B: [E, r, hidden] """ -from dataclasses import dataclass, field import math -from typing import Optional, List +from dataclasses import dataclass +from typing import List, Optional import torch import torch.nn as nn -from ....ops.group_gemm.kernel import compute_lora_scaling from ....lora.modules.base import LoraModule +from ....ops.group_gemm.kernel import compute_lora_scaling from ....utils import logging from ..activations import ACT2FN + logger = logging.get_logger(__name__) @@ -45,7 +46,6 @@ def __post_init__(self): self.target_modules = ["gate_proj", "up_proj", "down_proj"] - class MoEExpertsLoRA(LoraModule, nn.Module): """MoE experts with LoRA adapters. @@ -193,12 +193,11 @@ def forward( # Check EP — use unified dispatch/compute/combine path from xorl.distributed.parallel_state import get_parallel_state + parallel_state = get_parallel_state() if parallel_state.ep_enabled: - return self._ep_forward( - hidden_states, routing_weights, selected_experts, parallel_state - ) + return self._ep_forward(hidden_states, routing_weights, selected_experts, parallel_state) # Local path — registry-based from .backend import MOE_EXPERT_BACKENDS_LORA @@ -233,7 +232,7 @@ def _ep_forward( Uses the same dispatch/combine as ``MoEExperts._ep_forward()`` but routes to the LoRA-aware EP compute registry. """ - from .backend import EP_DISPATCH, EP_COMBINE, EP_EXPERT_COMPUTE_LORA + from .backend import EP_COMBINE, EP_DISPATCH, EP_EXPERT_COMPUTE_LORA if self.moe_implementation not in EP_EXPERT_COMPUTE_LORA: raise ValueError( @@ -242,8 +241,7 @@ def _ep_forward( ) if self.ep_dispatch not in EP_DISPATCH: raise ValueError( - f"ep_dispatch={self.ep_dispatch!r} is not available. " - f"Available: {list(EP_DISPATCH.keys())}" + f"ep_dispatch={self.ep_dispatch!r} is not available. Available: {list(EP_DISPATCH.keys())}" ) dispatch_fn = EP_DISPATCH[self.ep_dispatch] @@ -251,25 +249,27 @@ def _ep_forward( compute_fn = EP_EXPERT_COMPUTE_LORA[self.moe_implementation] # Step 1: Dispatch tokens to expert-owning ranks - dispatch_kwargs = self._build_dispatch_kwargs( - hidden_states, routing_weights, selected_experts, parallel_state - ) + dispatch_kwargs = self._build_dispatch_kwargs(hidden_states, routing_weights, selected_experts, parallel_state) permute_tokens, cumsum, ctx = dispatch_fn(**dispatch_kwargs) # Step 2: Expert computation with LoRA expert_output = compute_fn( - permute_tokens, cumsum, - self.gate_proj, self.up_proj, self.down_proj, - self.gate_proj_lora_A, self.gate_proj_lora_B, - self.up_proj_lora_A, self.up_proj_lora_B, - self.down_proj_lora_A, self.down_proj_lora_B, + permute_tokens, + cumsum, + self.gate_proj, + self.up_proj, + self.down_proj, + self.gate_proj_lora_A, + self.gate_proj_lora_B, + self.up_proj_lora_A, + self.up_proj_lora_B, + self.down_proj_lora_A, + self.down_proj_lora_B, self.scaling, ) # Step 3: Combine expert outputs back to original ranks - combine_kwargs = self._build_combine_kwargs( - expert_output, ctx, dispatch_kwargs, parallel_state - ) + combine_kwargs = self._build_combine_kwargs(expert_output, ctx, dispatch_kwargs, parallel_state) return combine_fn(**combine_kwargs) def _build_dispatch_kwargs(self, hidden_states, routing_weights, selected_experts, parallel_state): @@ -284,6 +284,7 @@ def _build_dispatch_kwargs(self, hidden_states, routing_weights, selected_expert kwargs["ep_group"] = parallel_state.ep_group elif self.ep_dispatch == "deepep": from xorl.distributed.moe.deepep import get_default_buffer + kwargs["buffer"] = get_default_buffer( ep_group=parallel_state.ep_group, buffer_size_gb=self.deepep_buffer_size_gb, @@ -439,10 +440,7 @@ def inject_lora_into_experts( block.experts = lora_experts - logger.debug( - f"Injected MoE LoRA with r={r}, alpha={lora_alpha}, " - f"target_modules={target_modules}" - ) + logger.debug(f"Injected MoE LoRA with r={r}, alpha={lora_alpha}, target_modules={target_modules}") # --------------------------------------------------------------------------- diff --git a/src/xorl/models/layers/moe/moe_block.py b/src/xorl/models/layers/moe/moe_block.py index b70d79c3..7a83ab0a 100644 --- a/src/xorl/models/layers/moe/moe_block.py +++ b/src/xorl/models/layers/moe/moe_block.py @@ -6,8 +6,8 @@ import torch.nn as nn import torch.nn.functional as F -from .router import TopKRouter from .experts import MoEExperts +from .router import TopKRouter from .routing_replay import RoutingReplay, get_replay_stage @@ -145,7 +145,7 @@ def forward(self, hidden_states: torch.Tensor): hidden_states = hidden_states.view(-1, hidden_dim) # Route (optionally upcast to fp32 for numerical alignment with SGLang) - if getattr(self, 'config', None) is not None and getattr(self.config, '_router_fp32', False): + if getattr(self, "config", None) is not None and getattr(self.config, "_router_fp32", False): router_logits = F.linear(hidden_states.float(), self.gate.weight.float()) else: router_logits = self.gate(hidden_states) @@ -188,9 +188,7 @@ def forward(self, hidden_states: torch.Tensor): ) else: # No replay active: use standard router - routing_weights, selected_experts = self.router( - router_logits, hidden_states.dtype - ) + routing_weights, selected_experts = self.router(router_logits, hidden_states.dtype) # When train_router is False, detach routing_weights so the gate is # only trained via auxiliary losses on router_logits, not through @@ -200,17 +198,11 @@ def forward(self, hidden_states: torch.Tensor): # Expert computation if self.moe_implementation == "eager": - final_hidden_states = self._eager_forward( - hidden_states, routing_weights, selected_experts - ) + final_hidden_states = self._eager_forward(hidden_states, routing_weights, selected_experts) else: - final_hidden_states = self.experts( - hidden_states, routing_weights, selected_experts - ) + final_hidden_states = self.experts(hidden_states, routing_weights, selected_experts) - final_hidden_states = final_hidden_states.reshape( - batch_size, sequence_length, hidden_dim - ) + final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) return final_hidden_states, router_logits def _eager_forward( @@ -223,21 +215,16 @@ def _eager_forward( hidden_dim = hidden_states.shape[-1] final_hidden_states = torch.zeros_like(hidden_states) - expert_mask = torch.nn.functional.one_hot( - selected_experts, num_classes=self.num_experts - ).permute(2, 1, 0) + expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) for expert_idx in range(self.num_experts): idx, top_x = torch.where(expert_mask[expert_idx]) current_state = hidden_states[None, top_x].reshape(-1, hidden_dim) current_hidden_states = ( - self.experts(current_state, expert_idx=expert_idx) - * routing_weights[top_x, idx, None] - ) - final_hidden_states.index_add_( - 0, top_x, current_hidden_states.to(hidden_states.dtype) + self.experts(current_state, expert_idx=expert_idx) * routing_weights[top_x, idx, None] ) + final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype)) return final_hidden_states diff --git a/src/xorl/models/layers/moe/router.py b/src/xorl/models/layers/moe/router.py index 8ada09d9..6b169fe2 100644 --- a/src/xorl/models/layers/moe/router.py +++ b/src/xorl/models/layers/moe/router.py @@ -45,9 +45,7 @@ def forward( selected_experts: ``(num_tokens, top_k)`` expert indices. """ routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) - routing_weights, selected_experts = torch.topk( - routing_weights, self.top_k, dim=-1 - ) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) if self.norm_topk_prob: routing_weights /= routing_weights.sum(dim=-1, keepdim=True) routing_weights = routing_weights.to(input_dtype) diff --git a/src/xorl/models/layers/moe/routing_replay.py b/src/xorl/models/layers/moe/routing_replay.py index 4f57c034..9c7ae3bb 100644 --- a/src/xorl/models/layers/moe/routing_replay.py +++ b/src/xorl/models/layers/moe/routing_replay.py @@ -75,9 +75,7 @@ def pop_forward_weights(self) -> Optional[torch.Tensor]: if not self.top_weights_list: return None # forward_index was already incremented by pop_forward, so use -1 - return self.top_weights_list[self.forward_index - 1].to( - torch.cuda.current_device(), non_blocking=True - ) + return self.top_weights_list[self.forward_index - 1].to(torch.cuda.current_device(), non_blocking=True) @torch.compiler.disable def pop_backward(self) -> torch.Tensor: @@ -91,9 +89,7 @@ def pop_backward_weights(self) -> Optional[torch.Tensor]: """Read routing weights for the last popped backward index, if available.""" if not self.top_weights_list: return None - return self.top_weights_list[self.backward_index - 1].to( - torch.cuda.current_device(), non_blocking=True - ) + return self.top_weights_list[self.backward_index - 1].to(torch.cuda.current_device(), non_blocking=True) @property def has_weights(self) -> bool: diff --git a/src/xorl/models/layers/rope.py b/src/xorl/models/layers/rope.py index 41de2fa3..5f5861d4 100644 --- a/src/xorl/models/layers/rope.py +++ b/src/xorl/models/layers/rope.py @@ -9,11 +9,11 @@ import math import warnings from functools import wraps -from typing import Optional import torch import torch.nn as nn + try: from flash_attn.layers.rotary import apply_rotary_emb as _flash_apply_rotary_emb except ImportError: @@ -27,6 +27,7 @@ # dynamic_rope_update decorator # --------------------------------------------------------------------------- + def dynamic_rope_update(rope_forward): """ Decorator function to update the RoPE parameters in the forward pass, if the model is using a dynamic RoPE @@ -123,6 +124,7 @@ def wrapper(self, x, position_ids, layer_type=None): # RoPE parameter computation functions # --------------------------------------------------------------------------- + def _compute_default_rope_parameters( config=None, device=None, @@ -158,9 +160,7 @@ def _compute_default_rope_parameters( head_dim = getattr(config, "head_dim", None) or (config.hidden_size // config.num_attention_heads) dim = int(head_dim * partial_rotary_factor) - inv_freq = 1.0 / ( - base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) - ) + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)) attention_factor = 1.0 return inv_freq, attention_factor @@ -418,6 +418,7 @@ def _compute_llama3_parameters( # RotaryEmbedding module # --------------------------------------------------------------------------- + class RotaryEmbedding(nn.Module): """Rotary Position Embedding layer. @@ -464,6 +465,7 @@ def forward(self, x, position_ids): # RoPE application helpers # --------------------------------------------------------------------------- + def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -524,6 +526,7 @@ def apply_rotary_pos_emb(q, k, cos, sin): # Deprecated helper (kept for backward compatibility) # --------------------------------------------------------------------------- + def rope_config_validation(config, ignore_keys=None): """ Deprecated function. Calls config.standardize_rope_params() and diff --git a/src/xorl/models/loader.py b/src/xorl/models/loader.py index 8dab1b6f..594d3723 100644 --- a/src/xorl/models/loader.py +++ b/src/xorl/models/loader.py @@ -8,7 +8,7 @@ ) from ..utils import logging -from .module_utils import init_empty_weights, all_ranks_load_weights +from .module_utils import all_ranks_load_weights, init_empty_weights from .registry import ModelRegistry diff --git a/src/xorl/models/module_utils.py b/src/xorl/models/module_utils.py index 971b52a5..a7c11234 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -1,4 +1,3 @@ - import json import os import re @@ -10,15 +9,13 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, Generator, List, Literal, Optional, Sequence, Set, Tuple, Union import torch - - +from safetensors import safe_open +from safetensors.torch import load_file, save_file from torch import distributed as dist from torch import nn from tqdm import tqdm from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME from transformers.utils.hub import cached_file, get_checkpoint_shard_files -from safetensors import safe_open -from safetensors.torch import load_file, save_file from xorl.distributed.parallel_state import get_parallel_state from xorl.lora.modules.linear import LoraLinear @@ -43,8 +40,8 @@ # Re-export checkpoint handler utilities for backward compatibility (used by tests) from xorl.models.checkpoint_handlers.buffers import ( # noqa: F401 ExpertWeightBuffer, - parse_expert_key, checkpoint_has_per_expert_weights, + parse_expert_key, ) @@ -95,7 +92,6 @@ def _get_checkpoint_keys(weights_path: str) -> Optional[Set[str]]: return None - @contextmanager def init_empty_weights(): """ @@ -285,13 +281,14 @@ def _load_state_dict(weights_path: str, **kwargs) -> List["StateDictIterator"]: Loads (sharded) state dict in transformers' format. """ import time + max_retries = 5 for attempt in range(max_retries): result = _try_load_state_dict(weights_path, **kwargs) if result is not None: return result if attempt < max_retries - 1: - retry_delay = 2 * (2 ** attempt) # 2, 4, 8, 16s + retry_delay = 2 * (2**attempt) # 2, 4, 8, 16s logger.warning( f"Cannot find checkpoint files in {weights_path} (attempt {attempt + 1}/{max_retries}). " f"Retrying in {retry_delay}s..." @@ -322,8 +319,7 @@ def _try_load_state_dict(weights_path: str, **kwargs): import time retry_delay = 5 logger.warning( - f"OSError getting shard files (attempt {attempt + 1}/{max_retries}): {e}. " - f"Retrying in {retry_delay}s..." + f"OSError resolving shard files (attempt {attempt + 1}/{max_retries}): {e}. Retrying in 5s..." ) time.sleep(retry_delay) else: @@ -429,9 +425,7 @@ def dispatch( to avoid interleaving no-op distribute_tensor calls (ep_fsdp mesh, size 1) with real NCCL collectives (fsdp mesh, size > 1). """ - assert tensor.device.type == "cpu", ( - f"DMA scheduler only handles CPU→CUDA, got {tensor.device} for {full_name}" - ) + assert tensor.device.type == "cpu", f"DMA scheduler only handles CPU→CUDA, got {tensor.device} for {full_name}" if tensor.dtype != orig_tensor.dtype: tensor = tensor.to(dtype=orig_tensor.dtype) pinned = tensor.pin_memory() @@ -457,9 +451,7 @@ def flush(self) -> None: if dtensor_factory is not None and hasattr(orig_tensor, "device_mesh"): device_mesh = getattr(orig_tensor, "device_mesh") placements = getattr(orig_tensor, "placements") - module._parameters[local_name].data.copy_( - dtensor_factory(gpu_temp, device_mesh, placements) - ) + module._parameters[local_name].data.copy_(dtensor_factory(gpu_temp, device_mesh, placements)) # distribute_tensor uses NCCL on a different CUDA stream. # Synchronize to prevent races with subsequent DMA transfers # on the copy-engine streams. @@ -505,11 +497,7 @@ def _dispatch_parameter( # would interleave no-op distribute_tensor calls (expert params on # ep_fsdp mesh, size 1) with real NCCL collectives (regular params on # fsdp mesh, size > 1), which can corrupt weights. - if ( - _active_dma_scheduler is not None - and tensor.device.type == "cpu" - and orig_tensor.device.type == "cuda" - ): + if _active_dma_scheduler is not None and tensor.device.type == "cpu" and orig_tensor.device.type == "cuda": _active_dma_scheduler.dispatch(module, local_name, tensor, orig_tensor, dtensor_factory, full_param_name) return @@ -537,6 +525,7 @@ def _dispatch_parameter( # CPU DTensor: copy shard directly into local tensor. # distribute_tensor doesn't support CPU mesh, so we manually shard and copy. from torch.distributed._tensor import Shard as DTShard + shard_dim = None for p in placements: if isinstance(p, DTShard): @@ -596,9 +585,10 @@ def _init_parameter( - lora_A: kaiming uniform initialization - lora_B: zeros (so LoRA has no effect at start) """ - import torch.nn as nn import math + import torch.nn as nn + # Check if this is a LoRA parameter and handle specially if "lora_A" in name or "lora_B" in name: # Navigate to the parameter @@ -664,7 +654,6 @@ def _convert_weight_key(key: str, model: "PreTrainedModel") -> str: return key - def _shrink_expert_params_for_ep(model: "nn.Module") -> None: """Shrink expert parameters to EP-local shapes before materialization. @@ -778,6 +767,7 @@ def all_ranks_load_weights( f"Retrying in {retry_delay} seconds..." ) import time + time.sleep(retry_delay) else: logger.error(f"Failed to load weights after {max_retries} attempts") @@ -928,7 +918,10 @@ def _dispatch_results(results): _active_dma_scheduler = None post_process_after_weight_loading( - model, buffer_dict, parameter_names_to_load, dtensor_factory, + model, + buffer_dict, + parameter_names_to_load, + dtensor_factory, qlora_skip_prefixes=_expected_skip_prefixes, qlora_skip_fn=_should_skip_qlora_expert_key, ) @@ -1070,8 +1063,7 @@ def rank0_load_and_broadcast_weights( start_time = time.perf_counter() dist.broadcast(tensor, src=0) logger.info_rank0( - f"{name=}, {shape=}, {dtype=}, broadcast time (ms): " - f"{1000 * (time.perf_counter() - start_time)}" + f"{name=}, {shape=}, {dtype=}, broadcast time (ms): {1000 * (time.perf_counter() - start_time)}" ) # Resolve _orig_mod prefix from torch.compile @@ -1336,7 +1328,6 @@ def save_model_weights( logger.warning(f"Model asset {model_asset} should implement `save_pretrained`.") - def save_model_assets(output_dir: Union[str, "os.PathLike"], model_assets: Sequence["ModelAssets"]): for model_asset in model_assets: if hasattr(model_asset, "save_pretrained"): diff --git a/src/xorl/models/transformers/__init__.py b/src/xorl/models/transformers/__init__.py index 6fde8dfc..830d6933 100644 --- a/src/xorl/models/transformers/__init__.py +++ b/src/xorl/models/transformers/__init__.py @@ -1,8 +1,8 @@ from . import ( qwen3, qwen3_5, - qwen3_moe, qwen3_5_moe, + qwen3_moe, ) diff --git a/src/xorl/models/transformers/qwen3/checkpoint_handler.py b/src/xorl/models/transformers/qwen3/checkpoint_handler.py index 8da80cb2..53dbb675 100644 --- a/src/xorl/models/transformers/qwen3/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3/checkpoint_handler.py @@ -7,10 +7,15 @@ from ...checkpoint_handlers.base import CheckpointHandler from ...checkpoint_handlers.buffers import ( - GateUpMergeBuffer, QKVMergeBuffer, QLoRAWeightBuffer, - QUANT_AUX_SUFFIX_PATTERN, FP8_AUX_SUFFIX_PATTERN, - QKV_PROJ_PATTERN, DENSE_GATE_UP_PATTERN, - OPROJ_WEIGHT_PATTERN, DENSE_DOWN_PROJ_PATTERN, + DENSE_DOWN_PROJ_PATTERN, + DENSE_GATE_UP_PATTERN, + FP8_AUX_SUFFIX_PATTERN, + OPROJ_WEIGHT_PATTERN, + QKV_PROJ_PATTERN, + QUANT_AUX_SUFFIX_PATTERN, + GateUpMergeBuffer, + QKVMergeBuffer, + QLoRAWeightBuffer, ) @@ -86,8 +91,12 @@ def _should_skip(key: str) -> bool: # Skip linear projection weight keys loaded by QLoRALinear directly. # Bias keys (.bias) are NOT skipped — they merge normally. if key.endswith(".weight"): - if (QKV_PROJ_PATTERN.match(key) or DENSE_GATE_UP_PATTERN.match(key) - or OPROJ_WEIGHT_PATTERN.match(key) or DENSE_DOWN_PROJ_PATTERN.match(key)): + if ( + QKV_PROJ_PATTERN.match(key) + or DENSE_GATE_UP_PATTERN.match(key) + or OPROJ_WEIGHT_PATTERN.match(key) + or DENSE_DOWN_PROJ_PATTERN.match(key) + ): return True return False @@ -101,9 +110,7 @@ def _is_excluded_module(self, key: str) -> bool: module_short_name = module_fqn.rsplit(".", 1)[-1] return module_short_name in self._exclude_modules - def on_load_weight( - self, key: str, tensor: torch.Tensor - ) -> List[Tuple[str, torch.Tensor]]: + def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: # Drop input_scale (unused by our quantization) if key.endswith(".input_scale"): return [] @@ -152,6 +159,7 @@ def on_load_weight( def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: import warnings + pending_gu = self._gate_up_buffer.get_pending() if pending_gu: warnings.warn(f"Incomplete gate/up merge pairs after loading: {pending_gu}") @@ -163,9 +171,7 @@ def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: self._qlora_buffer.set_inline_metadata() return [] - def on_save_weight( - self, param_name: str, tensor: torch.Tensor - ) -> List[Tuple[str, torch.Tensor]]: + def on_save_weight(self, param_name: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: # Split gate_up_proj -> gate_proj + up_proj if ".gate_up_proj." in param_name: prefix, suffix = param_name.rsplit(".gate_up_proj.", 1) diff --git a/src/xorl/models/transformers/qwen3/modeling_qwen3.py b/src/xorl/models/transformers/qwen3/modeling_qwen3.py index 93dc5ca6..4b5b441a 100644 --- a/src/xorl/models/transformers/qwen3/modeling_qwen3.py +++ b/src/xorl/models/transformers/qwen3/modeling_qwen3.py @@ -6,6 +6,11 @@ from xorl.distributed.parallel_state import get_parallel_state from xorl.distributed.sequence_parallel.strategy import get_cp_strategy from xorl.models.base import XorlPreTrainedModel +from xorl.models.checkpoint_handlers.buffers import ( + detect_prequantized_block_fp8_checkpoint, + detect_prequantized_checkpoint, + get_prequantized_exclude_modules, +) from xorl.models.layers import ACT2FN, RMSNorm, RotaryEmbedding from xorl.models.layers.attention import ( AttentionKwargs, @@ -15,14 +20,9 @@ ) from xorl.models.module_utils import GradientCheckpointingLayer from xorl.models.outputs import BaseModelOutput, CausalLMOutput -from xorl.models.checkpoint_handlers.buffers import ( - detect_prequantized_checkpoint, - detect_prequantized_block_fp8_checkpoint, - get_prequantized_exclude_modules, -) +from xorl.models.transformers.qwen3 import parallelize from xorl.models.transformers.qwen3.checkpoint_handler import Qwen3CheckpointHandler from xorl.models.transformers.qwen3.configuration_qwen3 import Qwen3Config -from xorl.models.transformers.qwen3 import parallelize from xorl.ops.fused_silu_and_mul import fused_silu_and_mul from xorl.utils import logging @@ -38,7 +38,7 @@ def __init__(self, config): self.gate_up_proj = nn.Linear(self.hidden_size, 2 * self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[config.hidden_act] - self._use_fused_silu = config.hidden_act == "silu" and not getattr(config, '_activation_native', False) + self._use_fused_silu = config.hidden_act == "silu" and not getattr(config, "_activation_native", False) def unfuse_for_tp(self): """Replace fused gate_up_proj with separate gate_proj and up_proj for tensor parallelism.""" @@ -82,8 +82,8 @@ def __init__(self, config: Qwen3Config, layer_idx: int): self.mlp = Qwen3MLP(config) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - if ( - config.sliding_window and not is_flash_attention(config._attn_implementation) + if config.sliding_window and not is_flash_attention( + config._attn_implementation ): # diff with Llama is this warning logger.warning_once( f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " @@ -257,7 +257,9 @@ def forward( # SP strategy handles slicing (sync: slice, async: keep full-length) ps = get_parallel_state() position_embeddings = get_cp_strategy(num_kv_heads=self.config.num_key_value_heads).prepare_position_embeddings( - position_embeddings, dim=1, sp_group=ps.sp_group, + position_embeddings, + dim=1, + sp_group=ps.sp_group, num_kv_heads=self.config.num_key_value_heads, ) diff --git a/src/xorl/models/transformers/qwen3/parallelize.py b/src/xorl/models/transformers/qwen3/parallelize.py index e5df8b99..7ee99cc9 100644 --- a/src/xorl/models/transformers/qwen3/parallelize.py +++ b/src/xorl/models/transformers/qwen3/parallelize.py @@ -1,6 +1,5 @@ """Parallelization plan and utilities for dense Qwen3 models.""" - # TP plan for the base model (Qwen3Model). # Keys use wildcard patterns relative to the base model prefix. TP_PLAN = { diff --git a/src/xorl/models/transformers/qwen3_5/__init__.py b/src/xorl/models/transformers/qwen3_5/__init__.py index d447924c..3ca2f676 100644 --- a/src/xorl/models/transformers/qwen3_5/__init__.py +++ b/src/xorl/models/transformers/qwen3_5/__init__.py @@ -7,6 +7,7 @@ Qwen3_5Model, ) + __all__ = [ "Qwen3_5Config", "Qwen3_5ForCausalLM", diff --git a/src/xorl/models/transformers/qwen3_5/checkpoint_handler.py b/src/xorl/models/transformers/qwen3_5/checkpoint_handler.py index 6bf0c0c2..163201b2 100644 --- a/src/xorl/models/transformers/qwen3_5/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3_5/checkpoint_handler.py @@ -6,14 +6,14 @@ from ...checkpoint_handlers.base import CheckpointHandler from ...checkpoint_handlers.buffers import ( - GateUpMergeBuffer, - QKVMergeBuffer, - QUANT_AUX_SUFFIX_PATTERN, - FP8_AUX_SUFFIX_PATTERN, - QKV_PROJ_PATTERN, + DENSE_DOWN_PROJ_PATTERN, DENSE_GATE_UP_PATTERN, + FP8_AUX_SUFFIX_PATTERN, OPROJ_WEIGHT_PATTERN, - DENSE_DOWN_PROJ_PATTERN, + QKV_PROJ_PATTERN, + QUANT_AUX_SUFFIX_PATTERN, + GateUpMergeBuffer, + QKVMergeBuffer, ) from ..qwen3_5_shared import ( is_excluded_module_key, @@ -55,7 +55,9 @@ def __init__( self._is_prequantized = is_prequantized self._exclude_modules = exclude_modules or set() - def _handle_linear_attention_weights(self, key: str, tensor: torch.Tensor) -> Optional[List[Tuple[str, torch.Tensor]]]: + def _handle_linear_attention_weights( + self, key: str, tensor: torch.Tensor + ) -> Optional[List[Tuple[str, torch.Tensor]]]: return map_qwen3_5_linear_attention_weight( key, tensor, @@ -108,9 +110,7 @@ def _is_excluded_module(self, key: str) -> bool: """Check if a key belongs to a module excluded from quantization.""" return is_excluded_module_key(key, self._exclude_modules) - def on_load_weight( - self, key: str, tensor: torch.Tensor - ) -> List[Tuple[str, torch.Tensor]]: + def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: # Pre-quantized: drop any quant auxiliary or weight keys that weren't # caught by get_skip_key_fn (safety net). # Excluded modules pass through as bf16 - don't drop them. @@ -165,9 +165,7 @@ def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: warnings.warn(f"Incomplete QKV merge groups after loading: {pending_qkv}") return [] - def on_save_weight( - self, param_name: str, tensor: torch.Tensor - ) -> List[Tuple[str, torch.Tensor]]: + def on_save_weight(self, param_name: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: # Split gate_up_proj -> gate_proj + up_proj if ".gate_up_proj." in param_name: prefix, suffix = param_name.rsplit(".gate_up_proj.", 1) diff --git a/src/xorl/models/transformers/qwen3_5/configuration_qwen3_5.py b/src/xorl/models/transformers/qwen3_5/configuration_qwen3_5.py index d3425ab0..2678284e 100644 --- a/src/xorl/models/transformers/qwen3_5/configuration_qwen3_5.py +++ b/src/xorl/models/transformers/qwen3_5/configuration_qwen3_5.py @@ -1,6 +1,7 @@ """Qwen3_5 model configuration""" from transformers.configuration_utils import PretrainedConfig + from xorl.models.layers import rope_config_validation from .parallelize import TP_PLAN @@ -19,7 +20,7 @@ def _cfg_to_dict(value): return None if isinstance(value, dict): return dict(value) - if hasattr(value, '__dict__'): + if hasattr(value, "__dict__"): return {k: v for k, v in vars(value).items()} return value @@ -32,13 +33,13 @@ def _split_mrope_fields(value): class Qwen3_5Config(PretrainedConfig): - model_type = 'xorl_qwen3_5' + model_type = "xorl_qwen3_5" base_model_tp_plan = TP_PLAN base_model_pp_plan = { - 'embed_tokens': (['input_ids'], ['inputs_embeds']), - 'layers': (['hidden_states', 'attention_mask'], ['hidden_states']), - 'norm': (['hidden_states'], ['hidden_states']), + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), } def __init__( @@ -50,7 +51,7 @@ def __init__( num_attention_heads=16, num_key_value_heads=4, head_dim=256, - hidden_act='silu', + hidden_act="silu", max_position_embeddings=32768, initializer_range=0.02, rms_norm_eps=1e-6, @@ -78,8 +79,8 @@ def __init__( eos_token_id=None, **kwargs, ): - kwargs['ignore_keys_at_rope_validation'] = {'mrope_section', 'mrope_interleaved'} - kwargs.setdefault('partial_rotary_factor', 0.25) + kwargs["ignore_keys_at_rope_validation"] = {"mrope_section", "mrope_interleaved"} + kwargs.setdefault("partial_rotary_factor", 0.25) self.pad_token_id = pad_token_id self.bos_token_id = bos_token_id @@ -109,24 +110,28 @@ def __init__( self.attention_dropout = attention_dropout self.full_attention_interval = full_attention_interval self.linear_num_key_heads = linear_num_key_heads if linear_num_key_heads is not None else num_attention_heads - self.linear_num_value_heads = linear_num_value_heads if linear_num_value_heads is not None else num_attention_heads + self.linear_num_value_heads = ( + linear_num_value_heads if linear_num_value_heads is not None else num_attention_heads + ) self.linear_key_head_dim = linear_key_head_dim if linear_key_head_dim is not None else 128 - self.linear_value_head_dim = linear_value_head_dim if linear_value_head_dim is not None else self.linear_key_head_dim + self.linear_value_head_dim = ( + linear_value_head_dim if linear_value_head_dim is not None else self.linear_key_head_dim + ) self.attn_output_gate = attn_output_gate self.linear_conv_kernel_dim = linear_conv_kernel_dim if layer_types is None: if full_attention_interval: layer_types = [ - 'full_attention' if (layer_idx + 1) % full_attention_interval == 0 else 'linear_attention' + "full_attention" if (layer_idx + 1) % full_attention_interval == 0 else "linear_attention" for layer_idx in range(num_hidden_layers) ] else: - layer_types = ['full_attention'] * num_hidden_layers + layer_types = ["full_attention"] * num_hidden_layers self.layer_types = list(layer_types) - if self._rope_scaling is not None and 'type' in self._rope_scaling: - self._rope_scaling['rope_type'] = self._rope_scaling['type'] - rope_config_validation(self, ignore_keys=kwargs.get('ignore_keys_at_rope_validation')) + if self._rope_scaling is not None and "type" in self._rope_scaling: + self._rope_scaling["rope_type"] = self._rope_scaling["type"] + rope_config_validation(self, ignore_keys=kwargs.get("ignore_keys_at_rope_validation")) super().__init__( pad_token_id=pad_token_id, @@ -139,71 +144,73 @@ def __init__( @property def rope_parameters(self): rope_params = { - 'rope_type': 'default', - 'rope_theta': self.rope_theta, + "rope_type": "default", + "rope_theta": self.rope_theta, } if self._rope_scaling is not None: rope_params.update(self._rope_scaling) - if 'type' in rope_params and 'rope_type' not in rope_params: - rope_params['rope_type'] = rope_params.pop('type') + if "type" in rope_params and "rope_type" not in rope_params: + rope_params["rope_type"] = rope_params.pop("type") return rope_params @rope_parameters.setter def rope_parameters(self, value): value_dict = _cfg_to_dict(value) if value_dict is not None and isinstance(value_dict, dict): - if 'rope_theta' in value_dict: - self.rope_theta = value_dict['rope_theta'] + if "rope_theta" in value_dict: + self.rope_theta = value_dict["rope_theta"] cleaned_value, rope_mrope_interleaved, rope_mrope_section = _split_mrope_fields(value_dict) - self.mrope_interleaved = rope_mrope_interleaved or getattr(self, 'mrope_interleaved', False) + self.mrope_interleaved = rope_mrope_interleaved or getattr(self, "mrope_interleaved", False) if rope_mrope_section is not None: self.mrope_section = rope_mrope_section self._rope_scaling = cleaned_value @classmethod def from_hf_config(cls, hf_config): - text_config = getattr(hf_config, 'text_config', hf_config) - rope_params = getattr(text_config, 'rope_parameters', None) + text_config = getattr(hf_config, "text_config", hf_config) + rope_params = getattr(text_config, "rope_parameters", None) if rope_params is None: - rope_params = getattr(text_config, 'rope_scaling', None) + rope_params = getattr(text_config, "rope_scaling", None) - hidden_size = getattr(text_config, 'hidden_size') - num_attention_heads = getattr(text_config, 'num_attention_heads') - head_dim = getattr(text_config, 'head_dim', hidden_size // num_attention_heads) + hidden_size = getattr(text_config, "hidden_size") + num_attention_heads = getattr(text_config, "num_attention_heads") + head_dim = getattr(text_config, "head_dim", hidden_size // num_attention_heads) return cls( - vocab_size=getattr(text_config, 'vocab_size', getattr(hf_config, 'vocab_size', 248320)), + vocab_size=getattr(text_config, "vocab_size", getattr(hf_config, "vocab_size", 248320)), hidden_size=hidden_size, - intermediate_size=getattr(text_config, 'intermediate_size'), - num_hidden_layers=getattr(text_config, 'num_hidden_layers'), + intermediate_size=getattr(text_config, "intermediate_size"), + num_hidden_layers=getattr(text_config, "num_hidden_layers"), num_attention_heads=num_attention_heads, - num_key_value_heads=getattr(text_config, 'num_key_value_heads', num_attention_heads), + num_key_value_heads=getattr(text_config, "num_key_value_heads", num_attention_heads), head_dim=head_dim, - hidden_act=getattr(text_config, 'hidden_act', 'silu'), - max_position_embeddings=getattr(text_config, 'max_position_embeddings'), - initializer_range=getattr(text_config, 'initializer_range', 0.02), - rms_norm_eps=getattr(text_config, 'rms_norm_eps', 1e-6), - use_cache=getattr(text_config, 'use_cache', True), - tie_word_embeddings=getattr(hf_config, 'tie_word_embeddings', getattr(text_config, 'tie_word_embeddings', True)), - rope_theta=_cfg_get(rope_params, 'rope_theta', getattr(text_config, 'rope_theta', 10000.0)), + hidden_act=getattr(text_config, "hidden_act", "silu"), + max_position_embeddings=getattr(text_config, "max_position_embeddings"), + initializer_range=getattr(text_config, "initializer_range", 0.02), + rms_norm_eps=getattr(text_config, "rms_norm_eps", 1e-6), + use_cache=getattr(text_config, "use_cache", True), + tie_word_embeddings=getattr( + hf_config, "tie_word_embeddings", getattr(text_config, "tie_word_embeddings", True) + ), + rope_theta=_cfg_get(rope_params, "rope_theta", getattr(text_config, "rope_theta", 10000.0)), rope_scaling=_cfg_to_dict(rope_params), - attention_bias=getattr(text_config, 'attention_bias', False), - attention_dropout=getattr(text_config, 'attention_dropout', 0.0), - layer_types=getattr(text_config, 'layer_types', None), - full_attention_interval=getattr(text_config, 'full_attention_interval', None), - linear_num_key_heads=getattr(text_config, 'linear_num_key_heads', None), - linear_num_value_heads=getattr(text_config, 'linear_num_value_heads', None), - linear_key_head_dim=getattr(text_config, 'linear_key_head_dim', None), - linear_value_head_dim=getattr(text_config, 'linear_value_head_dim', None), - attn_output_gate=getattr(text_config, 'attn_output_gate', True), - linear_conv_kernel_dim=getattr(text_config, 'linear_conv_kernel_dim', 4), - mrope_interleaved=_cfg_get(rope_params, 'mrope_interleaved', False), - mrope_section=_cfg_get(rope_params, 'mrope_section', None), - pad_token_id=getattr(text_config, 'pad_token_id', getattr(hf_config, 'pad_token_id', None)), - bos_token_id=getattr(text_config, 'bos_token_id', getattr(hf_config, 'bos_token_id', None)), - eos_token_id=getattr(text_config, 'eos_token_id', getattr(hf_config, 'eos_token_id', None)), - architectures=['Qwen3_5ForConditionalGeneration'], + attention_bias=getattr(text_config, "attention_bias", False), + attention_dropout=getattr(text_config, "attention_dropout", 0.0), + layer_types=getattr(text_config, "layer_types", None), + full_attention_interval=getattr(text_config, "full_attention_interval", None), + linear_num_key_heads=getattr(text_config, "linear_num_key_heads", None), + linear_num_value_heads=getattr(text_config, "linear_num_value_heads", None), + linear_key_head_dim=getattr(text_config, "linear_key_head_dim", None), + linear_value_head_dim=getattr(text_config, "linear_value_head_dim", None), + attn_output_gate=getattr(text_config, "attn_output_gate", True), + linear_conv_kernel_dim=getattr(text_config, "linear_conv_kernel_dim", 4), + mrope_interleaved=_cfg_get(rope_params, "mrope_interleaved", False), + mrope_section=_cfg_get(rope_params, "mrope_section", None), + pad_token_id=getattr(text_config, "pad_token_id", getattr(hf_config, "pad_token_id", None)), + bos_token_id=getattr(text_config, "bos_token_id", getattr(hf_config, "bos_token_id", None)), + eos_token_id=getattr(text_config, "eos_token_id", getattr(hf_config, "eos_token_id", None)), + architectures=["Qwen3_5ForConditionalGeneration"], ) -__all__ = ['Qwen3_5Config'] +__all__ = ["Qwen3_5Config"] diff --git a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py index ee2039fa..7bf43347 100644 --- a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py +++ b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py @@ -6,6 +6,11 @@ from xorl.distributed.parallel_state import get_parallel_state from xorl.distributed.sequence_parallel.strategy import get_cp_strategy from xorl.models.base import XorlPreTrainedModel +from xorl.models.checkpoint_handlers.buffers import ( + detect_prequantized_block_fp8_checkpoint, + detect_prequantized_checkpoint, + get_prequantized_exclude_modules, +) from xorl.models.layers import ACT2FN, RotaryEmbedding from xorl.models.layers.attention import ( AttentionKwargs, @@ -16,11 +21,9 @@ from xorl.models.layers.attention.backend.eager import eager_attention_forward from xorl.models.module_utils import GradientCheckpointingLayer from xorl.models.outputs import BaseModelOutput, CausalLMOutput -from xorl.models.checkpoint_handlers.buffers import ( - detect_prequantized_checkpoint, - detect_prequantized_block_fp8_checkpoint, - get_prequantized_exclude_modules, -) +from xorl.models.transformers.qwen3_5 import parallelize +from xorl.models.transformers.qwen3_5.checkpoint_handler import Qwen3_5CheckpointHandler +from xorl.models.transformers.qwen3_5.configuration_qwen3_5 import Qwen3_5Config from xorl.models.transformers.qwen3_5_shared import ( LINEAR_ATTENTION_RING_UNSUPPORTED_MESSAGE, QWEN3_5_CHECKPOINT_CONVERSION_MAPPING, @@ -28,12 +31,9 @@ has_linear_attention_layers, qwen3_5_apply_rotary_pos_emb, ) -from xorl.models.transformers.qwen3_5.checkpoint_handler import Qwen3_5CheckpointHandler -from xorl.models.transformers.qwen3_5.configuration_qwen3_5 import Qwen3_5Config -from xorl.models.transformers.qwen3_5 import parallelize +from xorl.ops.fused_silu_and_mul import fused_silu_and_mul from xorl.ops.linear_attention import GatedDeltaNet from xorl.ops.linear_attention.ops.cp import build_linear_attention_cp_context -from xorl.ops.fused_silu_and_mul import fused_silu_and_mul from xorl.utils import logging @@ -185,7 +185,9 @@ def forward( del position_ids, past_key_values attn_strategy = get_cp_strategy() query_states, key_states, value_states = attn_strategy.project_qkv(self, hidden_states, position_embeddings) - attn_output = attn_strategy.compute_attention(self, query_states, key_states, value_states, attention_mask, **kwargs) + attn_output = attn_strategy.compute_attention( + self, query_states, key_states, value_states, attention_mask, **kwargs + ) attn_output = attn_strategy.project_output(self, attn_output) return attn_output, None @@ -216,9 +218,7 @@ def __init__(self, config: Qwen3_5Config, layer_idx: int): self.mlp = Qwen3_5MLP(config) self.input_layernorm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - if ( - config.sliding_window and not is_flash_attention(config._attn_implementation) - ): + if config.sliding_window and not is_flash_attention(config._attn_implementation): logger.warning_once( f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " "unexpected results may be encountered." diff --git a/src/xorl/models/transformers/qwen3_5/parallelize.py b/src/xorl/models/transformers/qwen3_5/parallelize.py index 26fb8b8f..da0549b8 100644 --- a/src/xorl/models/transformers/qwen3_5/parallelize.py +++ b/src/xorl/models/transformers/qwen3_5/parallelize.py @@ -1,6 +1,5 @@ """Parallelization plan and utilities for dense Qwen3_5 models.""" - # TP plan for the base model (Qwen3_5Model). # Keys use wildcard patterns relative to the base model prefix. TP_PLAN = { diff --git a/src/xorl/models/transformers/qwen3_5_moe/__init__.py b/src/xorl/models/transformers/qwen3_5_moe/__init__.py index 290097c0..63315ad4 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/__init__.py +++ b/src/xorl/models/transformers/qwen3_5_moe/__init__.py @@ -7,6 +7,7 @@ Qwen3_5MoeModel, ) + __all__ = [ "Qwen3_5MoeConfig", "Qwen3_5MoeForCausalLM", diff --git a/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py b/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py index 3e9db4d1..b8307ce7 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py @@ -7,11 +7,17 @@ from ...checkpoint_handlers.base import CheckpointHandler from ...checkpoint_handlers.buffers import ( - ExpertWeightBuffer, GateUpMergeBuffer, QKVMergeBuffer, - parse_expert_key, EXPERT_QUANT_AUX_PATTERN, QUANT_AUX_SUFFIX_PATTERN, + DENSE_DOWN_PROJ_PATTERN, + DENSE_GATE_UP_PATTERN, + EXPERT_QUANT_AUX_PATTERN, FP8_AUX_SUFFIX_PATTERN, - QKV_PROJ_PATTERN, DENSE_GATE_UP_PATTERN, - OPROJ_WEIGHT_PATTERN, DENSE_DOWN_PROJ_PATTERN, + OPROJ_WEIGHT_PATTERN, + QKV_PROJ_PATTERN, + QUANT_AUX_SUFFIX_PATTERN, + ExpertWeightBuffer, + GateUpMergeBuffer, + QKVMergeBuffer, + parse_expert_key, ) from ..qwen3_5_shared import ( is_excluded_module_key, @@ -65,9 +71,7 @@ def __init__( # Disable expert buffer when pre-quantized: expert weights are loaded # directly by QLoRAMoeExperts, not merged through the checkpoint handler. if checkpoint_has_per_expert and not is_prequantized: - self._expert_buffer = ExpertWeightBuffer( - num_experts, ep_rank=ep_rank, ep_size=ep_size - ) + self._expert_buffer = ExpertWeightBuffer(num_experts, ep_rank=ep_rank, ep_size=ep_size) self._qkv_buffer: Optional[QKVMergeBuffer] = None if not skip_qkv_merge: self._qkv_buffer = QKVMergeBuffer() @@ -86,17 +90,13 @@ def __init__( self._expert_start = ep_rank * self._local_num_experts self._expert_end = self._expert_start + self._local_num_experts - _FUSED_EXPERT_GATE_UP_PATTERN = re.compile( - r"^model\.layers\.(\d+)\.mlp\.experts\.gate_up_proj(?:\.weight)?$" - ) - _FUSED_EXPERT_DOWN_PATTERN = re.compile( - r"^model\.layers\.(\d+)\.mlp\.experts\.down_proj(?:\.weight)?$" - ) + _FUSED_EXPERT_GATE_UP_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.gate_up_proj(?:\.weight)?$") + _FUSED_EXPERT_DOWN_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.down_proj(?:\.weight)?$") def _slice_expert_tensor_for_ep(self, tensor: torch.Tensor) -> torch.Tensor: if self._ep_size == 1: return tensor - return tensor[self._expert_start:self._expert_end].contiguous() + return tensor[self._expert_start : self._expert_end].contiguous() def _handle_fused_expert_weights(self, key: str, tensor: torch.Tensor) -> Optional[List[Tuple[str, torch.Tensor]]]: gate_up_match = self._FUSED_EXPERT_GATE_UP_PATTERN.match(key) @@ -122,7 +122,9 @@ def _handle_fused_expert_weights(self, key: str, tensor: torch.Tensor) -> Option return None - def _handle_linear_attention_weights(self, key: str, tensor: torch.Tensor) -> Optional[List[Tuple[str, torch.Tensor]]]: + def _handle_linear_attention_weights( + self, key: str, tensor: torch.Tensor + ) -> Optional[List[Tuple[str, torch.Tensor]]]: return map_qwen3_5_linear_attention_weight( key, tensor, @@ -138,10 +140,8 @@ def get_skip_key_fn(self) -> Optional[Callable[[str], bool]]: - All quantized auxiliary keys when is_prequantized=True (weight_scale, weight_scale_2, input_scale) — these are loaded directly by QLoRA modules """ - has_ep_filter = ( - self._expert_buffer is not None - and not (self._expert_buffer.expert_start == 0 - and self._expert_buffer.expert_end == self._expert_buffer.num_experts) + has_ep_filter = self._expert_buffer is not None and not ( + self._expert_buffer.expert_start == 0 and self._expert_buffer.expert_end == self._expert_buffer.num_experts ) if not has_ep_filter and not self._is_prequantized: @@ -176,8 +176,12 @@ def _should_skip(key: str) -> bool: # Skip all linear projection weight keys that are loaded by QLoRALinear # directly. Bias keys (.bias) are NOT skipped (merged normally). if key.endswith(".weight"): - if (QKV_PROJ_PATTERN.match(key) or DENSE_GATE_UP_PATTERN.match(key) - or OPROJ_WEIGHT_PATTERN.match(key) or DENSE_DOWN_PROJ_PATTERN.match(key)): + if ( + QKV_PROJ_PATTERN.match(key) + or DENSE_GATE_UP_PATTERN.match(key) + or OPROJ_WEIGHT_PATTERN.match(key) + or DENSE_DOWN_PROJ_PATTERN.match(key) + ): return True # Skip out-of-range expert keys for EP (non-prequantized path) if has_ep_filter: @@ -193,9 +197,7 @@ def _is_excluded_module(self, key: str) -> bool: """Check if a key belongs to a module excluded from quantization.""" return is_excluded_module_key(key, self._exclude_modules) - def on_load_weight( - self, key: str, tensor: torch.Tensor - ) -> List[Tuple[str, torch.Tensor]]: + def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: # 0. Pre-quantized: skip quantized auxiliary keys and quantized weight keys # that weren't caught by get_skip_key_fn (e.g., when skip_key_fn wasn't used) # Excluded modules pass through as bf16 — don't drop them. @@ -266,9 +268,7 @@ def on_load_weight( # 6. Passthrough return [(key, tensor)] - def on_skip_weight( - self, key: str - ) -> List[Tuple[str, torch.Tensor]]: + def on_skip_weight(self, key: str) -> List[Tuple[str, torch.Tensor]]: """Count a skipped expert key so completion tracking stays correct.""" if self._expert_buffer is not None: parsed = parse_expert_key(key) @@ -283,6 +283,7 @@ def on_skip_weight( def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: import warnings + if self._expert_buffer is not None: pending = self._expert_buffer.get_pending_counts() if pending: @@ -297,9 +298,7 @@ def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: warnings.warn(f"Incomplete QKV merge groups after loading: {pending_qkv}") return [] - def on_save_weight( - self, param_name: str, tensor: torch.Tensor - ) -> List[Tuple[str, torch.Tensor]]: + def on_save_weight(self, param_name: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: # Split gate_up_proj -> gate_proj + up_proj (shared expert / dense layers) if ".gate_up_proj." in param_name: prefix, suffix = param_name.rsplit(".gate_up_proj.", 1) diff --git a/src/xorl/models/transformers/qwen3_5_moe/configuration_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/configuration_qwen3_5_moe.py index ec9d36f7..a028859d 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/configuration_qwen3_5_moe.py +++ b/src/xorl/models/transformers/qwen3_5_moe/configuration_qwen3_5_moe.py @@ -14,6 +14,7 @@ """Qwen3_5Moe model configuration""" from transformers.configuration_utils import PretrainedConfig + from xorl.models.layers import rope_config_validation from .parallelize import TP_PLAN @@ -226,7 +227,9 @@ def from_hf_config(cls, hf_config): initializer_range=getattr(text_config, "initializer_range", 0.02), rms_norm_eps=getattr(text_config, "rms_norm_eps", 1e-6), use_cache=getattr(text_config, "use_cache", False), - tie_word_embeddings=getattr(hf_config, "tie_word_embeddings", getattr(text_config, "tie_word_embeddings", False)), + tie_word_embeddings=getattr( + hf_config, "tie_word_embeddings", getattr(text_config, "tie_word_embeddings", False) + ), rope_theta=_cfg_get(rope_params, "rope_theta", getattr(text_config, "rope_theta", 10000.0)), rope_scaling=_cfg_to_dict(rope_params), attention_bias=getattr(text_config, "attention_bias", False), diff --git a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py index 2c19778b..c37b3555 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -4,20 +4,24 @@ import torch from torch import nn +from xorl.distributed.moe.deepep import sync_pending_combine from xorl.distributed.parallel_state import get_parallel_state from xorl.distributed.sequence_parallel.strategy import get_cp_strategy from xorl.models.base import XorlPreTrainedModel +from xorl.models.checkpoint_handlers.buffers import ( + checkpoint_has_per_expert_weights, + detect_prequantized_checkpoint, + get_prequantized_exclude_modules, +) from xorl.models.layers import ACT2FN, RotaryEmbedding from xorl.models.layers.attention import AttentionKwargs, update_causal_mask from xorl.models.layers.attention.backend import ATTENTION_FUNCTIONS from xorl.models.layers.attention.backend.eager import eager_attention_forward from xorl.models.layers.moe import MoEBlock from xorl.models.outputs import MoeCausalLMOutput, MoeModelOutput -from xorl.models.checkpoint_handlers.buffers import ( - checkpoint_has_per_expert_weights, - detect_prequantized_checkpoint, - get_prequantized_exclude_modules, -) +from xorl.models.transformers.qwen3_5_moe import parallelize +from xorl.models.transformers.qwen3_5_moe.checkpoint_handler import Qwen3_5MoeCheckpointHandler +from xorl.models.transformers.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig from xorl.models.transformers.qwen3_5_shared import ( LINEAR_ATTENTION_RING_UNSUPPORTED_MESSAGE, QWEN3_5_CHECKPOINT_CONVERSION_MAPPING, @@ -25,24 +29,21 @@ has_linear_attention_layers, qwen3_5_apply_rotary_pos_emb, ) -from xorl.models.transformers.qwen3_5_moe.checkpoint_handler import Qwen3_5MoeCheckpointHandler -from xorl.models.transformers.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig -from xorl.models.transformers.qwen3_5_moe import parallelize -from xorl.distributed.moe.deepep import sync_pending_combine from xorl.ops.fused_silu_and_mul import fused_silu_and_mul -from xorl.utils import logging from xorl.ops.linear_attention import GatedDeltaNet from xorl.ops.linear_attention.ops.cp import build_linear_attention_cp_context +from xorl.utils import logging + logger = logging.get_logger(__name__) def _adapt_qwen3_5_moe_config(config): - if hasattr(config, 'text_config'): + if hasattr(config, "text_config"): return Qwen3_5MoeConfig.from_hf_config(config) if isinstance(config, Qwen3_5MoeConfig): return config - if getattr(config, 'model_type', None) in {'qwen3_5_moe', 'qwen3_5_moe_text'}: + if getattr(config, "model_type", None) in {"qwen3_5_moe", "qwen3_5_moe_text"}: return Qwen3_5MoeConfig.from_hf_config(config) return config @@ -61,7 +62,7 @@ def __init__(self, config, intermediate_size=None): self.gate_up_proj = nn.Linear(self.hidden_size, 2 * self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[config.hidden_act] - self._use_fused_silu = config.hidden_act == 'silu' + self._use_fused_silu = config.hidden_act == "silu" def unfuse_for_tp(self): device = self.gate_up_proj.weight.device @@ -71,7 +72,7 @@ def unfuse_for_tp(self): del self.gate_up_proj def forward(self, x): - if hasattr(self, 'gate_up_proj'): + if hasattr(self, "gate_up_proj"): if self._use_fused_silu: x = fused_silu_and_mul(self.gate_up_proj(x)) else: @@ -183,13 +184,15 @@ def forward( del position_ids, past_key_values attn_strategy = get_cp_strategy() query_states, key_states, value_states = attn_strategy.project_qkv(self, hidden_states, position_embeddings) - attn_output = attn_strategy.compute_attention(self, query_states, key_states, value_states, attention_mask, **kwargs) + attn_output = attn_strategy.compute_attention( + self, query_states, key_states, value_states, attention_mask, **kwargs + ) attn_output = attn_strategy.project_output(self, attn_output) return attn_output, None class Qwen3_5MoeSparseMoeBlock(MoEBlock): - def __init__(self, config, moe_implementation='triton'): + def __init__(self, config, moe_implementation="triton"): super().__init__( hidden_size=config.hidden_size, num_experts=config.num_experts, @@ -200,10 +203,10 @@ def __init__(self, config, moe_implementation='triton'): moe_implementation=moe_implementation, ) self.config = config - self.experts.ep_dispatch = getattr(config, '_ep_dispatch', 'alltoall') - self.experts.deepep_buffer_size_gb = getattr(config, '_deepep_buffer_size_gb', 2.0) - self.experts.deepep_num_sms = getattr(config, '_deepep_num_sms', 20) - self.experts.deepep_async_combine = getattr(config, '_deepep_async_combine', False) + self.experts.ep_dispatch = getattr(config, "_ep_dispatch", "alltoall") + self.experts.deepep_buffer_size_gb = getattr(config, "_deepep_buffer_size_gb", 2.0) + self.experts.deepep_num_sms = getattr(config, "_deepep_num_sms", 20) + self.experts.deepep_async_combine = getattr(config, "_deepep_async_combine", False) self.shared_expert = Qwen3_5MoeMLP(config, intermediate_size=config.shared_expert_intermediate_size) self.shared_expert_gate = nn.Linear(config.hidden_size, 1, bias=False) @@ -218,10 +221,10 @@ def forward(self, hidden_states: torch.Tensor): QWEN3_5_MOE_CLASSES = { - 'eager': partial(Qwen3_5MoeSparseMoeBlock, moe_implementation='eager'), - 'triton': partial(Qwen3_5MoeSparseMoeBlock, moe_implementation='triton'), - 'native': partial(Qwen3_5MoeSparseMoeBlock, moe_implementation='native'), - 'quack': partial(Qwen3_5MoeSparseMoeBlock, moe_implementation='quack'), + "eager": partial(Qwen3_5MoeSparseMoeBlock, moe_implementation="eager"), + "triton": partial(Qwen3_5MoeSparseMoeBlock, moe_implementation="triton"), + "native": partial(Qwen3_5MoeSparseMoeBlock, moe_implementation="native"), + "quack": partial(Qwen3_5MoeSparseMoeBlock, moe_implementation="quack"), } @@ -229,17 +232,17 @@ class Qwen3_5MoeDecoderLayer(nn.Module): def __init__(self, config: Qwen3_5MoeConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size - self.layer_type = config.layer_types[layer_idx] if layer_idx < len(config.layer_types) else 'full_attention' + self.layer_type = config.layer_types[layer_idx] if layer_idx < len(config.layer_types) else "full_attention" self.self_attn = None self.linear_attn = None - if self.layer_type == 'linear_attention': + if self.layer_type == "linear_attention": self.linear_attn = GatedDeltaNet( hidden_size=config.hidden_size, expand_v=config.linear_value_head_dim / config.linear_key_head_dim, head_dim=config.linear_key_head_dim, num_heads=config.linear_num_key_heads, num_v_heads=config.linear_num_value_heads, - mode='chunk', + mode="chunk", use_gate=config.attn_output_gate, use_short_conv=True, conv_size=config.linear_conv_kernel_dim, @@ -254,7 +257,7 @@ def __init__(self, config: Qwen3_5MoeConfig, layer_idx: int): if (layer_idx not in config.mlp_only_layers) and ( config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0 ): - moe_implementation = getattr(config, '_moe_implementation', 'triton') + moe_implementation = getattr(config, "_moe_implementation", "triton") self.mlp = QWEN3_5_MOE_CLASSES[moe_implementation](config) else: self.mlp = Qwen3_5MoeMLP(config, intermediate_size=config.intermediate_size) @@ -276,8 +279,8 @@ def forward( if self.linear_attn is not None: linear_kwargs = {} - if kwargs.get('cu_seq_lens_q') is not None: - linear_kwargs['cu_seqlens'] = kwargs.get('cu_seq_lens_q') + if kwargs.get("cu_seq_lens_q") is not None: + linear_kwargs["cu_seqlens"] = kwargs.get("cu_seq_lens_q") cp_context = build_linear_attention_cp_context( kwargs.get("cu_seq_lens_q"), conv1d_kernel_size=self.linear_attn.conv_size if self.linear_attn.use_short_conv else None, @@ -325,8 +328,8 @@ def forward( class Qwen3_5MoePreTrainedModel(XorlPreTrainedModel): config_class = Qwen3_5MoeConfig - base_model_prefix = 'model' - _no_split_modules = ['Qwen3_5MoeDecoderLayer'] + base_model_prefix = "model" + _no_split_modules = ["Qwen3_5MoeDecoderLayer"] _checkpoint_conversion_mapping = QWEN3_5_CHECKPOINT_CONVERSION_MAPPING _checkpoint_skip_key_patterns = QWEN3_5_CHECKPOINT_SKIP_KEY_PATTERNS @@ -354,20 +357,20 @@ def get_parallel_plan(self): return parallelize.get_ep_plan() def get_checkpoint_handler(self, **kwargs): - checkpoint_keys = kwargs.get('checkpoint_keys', set()) - weights_path = kwargs.get('weights_path', None) - ep_rank = kwargs.get('ep_rank', 0) - ep_size = kwargs.get('ep_size', 1) - is_broadcast = kwargs.get('is_broadcast', False) + checkpoint_keys = kwargs.get("checkpoint_keys", set()) + weights_path = kwargs.get("weights_path", None) + ep_rank = kwargs.get("ep_rank", 0) + ep_size = kwargs.get("ep_size", 1) + is_broadcast = kwargs.get("is_broadcast", False) has_per_expert = checkpoint_has_per_expert_weights(checkpoint_keys) if checkpoint_keys else True is_prequantized = detect_prequantized_checkpoint(weights_path) - exclude_modules = getattr(self, '_qlora_exclude_modules', None) + exclude_modules = getattr(self, "_qlora_exclude_modules", None) if exclude_modules is None: exclude_modules = get_prequantized_exclude_modules(weights_path) if is_prequantized else set() if is_broadcast: ep_rank, ep_size = 0, 1 - unfused = getattr(self, '_unfused_for_tp', False) - head_dim = getattr(self.config, 'head_dim', self.config.hidden_size // self.config.num_attention_heads) + unfused = getattr(self, "_unfused_for_tp", False) + head_dim = getattr(self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads) return Qwen3_5MoeCheckpointHandler( num_experts=self.config.num_experts, num_attention_heads=self.config.num_attention_heads, @@ -392,7 +395,9 @@ def __init__(self, config: Qwen3_5MoeConfig): self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) - self.layers = nn.ModuleList([Qwen3_5MoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.layers = nn.ModuleList( + [Qwen3_5MoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) self.norm = Qwen3_5MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = RotaryEmbedding(config=config) self.gradient_checkpointing = False @@ -417,11 +422,13 @@ def forward( **kwargs: Unpack[AttentionKwargs], ) -> MoeModelOutput: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_router_logits = output_router_logits if output_router_logits is not None else self.config.output_router_logits + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) if self.embed_tokens is not None: if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError('You must specify exactly one of input_ids or inputs_embeds') + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) hidden_states = inputs_embeds @@ -497,12 +504,14 @@ def forward( all_router_logits += (layer_outputs[-1],) hidden_states = self.norm(hidden_states) if self.norm is not None else hidden_states - return MoeModelOutput(last_hidden_state=hidden_states, attentions=all_self_attns, router_logits=all_router_logits) + return MoeModelOutput( + last_hidden_state=hidden_states, attentions=all_self_attns, router_logits=all_router_logits + ) class Qwen3_5MoeForCausalLM(Qwen3_5MoePreTrainedModel): - _tied_weights_keys = {'lm_head.weight': 'model.embed_tokens.weight'} - _pp_plan = {'lm_head': (['hidden_states'], ['logits'])} + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _tp_plan = parallelize.MODEL_TP_PLAN def __init__(self, config): @@ -539,11 +548,11 @@ def get_decoder(self): def get_pp_module_config(self): return { - 'input_fqns': ['model.embed_tokens'], - 'layer_prefix': 'model.layers', - 'output_fqns': ['model.norm', 'lm_head'], - 'always_keep_fqns': ['model.rotary_emb'], - 'num_layers': self.config.num_hidden_layers, + "input_fqns": ["model.embed_tokens"], + "layer_prefix": "model.layers", + "output_fqns": ["model.norm", "lm_head"], + "always_keep_fqns": ["model.rotary_emb"], + "num_layers": self.config.num_hidden_layers, } def forward( @@ -573,9 +582,9 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3_5MoeForCausalLM): ModelClass = [Qwen3_5MoeForCausalLM, Qwen3_5MoeForConditionalGeneration] __all__ = [ - 'Qwen3_5MoeForCausalLM', - 'Qwen3_5MoeForConditionalGeneration', - 'Qwen3_5MoeModel', - 'Qwen3_5MoePreTrainedModel', - 'Qwen3_5MoeSparseMoeBlock', + "Qwen3_5MoeForCausalLM", + "Qwen3_5MoeForConditionalGeneration", + "Qwen3_5MoeModel", + "Qwen3_5MoePreTrainedModel", + "Qwen3_5MoeSparseMoeBlock", ] diff --git a/src/xorl/models/transformers/qwen3_5_shared.py b/src/xorl/models/transformers/qwen3_5_shared.py index 73c5f070..d93a0860 100644 --- a/src/xorl/models/transformers/qwen3_5_shared.py +++ b/src/xorl/models/transformers/qwen3_5_shared.py @@ -6,6 +6,7 @@ import torch + QWEN3_5_CHECKPOINT_CONVERSION_MAPPING = { r"^model\.language_model\.": "model.", r"^language_model\.": "model.", @@ -17,33 +18,15 @@ r"^mtp\.", ] -_LINEAR_ATTN_QKV_PATTERN = re.compile( - r"^model\.layers\.(\d+)\.linear_attn\.in_proj_qkv\.weight$" -) -_LINEAR_ATTN_Z_PATTERN = re.compile( - r"^model\.layers\.(\d+)\.linear_attn\.in_proj_z\.weight$" -) -_LINEAR_ATTN_B_PATTERN = re.compile( - r"^model\.layers\.(\d+)\.linear_attn\.in_proj_b\.weight$" -) -_LINEAR_ATTN_A_PATTERN = re.compile( - r"^model\.layers\.(\d+)\.linear_attn\.in_proj_a\.weight$" -) -_LINEAR_ATTN_CONV_PATTERN = re.compile( - r"^model\.layers\.(\d+)\.linear_attn\.conv1d\.weight$" -) -_LINEAR_ATTN_OUT_PATTERN = re.compile( - r"^model\.layers\.(\d+)\.linear_attn\.out_proj\.weight$" -) -_LINEAR_ATTN_NORM_PATTERN = re.compile( - r"^model\.layers\.(\d+)\.linear_attn\.norm\.weight$" -) -_LINEAR_ATTN_DT_PATTERN = re.compile( - r"^model\.layers\.(\d+)\.linear_attn\.dt_bias$" -) -_LINEAR_ATTN_A_LOG_PATTERN = re.compile( - r"^model\.layers\.(\d+)\.linear_attn\.A_log$" -) +_LINEAR_ATTN_QKV_PATTERN = re.compile(r"^model\.layers\.(\d+)\.linear_attn\.in_proj_qkv\.weight$") +_LINEAR_ATTN_Z_PATTERN = re.compile(r"^model\.layers\.(\d+)\.linear_attn\.in_proj_z\.weight$") +_LINEAR_ATTN_B_PATTERN = re.compile(r"^model\.layers\.(\d+)\.linear_attn\.in_proj_b\.weight$") +_LINEAR_ATTN_A_PATTERN = re.compile(r"^model\.layers\.(\d+)\.linear_attn\.in_proj_a\.weight$") +_LINEAR_ATTN_CONV_PATTERN = re.compile(r"^model\.layers\.(\d+)\.linear_attn\.conv1d\.weight$") +_LINEAR_ATTN_OUT_PATTERN = re.compile(r"^model\.layers\.(\d+)\.linear_attn\.out_proj\.weight$") +_LINEAR_ATTN_NORM_PATTERN = re.compile(r"^model\.layers\.(\d+)\.linear_attn\.norm\.weight$") +_LINEAR_ATTN_DT_PATTERN = re.compile(r"^model\.layers\.(\d+)\.linear_attn\.dt_bias$") +_LINEAR_ATTN_A_LOG_PATTERN = re.compile(r"^model\.layers\.(\d+)\.linear_attn\.A_log$") LINEAR_ATTENTION_RING_UNSUPPORTED_MESSAGE = ( "Native FLA CP for Qwen3.5 linear_attention currently supports Ulysses-only CP only; " @@ -91,6 +74,63 @@ def has_linear_attention_layers(config: object) -> bool: return any(layer_type == "linear_attention" for layer_type in getattr(config, "layer_types", [])) +_LINEAR_ATTN_SPLIT_PATTERN = re.compile( + r"^(model\.layers\.(\d+)\.linear_attn)\.(q_proj|k_proj|v_proj|g_proj|a_proj|b_proj|" + r"q_conv1d|k_conv1d|v_conv1d|o_proj|o_norm|dt_bias|A_log)\.(weight|bias)$" +) +_LINEAR_ATTN_SPLIT_PATTERN_NO_SUFFIX = re.compile(r"^(model\.layers\.(\d+)\.linear_attn)\.(dt_bias|A_log)$") + +_SPLIT_TO_HF_RENAME = { + "g_proj": "in_proj_z", + "a_proj": "in_proj_a", + "b_proj": "in_proj_b", + "o_proj": "out_proj", + "o_norm": "norm", +} + +_SPLIT_QKV_PARTS = {"q_proj", "k_proj", "v_proj"} +_SPLIT_CONV_PARTS = {"q_conv1d", "k_conv1d", "v_conv1d"} + + +def remap_linear_attention_params_for_inference( + buffer: list[tuple[str, "torch.Tensor"]], +) -> list[tuple[str, "torch.Tensor"]]: + fuse_groups: dict[str, dict[str, "torch.Tensor"]] = {} + result: list[tuple[str, "torch.Tensor"]] = [] + + for name, tensor in buffer: + m = _LINEAR_ATTN_SPLIT_PATTERN.match(name) + if m is None: + m = _LINEAR_ATTN_SPLIT_PATTERN_NO_SUFFIX.match(name) + if m is None: + result.append((name, tensor)) + continue + + prefix = m.group(1) + proj = m.group(3) + rest = name[m.end(3) :] + + if proj in _SPLIT_QKV_PARTS: + key = f"{prefix}.in_proj_qkv{rest}" + fuse_groups.setdefault(key, {})[proj] = tensor + elif proj in _SPLIT_CONV_PARTS: + key = f"{prefix}.conv1d{rest}" + fuse_groups.setdefault(key, {})[proj] = tensor + elif proj in _SPLIT_TO_HF_RENAME: + result.append((f"{prefix}.{_SPLIT_TO_HF_RENAME[proj]}{rest}", tensor)) + else: + result.append((f"{prefix}.{proj}{rest}", tensor)) + + for fused_name, parts in fuse_groups.items(): + if "q_proj" in parts: + ordered = [parts["q_proj"], parts["k_proj"], parts["v_proj"]] + else: + ordered = [parts["q_conv1d"], parts["k_conv1d"], parts["v_conv1d"]] + result.append((fused_name, torch.cat(ordered, dim=0))) + + return result + + def map_qwen3_5_linear_attention_weight( key: str, tensor: torch.Tensor, diff --git a/src/xorl/models/transformers/qwen3_moe/__init__.py b/src/xorl/models/transformers/qwen3_moe/__init__.py index 2ad67e52..8d4218f8 100644 --- a/src/xorl/models/transformers/qwen3_moe/__init__.py +++ b/src/xorl/models/transformers/qwen3_moe/__init__.py @@ -1,24 +1,24 @@ """Qwen3 MoE model.""" -from .configuration_qwen3_moe import Qwen3MoeConfig -from .modeling_qwen3_moe import ( - Qwen3MoeSparseExperts, - Qwen3MoeTritonExperts, - Qwen3MoeForCausalLM, - Qwen3MoeModel, - Qwen3MoeSparseTritonMoeBlock, - Qwen3MoeSparseNativeMoeBlock, -) - # Canonical MoE layer abstractions from ...layers.moe import ( + MOE_EXPERT_BACKENDS, MoEBlock, MoEExperts, MoEExpertsLoRA, MoELoRAConfig, TopKRouter, - MOE_EXPERT_BACKENDS, ) +from .configuration_qwen3_moe import Qwen3MoeConfig +from .modeling_qwen3_moe import ( + Qwen3MoeForCausalLM, + Qwen3MoeModel, + Qwen3MoeSparseExperts, + Qwen3MoeSparseNativeMoeBlock, + Qwen3MoeSparseTritonMoeBlock, + Qwen3MoeTritonExperts, +) + __all__ = [ # Base model diff --git a/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py b/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py index 63f55e5f..57ea4c6c 100644 --- a/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py @@ -7,13 +7,20 @@ from ...checkpoint_handlers.base import CheckpointHandler from ...checkpoint_handlers.buffers import ( - ExpertWeightBuffer, GateUpMergeBuffer, QKVMergeBuffer, - QLoRAWeightBuffer, QLoRAExpertBuffer, - parse_expert_key, parse_expert_full_key, - EXPERT_QUANT_AUX_PATTERN, QUANT_AUX_SUFFIX_PATTERN, + DENSE_DOWN_PROJ_PATTERN, + DENSE_GATE_UP_PATTERN, + EXPERT_QUANT_AUX_PATTERN, FP8_AUX_SUFFIX_PATTERN, - QKV_PROJ_PATTERN, DENSE_GATE_UP_PATTERN, - OPROJ_WEIGHT_PATTERN, DENSE_DOWN_PROJ_PATTERN, + OPROJ_WEIGHT_PATTERN, + QKV_PROJ_PATTERN, + QUANT_AUX_SUFFIX_PATTERN, + ExpertWeightBuffer, + GateUpMergeBuffer, + QKVMergeBuffer, + QLoRAExpertBuffer, + QLoRAWeightBuffer, + parse_expert_full_key, + parse_expert_key, ) @@ -63,7 +70,10 @@ def __init__( # Non-QLoRA expert stacking (disabled when pre-quantized — use QLoRAExpertBuffer instead) if checkpoint_has_per_expert and not is_prequantized: self._expert_buffer = ExpertWeightBuffer( - num_experts, ep_rank=ep_rank, ep_size=ep_size, device=device, + num_experts, + ep_rank=ep_rank, + ep_size=ep_size, + device=device, ) self._qkv_buffer: Optional[QKVMergeBuffer] = None if not skip_qkv_merge: @@ -83,7 +93,10 @@ def __init__( self._qlora_expert_buffer = None if is_prequantized and model is not None: self._qlora_expert_buffer = QLoRAExpertBuffer( - model, ep_rank=ep_rank, ep_size=ep_size, num_experts=num_experts, + model, + ep_rank=ep_rank, + ep_size=ep_size, + num_experts=num_experts, ) def get_skip_key_fn(self) -> Optional[Callable[[str], bool]]: @@ -96,15 +109,12 @@ def get_skip_key_fn(self) -> Optional[Callable[[str], bool]]: When inline buffers are active, in-range keys flow through on_load_weight. """ - has_ep_filter = ( - self._expert_buffer is not None - and not (self._expert_buffer.expert_start == 0 - and self._expert_buffer.expert_end == self._expert_buffer.num_experts) + has_ep_filter = self._expert_buffer is not None and not ( + self._expert_buffer.expert_start == 0 and self._expert_buffer.expert_end == self._expert_buffer.num_experts ) - has_expert_ep_filter = ( - self._qlora_expert_buffer is not None - and not (self._qlora_expert_buffer.expert_start == 0 - and self._qlora_expert_buffer.expert_end == self._qlora_expert_buffer._num_experts) + has_expert_ep_filter = self._qlora_expert_buffer is not None and not ( + self._qlora_expert_buffer.expert_start == 0 + and self._qlora_expert_buffer.expert_end == self._qlora_expert_buffer._num_experts ) if not has_ep_filter and not has_expert_ep_filter and not self._is_prequantized: @@ -152,8 +162,12 @@ def _should_skip(key: str) -> bool: if FP8_AUX_SUFFIX_PATTERN.search(key): return True if key.endswith(".weight"): - if (QKV_PROJ_PATTERN.match(key) or DENSE_GATE_UP_PATTERN.match(key) - or OPROJ_WEIGHT_PATTERN.match(key) or DENSE_DOWN_PROJ_PATTERN.match(key)): + if ( + QKV_PROJ_PATTERN.match(key) + or DENSE_GATE_UP_PATTERN.match(key) + or OPROJ_WEIGHT_PATTERN.match(key) + or DENSE_DOWN_PROJ_PATTERN.match(key) + ): return True # Skip out-of-range expert keys for EP (non-prequantized path) @@ -174,9 +188,7 @@ def _is_excluded_module(self, key: str) -> bool: module_short_name = module_fqn.rsplit(".", 1)[-1] return module_short_name in self._exclude_modules - def on_load_weight( - self, key: str, tensor: torch.Tensor - ) -> List[Tuple[str, torch.Tensor]]: + def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: # Drop input_scale (unused by our quantization) if key.endswith(".input_scale"): return [] @@ -251,9 +263,7 @@ def on_load_weight( # 4. Passthrough return [(key, tensor)] - def on_skip_weight( - self, key: str - ) -> List[Tuple[str, torch.Tensor]]: + def on_skip_weight(self, key: str) -> List[Tuple[str, torch.Tensor]]: """Count a skipped expert key so completion tracking stays correct.""" # QLoRA expert buffer: count skipped out-of-range expert keys if self._qlora_expert_buffer is not None: @@ -272,6 +282,7 @@ def on_skip_weight( def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: import warnings + if self._expert_buffer is not None: pending = self._expert_buffer.get_pending_counts() if pending: @@ -291,15 +302,12 @@ def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: pending_exp = self._qlora_expert_buffer.get_pending() if pending_exp: warnings.warn( - f"Incomplete QLoRA expert weights after loading " - f"(will fall back to deferred loading): {pending_exp}" + f"Incomplete QLoRA expert weights after loading (will fall back to deferred loading): {pending_exp}" ) self._qlora_expert_buffer.set_inline_metadata() return [] - def on_save_weight( - self, param_name: str, tensor: torch.Tensor - ) -> List[Tuple[str, torch.Tensor]]: + def on_save_weight(self, param_name: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: # Split gate_up_proj -> gate_proj + up_proj (shared expert / dense layers) if ".gate_up_proj." in param_name: prefix, suffix = param_name.rsplit(".gate_up_proj.", 1) diff --git a/src/xorl/models/transformers/qwen3_moe/configuration_qwen3_moe.py b/src/xorl/models/transformers/qwen3_moe/configuration_qwen3_moe.py index fa638c25..62778afc 100644 --- a/src/xorl/models/transformers/qwen3_moe/configuration_qwen3_moe.py +++ b/src/xorl/models/transformers/qwen3_moe/configuration_qwen3_moe.py @@ -14,6 +14,7 @@ """Qwen3Moe model configuration""" from transformers.configuration_utils import PretrainedConfig + from xorl.models.layers import rope_config_validation from ....utils import logging @@ -125,7 +126,7 @@ class Qwen3MoeConfig(PretrainedConfig): norm_topk_prob (`bool`, *optional*, defaults to `False`): Whether to normalize the topk probabilities. output_router_logits (`bool`, *optional*, defaults to `False`): - Whether or not the router logits should be returned by the model. Enabeling this will also + Whether or not the router logits should be returned by the model. Enabling this will also allow the model to output the auxiliary loss, including load balancing loss and router z-loss. router_aux_loss_coef (`float`, *optional*, defaults to 0.001): The aux loss factor for the total loss. diff --git a/src/xorl/ops/__init__.py b/src/xorl/ops/__init__.py index e6ca3393..6b945ea4 100644 --- a/src/xorl/ops/__init__.py +++ b/src/xorl/ops/__init__.py @@ -1,17 +1,17 @@ -from .moe.triton import TritonMoeExpertsFunction, triton_moe_forward -from .moe.triton_lora import ( - TritonMoeExpertsLoRAFunction, - triton_moe_lora_forward, +from .linear_attention import ( + GatedDeltaNet, + chunk_gated_delta_rule, + fused_recurrent_gated_delta_rule, ) -from .moe.quack import quack_moe_forward from .loss import ( causallm_loss_function, importance_sampling_loss_function, ) -from .linear_attention import ( - GatedDeltaNet, - chunk_gated_delta_rule, - fused_recurrent_gated_delta_rule, +from .moe.quack import quack_moe_forward +from .moe.triton import TritonMoeExpertsFunction, triton_moe_forward +from .moe.triton_lora import ( + TritonMoeExpertsLoRAFunction, + triton_moe_lora_forward, ) diff --git a/src/xorl/ops/ep_kernels/__init__.py b/src/xorl/ops/ep_kernels/__init__.py index b453fb75..ba136a32 100644 --- a/src/xorl/ops/ep_kernels/__init__.py +++ b/src/xorl/ops/ep_kernels/__init__.py @@ -1,14 +1,14 @@ +from .deepep_counting_sort import ( + build_sorted_indices, + group_tokens_by_expert, + group_tokens_by_expert_v2, +) from .deepep_scatter_gather import ( DeepEPScatter, DeepEPWeightedGather, deepep_scatter, deepep_weighted_gather, ) -from .deepep_counting_sort import ( - build_sorted_indices, - group_tokens_by_expert, - group_tokens_by_expert_v2, -) __all__ = [ diff --git a/src/xorl/ops/ep_kernels/deepep_scatter_gather.py b/src/xorl/ops/ep_kernels/deepep_scatter_gather.py index fbfce8b4..9c50550b 100644 --- a/src/xorl/ops/ep_kernels/deepep_scatter_gather.py +++ b/src/xorl/ops/ep_kernels/deepep_scatter_gather.py @@ -63,9 +63,7 @@ def forward(ctx, tokens, sorted_indices): return output BLOCK_H = min(triton.next_power_of_2(hidden_dim), 1024) grid = (num_gather,) - _gather_by_index_kernel[grid]( - output, tokens, sorted_indices, num_gather, hidden_dim, BLOCK_H=BLOCK_H - ) + _gather_by_index_kernel[grid](output, tokens, sorted_indices, num_gather, hidden_dim, BLOCK_H=BLOCK_H) ctx.save_for_backward(sorted_indices) ctx.num_tokens = tokens.shape[0] ctx.hidden_dim = hidden_dim diff --git a/src/xorl/ops/group_gemm/kernel/__init__.py b/src/xorl/ops/group_gemm/kernel/__init__.py index b392155f..c4fe4487 100644 --- a/src/xorl/ops/group_gemm/kernel/__init__.py +++ b/src/xorl/ops/group_gemm/kernel/__init__.py @@ -1,6 +1,14 @@ # Group GEMM kernels from .group_gemm import group_gemm_same_mn, group_gemm_same_nk -from .quack import quack_group_gemm_same_mn, quack_group_gemm_same_nk + +# LoRA utilities +from .lora_utils import ( + compute_lora_scaling, + get_lora_delta_weight_stacked, + init_lora_weights_stacked, + merge_lora_weights_stacked, + unmerge_lora_weights_stacked, +) # MoE operations from .moe import ( @@ -10,15 +18,8 @@ moe_index_compute, moe_scatter, ) +from .quack import quack_group_gemm_same_mn, quack_group_gemm_same_nk -# LoRA utilities -from .lora_utils import ( - init_lora_weights_stacked, - compute_lora_scaling, - merge_lora_weights_stacked, - unmerge_lora_weights_stacked, - get_lora_delta_weight_stacked, -) __all__ = [ # Group GEMM diff --git a/src/xorl/ops/group_gemm/kernel/group_gemm.py b/src/xorl/ops/group_gemm/kernel/group_gemm.py index a235bd7e..aaed7cef 100644 --- a/src/xorl/ops/group_gemm/kernel/group_gemm.py +++ b/src/xorl/ops/group_gemm/kernel/group_gemm.py @@ -25,7 +25,6 @@ get_pid_mn, make_blocked, ) -from ..utils.pretuned import algo_key_scaled, pretuned def _get_cuda_autotune_config(): @@ -47,6 +46,7 @@ def _get_cuda_autotune_config(): ), ] + @triton.autotune( configs=_get_cuda_autotune_config(), key=["N", "K"], diff --git a/src/xorl/ops/group_gemm/kernel/lora_utils.py b/src/xorl/ops/group_gemm/kernel/lora_utils.py index 2f706907..b6c6a354 100644 --- a/src/xorl/ops/group_gemm/kernel/lora_utils.py +++ b/src/xorl/ops/group_gemm/kernel/lora_utils.py @@ -5,7 +5,7 @@ """ import math -from typing import Tuple, Optional +from typing import Optional, Tuple import torch import torch.nn as nn diff --git a/src/xorl/ops/group_gemm/kernel/moe.py b/src/xorl/ops/group_gemm/kernel/moe.py index d4fcb2d9..a0d19fa5 100644 --- a/src/xorl/ops/group_gemm/kernel/moe.py +++ b/src/xorl/ops/group_gemm/kernel/moe.py @@ -347,7 +347,7 @@ def _moe_index_compute_kernel( slot_ids = ( # Reserve last `slots_to_reserve` slots for us. tl.atomic_add(temp_histogram_cumsum_ptr + expert_id, -slots_to_reserve, sem="relaxed") - # `atomic_add` returns old value, so we need to do substraction again. + # `atomic_add` returns old value, so we need to do subtraction again. - slots_to_reserve # Local offset for each token in `expert_ids`. + tl.cumsum(one_if_expert_id_matches) diff --git a/src/xorl/ops/group_gemm/kernel/quack.py b/src/xorl/ops/group_gemm/kernel/quack.py index 011dc912..66283257 100644 --- a/src/xorl/ops/group_gemm/kernel/quack.py +++ b/src/xorl/ops/group_gemm/kernel/quack.py @@ -18,10 +18,11 @@ kernels to the Xorl group GEMM API. Requires SM90+ (H100) or SM100+ (B200). """ +import os from typing import Optional -import os import torch + from xorl.ops.quack.gemm_interface import gemm as quack_gemm diff --git a/src/xorl/ops/linear_attention/__init__.py b/src/xorl/ops/linear_attention/__init__.py index 3acabf21..a10dcd3a 100644 --- a/src/xorl/ops/linear_attention/__init__.py +++ b/src/xorl/ops/linear_attention/__init__.py @@ -3,6 +3,7 @@ from .layers.gated_deltanet import GatedDeltaNet from .ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule + __all__ = [ "GatedDeltaNet", "chunk_gated_delta_rule", diff --git a/src/xorl/ops/linear_attention/layers/__init__.py b/src/xorl/ops/linear_attention/layers/__init__.py index fb4fb7ab..b91b02dd 100644 --- a/src/xorl/ops/linear_attention/layers/__init__.py +++ b/src/xorl/ops/linear_attention/layers/__init__.py @@ -1,3 +1,4 @@ from .gated_deltanet import GatedDeltaNet + __all__ = ["GatedDeltaNet"] diff --git a/src/xorl/ops/linear_attention/layers/gated_deltanet.py b/src/xorl/ops/linear_attention/layers/gated_deltanet.py index 1cea66af..2e832db2 100644 --- a/src/xorl/ops/linear_attention/layers/gated_deltanet.py +++ b/src/xorl/ops/linear_attention/layers/gated_deltanet.py @@ -2,7 +2,6 @@ # Adapted from flash-linear-attention/fla/layers/gated_deltanet.py. # Portions of this file are adapted from flash-linear-attention, Copyright (c) 2023-2025 Songlin Yang, licensed under the MIT License. - import math import warnings from typing import Any @@ -144,7 +143,11 @@ def forward( batch_size, q_len, _ = hidden_states.shape cp_context = kwargs.get("cp_context") - mode = self.mode if cp_context is not None else ("fused_recurrent" if (q_len <= 64 and not self.training) else self.mode) + mode = ( + self.mode + if cp_context is not None + else ("fused_recurrent" if (q_len <= 64 and not self.training) else self.mode) + ) if self.training and mode != "chunk": raise AssertionError("Only chunk mode is supported in training.") diff --git a/src/xorl/ops/linear_attention/layers/utils.py b/src/xorl/ops/linear_attention/layers/utils.py index 659e2e18..a7016142 100644 --- a/src/xorl/ops/linear_attention/layers/utils.py +++ b/src/xorl/ops/linear_attention/layers/utils.py @@ -2,7 +2,6 @@ # Adapted from flash-linear-attention/fla/layers/utils.py. # Portions of this file are adapted from flash-linear-attention, Copyright (c) 2023-2025 Songlin Yang, licensed under the MIT License. - import torch from einops import rearrange, repeat diff --git a/src/xorl/ops/linear_attention/modules/__init__.py b/src/xorl/ops/linear_attention/modules/__init__.py index d6bfebaa..8dbea46c 100644 --- a/src/xorl/ops/linear_attention/modules/__init__.py +++ b/src/xorl/ops/linear_attention/modules/__init__.py @@ -2,6 +2,7 @@ from .layernorm import RMSNorm from .short_conv import ShortConvolution + __all__ = [ "FusedRMSNormGated", "RMSNorm", diff --git a/src/xorl/ops/linear_attention/modules/fused_norm_gate.py b/src/xorl/ops/linear_attention/modules/fused_norm_gate.py index 3e11352c..1ba9920e 100644 --- a/src/xorl/ops/linear_attention/modules/fused_norm_gate.py +++ b/src/xorl/ops/linear_attention/modules/fused_norm_gate.py @@ -2,7 +2,6 @@ # Minimal local RMSNorm-gate path adapted from flash-linear-attention/fla/modules/fused_norm_gate.py. # Portions of this file are adapted from flash-linear-attention, Copyright (c) 2023-2025 Songlin Yang, licensed under the MIT License. - import torch import torch.nn as nn diff --git a/src/xorl/ops/linear_attention/modules/l2norm.py b/src/xorl/ops/linear_attention/modules/l2norm.py index 4a4371f0..ea5ed03d 100644 --- a/src/xorl/ops/linear_attention/modules/l2norm.py +++ b/src/xorl/ops/linear_attention/modules/l2norm.py @@ -7,6 +7,7 @@ from xorl.ops.linear_attention.utils import IS_AMD, autotune_cache_kwargs, input_guard + BT_LIST = [8, 16, 32, 64, 128] NUM_WARPS_AUTOTUNE = [1, 2, 4, 8, 16] if IS_AMD else [1, 2, 4, 8, 16, 32] diff --git a/src/xorl/ops/linear_attention/modules/layernorm.py b/src/xorl/ops/linear_attention/modules/layernorm.py index f33f362c..5155eb8a 100644 --- a/src/xorl/ops/linear_attention/modules/layernorm.py +++ b/src/xorl/ops/linear_attention/modules/layernorm.py @@ -2,7 +2,6 @@ # Minimal local RMSNorm adapted from flash-linear-attention/fla/modules/layernorm.py. # Portions of this file are adapted from flash-linear-attention, Copyright (c) 2023-2025 Songlin Yang, licensed under the MIT License. - import torch import torch.nn as nn diff --git a/src/xorl/ops/linear_attention/ops/__init__.py b/src/xorl/ops/linear_attention/ops/__init__.py index 4c4e5040..ebe4f085 100644 --- a/src/xorl/ops/linear_attention/ops/__init__.py +++ b/src/xorl/ops/linear_attention/ops/__init__.py @@ -1,3 +1,4 @@ from .gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule + __all__ = ["chunk_gated_delta_rule", "fused_recurrent_gated_delta_rule"] diff --git a/src/xorl/ops/linear_attention/ops/common/chunk_delta_h.py b/src/xorl/ops/linear_attention/ops/common/chunk_delta_h.py index 06dba999..fcad84f5 100644 --- a/src/xorl/ops/linear_attention/ops/common/chunk_delta_h.py +++ b/src/xorl/ops/linear_attention/ops/common/chunk_delta_h.py @@ -9,29 +9,32 @@ from xorl.ops.linear_attention.ops.utils.op import exp, exp2 from xorl.ops.linear_attention.utils import IS_NVIDIA_HOPPER, USE_CUDA_GRAPH, autotune_cache_kwargs, check_shared_mem + NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_GK': lambda args: args['gk'] is not None, - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'SAVE_NEW_VALUE': lambda args: args['v_new'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "USE_INITIAL_STATE": lambda args: args["h0"] is not None, + "STORE_FINAL_STATE": lambda args: args["ht"] is not None, + "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ - triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) + triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4] - for num_stages in ([4, 3, 2] if check_shared_mem('ampere') else [2, 1]) - for BV in ([32, 64] if check_shared_mem('ada') else [32]) + for num_stages in ([4, 3, 2] if check_shared_mem("ampere") else [2, 1]) + for BV in ([32, 64] if check_shared_mem("ada") else [32]) ], - key=['H', 'K', 'V', 'BT', 'USE_EXP2'], + key=["H", "K", "V", "BT", "USE_EXP2"], use_cuda_graph=USE_CUDA_GRAPH, **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( k, v, @@ -80,7 +83,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( b_h4 = tl.zeros([64, BV], dtype=tl.float32) # calculate offset - h += (boh * H + i_h).to(tl.int64) * K*V + h += (boh * H + i_h).to(tl.int64) * K * V v += (bos * H + i_h).to(tl.int64) * V k += (bos * H + i_h).to(tl.int64) * K w += (bos * H + i_h).to(tl.int64) * K @@ -88,9 +91,9 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( v_new += (bos * H + i_h).to(tl.int64) * V if USE_INITIAL_STATE: - h0 = h0 + i_nh * K*V + h0 = h0 + i_nh * K * V if STORE_FINAL_STATE: - ht = ht + i_nh * K*V + ht = ht + i_nh * K * V # load initial state if USE_INITIAL_STATE: @@ -109,38 +112,38 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( # main recurrence for i_t in range(NT): i_t_int64 = i_t.to(tl.int64) - p_h1 = tl.make_block_ptr(h + i_t_int64 * H*K*V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + p_h1 = tl.make_block_ptr(h + i_t_int64 * H * K * V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) if K > 64: - p_h2 = tl.make_block_ptr(h + i_t_int64 * H*K*V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + p_h2 = tl.make_block_ptr(h + i_t_int64 * H * K * V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) if K > 128: - p_h3 = tl.make_block_ptr(h + i_t_int64 * H*K*V, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + p_h3 = tl.make_block_ptr(h + i_t_int64 * H * K * V, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) if K > 192: - p_h4 = tl.make_block_ptr(h + i_t_int64 * H*K*V, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + p_h4 = tl.make_block_ptr(h + i_t_int64 * H * K * V, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) - p_w = tl.make_block_ptr(w, (T, K), (H*K, 1), (i_t * BT, 0), (BT, 64), (1, 0)) + p_w = tl.make_block_ptr(w, (T, K), (H * K, 1), (i_t * BT, 0), (BT, 64), (1, 0)) b_w = tl.load(p_w, boundary_check=(0, 1)) b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) if K > 64: - p_w = tl.make_block_ptr(w, (T, K), (H*K, 1), (i_t * BT, 64), (BT, 64), (1, 0)) + p_w = tl.make_block_ptr(w, (T, K), (H * K, 1), (i_t * BT, 64), (BT, 64), (1, 0)) b_w = tl.load(p_w, boundary_check=(0, 1)) b_v += tl.dot(b_w, b_h2.to(b_w.dtype)) if K > 128: - p_w = tl.make_block_ptr(w, (T, K), (H*K, 1), (i_t * BT, 128), (BT, 64), (1, 0)) + p_w = tl.make_block_ptr(w, (T, K), (H * K, 1), (i_t * BT, 128), (BT, 64), (1, 0)) b_w = tl.load(p_w, boundary_check=(0, 1)) b_v += tl.dot(b_w, b_h3.to(b_w.dtype)) if K > 192: - p_w = tl.make_block_ptr(w, (T, K), (H*K, 1), (i_t * BT, 192), (BT, 64), (1, 0)) + p_w = tl.make_block_ptr(w, (T, K), (H * K, 1), (i_t * BT, 192), (BT, 64), (1, 0)) b_w = tl.load(p_w, boundary_check=(0, 1)) b_v += tl.dot(b_w, b_h4.to(b_w.dtype)) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v = tl.make_block_ptr(v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v if SAVE_NEW_VALUE: - p_v = tl.make_block_ptr(v_new, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v = tl.make_block_ptr(v_new, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) tl.store(p_v, b_v.to(p_v.dtype.element_ty), boundary_check=(0, 1)) last_idx = min((i_t + 1) * BT, T) - 1 @@ -165,28 +168,36 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( if USE_GK: o_k1 = tl.arange(0, 64) - b_gk_last1 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) + b_gk_last1 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k1, mask=(o_k1 < K), other=0.0).to( + tl.float32 + ) if USE_EXP2: b_h1 *= exp2(b_gk_last1)[:, None] else: b_h1 *= exp(b_gk_last1)[:, None] if K > 64: o_k2 = 64 + o_k1 - b_gk_last2 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) + b_gk_last2 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k2, mask=(o_k2 < K), other=0.0).to( + tl.float32 + ) if USE_EXP2: b_h2 *= exp2(b_gk_last2)[:, None] else: b_h2 *= exp(b_gk_last2)[:, None] if K > 128: o_k3 = 128 + o_k1 - b_gk_last3 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) + b_gk_last3 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k3, mask=(o_k3 < K), other=0.0).to( + tl.float32 + ) if USE_EXP2: b_h3 *= exp2(b_gk_last3)[:, None] else: b_h3 *= exp(b_gk_last3)[:, None] if K > 192: o_k4 = 192 + o_k1 - b_gk_last4 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) + b_gk_last4 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k4, mask=(o_k4 < K), other=0.0).to( + tl.float32 + ) if USE_EXP2: b_h4 *= exp2(b_gk_last4)[:, None] else: @@ -194,19 +205,19 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( b_v = b_v.to(k.dtype.element_ty) - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_t * BT), (64, BT), (0, 1)) + p_k = tl.make_block_ptr(k, (K, T), (1, H * K), (0, i_t * BT), (64, BT), (0, 1)) b_k = tl.load(p_k, boundary_check=(0, 1)) b_h1 += tl.dot(b_k, b_v) if K > 64: - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (64, i_t * BT), (64, BT), (0, 1)) + p_k = tl.make_block_ptr(k, (K, T), (1, H * K), (64, i_t * BT), (64, BT), (0, 1)) b_k = tl.load(p_k, boundary_check=(0, 1)) b_h2 += tl.dot(b_k, b_v) if K > 128: - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (128, i_t * BT), (64, BT), (0, 1)) + p_k = tl.make_block_ptr(k, (K, T), (1, H * K), (128, i_t * BT), (64, BT), (0, 1)) b_k = tl.load(p_k, boundary_check=(0, 1)) b_h3 += tl.dot(b_k, b_v) if K > 192: - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (192, i_t * BT), (64, BT), (0, 1)) + p_k = tl.make_block_ptr(k, (K, T), (1, H * K), (192, i_t * BT), (64, BT), (0, 1)) b_k = tl.load(p_k, boundary_check=(0, 1)) b_h4 += tl.dot(b_k, b_v) @@ -224,25 +235,27 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_GK': lambda args: args['gk'] is not None, - 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, - 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "USE_INITIAL_STATE": lambda args: args["dh0"] is not None, + "USE_FINAL_STATE_GRADIENT": lambda args: args["dht"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ - triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) + triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4] - for num_stages in ([4, 3, 2] if check_shared_mem('ampere') else [1]) - for BV in ([64, 32] if check_shared_mem('ada') else [32]) + for num_stages in ([4, 3, 2] if check_shared_mem("ampere") else [1]) + for BV in ([64, 32] if check_shared_mem("ada") else [32]) ], - key=['H', 'K', 'V', 'BT', 'BV', 'USE_G', 'USE_EXP2'], + key=["H", "K", "V", "BT", "BV", "USE_G", "USE_EXP2"], use_cuda_graph=USE_CUDA_GRAPH, **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( q, k, @@ -299,14 +312,14 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( do += (bos * H + i_h).to(tl.int64) * V dv += (bos * H + i_h).to(tl.int64) * V dv2 += (bos * H + i_h).to(tl.int64) * V - dh += (boh * H + i_h).to(tl.int64) * K*V + dh += (boh * H + i_h).to(tl.int64) * K * V if USE_GK: gk += (bos * H + i_h).to(tl.int64) * K if USE_INITIAL_STATE: - dh0 += i_nh * K*V + dh0 += i_nh * K * V if USE_FINAL_STATE_GRADIENT: - dht += i_nh * K*V + dht += i_nh * K * V if USE_FINAL_STATE_GRADIENT: p_dht1 = tl.make_block_ptr(dht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) @@ -323,16 +336,16 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( for i_t in range(NT - 1, -1, -1): i_t_int64 = i_t.to(tl.int64) - p_dh1 = tl.make_block_ptr(dh + i_t_int64*H*K*V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + p_dh1 = tl.make_block_ptr(dh + i_t_int64 * H * K * V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) if K > 64: - p_dh2 = tl.make_block_ptr(dh + i_t_int64*H*K*V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + p_dh2 = tl.make_block_ptr(dh + i_t_int64 * H * K * V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) if K > 128: - p_dh3 = tl.make_block_ptr(dh + i_t_int64*H*K*V, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + p_dh3 = tl.make_block_ptr(dh + i_t_int64 * H * K * V, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) if K > 192: - p_dh4 = tl.make_block_ptr(dh + i_t_int64*H*K*V, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + p_dh4 = tl.make_block_ptr(dh + i_t_int64 * H * K * V, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) last_idx = min((i_t + 1) * BT, T) - 1 @@ -347,42 +360,42 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( bg_last_exp = exp(bg_last) b_g_exp = exp(b_g) - p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv2 = tl.make_block_ptr(dv2, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv2 = tl.make_block_ptr(dv2, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) b_do = tl.load(p_do, boundary_check=(0, 1)) # Update dv - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, 64), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT, 0), (BT, 64), (1, 0)) b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_GK: o_k1 = tl.arange(0, 64) - b_gk_last1 = tl.load(gk + last_idx * H*K + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) + b_gk_last1 = tl.load(gk + last_idx * H * K + o_k1, mask=(o_k1 < K), other=0.0).to(tl.float32) b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) if K > 64: - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 64), (BT, 64), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT, 64), (BT, 64), (1, 0)) b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_GK: o_k2 = 64 + o_k1 - b_gk_last2 = tl.load(gk + last_idx * H*K + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) + b_gk_last2 = tl.load(gk + last_idx * H * K + o_k2, mask=(o_k2 < K), other=0.0).to(tl.float32) b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype)) if K > 128: - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 128), (BT, 64), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT, 128), (BT, 64), (1, 0)) b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_GK: o_k3 = 128 + o_k1 - b_gk_last3 = tl.load(gk + last_idx * H*K + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) + b_gk_last3 = tl.load(gk + last_idx * H * K + o_k3, mask=(o_k3 < K), other=0.0).to(tl.float32) b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype)) if K > 192: - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 192), (BT, 64), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT, 192), (BT, 64), (1, 0)) b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_GK: o_k4 = 192 + o_k1 - b_gk_last4 = tl.load(gk + last_idx * H*K + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) + b_gk_last4 = tl.load(gk + last_idx * H * K + o_k4, mask=(o_k4 < K), other=0.0).to(tl.float32) b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) if USE_G: @@ -395,8 +408,8 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) # Update dh - p_w = tl.make_block_ptr(w, (K, T), (1, H*K), (0, i_t * BT), (64, BT), (0, 1)) - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (0, i_t * BT), (64, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (K, T), (1, H * K), (0, i_t * BT), (64, BT), (0, 1)) + p_q = tl.make_block_ptr(q, (K, T), (1, H * K), (0, i_t * BT), (64, BT), (0, 1)) b_w = tl.load(p_w, boundary_check=(0, 1)) b_q = tl.load(p_q, boundary_check=(0, 1)) if USE_G: @@ -409,8 +422,8 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( b_dh1 *= exp(b_gk_last1[:, None]) b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) if K > 64: - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (64, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, H*K), (64, i_t * BT), (64, BT), (0, 1)) + p_q = tl.make_block_ptr(q, (K, T), (1, H * K), (64, i_t * BT), (64, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (K, T), (1, H * K), (64, i_t * BT), (64, BT), (0, 1)) b_q = tl.load(p_q, boundary_check=(0, 1)) b_w = tl.load(p_w, boundary_check=(0, 1)) if USE_G: @@ -423,8 +436,8 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( b_dh2 *= exp(b_gk_last2[:, None]) b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) if K > 128: - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (128, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, H*K), (128, i_t * BT), (64, BT), (0, 1)) + p_q = tl.make_block_ptr(q, (K, T), (1, H * K), (128, i_t * BT), (64, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (K, T), (1, H * K), (128, i_t * BT), (64, BT), (0, 1)) b_q = tl.load(p_q, boundary_check=(0, 1)) b_w = tl.load(p_w, boundary_check=(0, 1)) if USE_G: @@ -437,8 +450,8 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( b_dh3 *= exp(b_gk_last3[:, None]) b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) if K > 192: - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (192, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, H*K), (192, i_t * BT), (64, BT), (0, 1)) + p_q = tl.make_block_ptr(q, (K, T), (1, H * K), (192, i_t * BT), (64, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (K, T), (1, H * K), (192, i_t * BT), (64, BT), (0, 1)) b_q = tl.load(p_q, boundary_check=(0, 1)) b_w = tl.load(p_w, boundary_check=(0, 1)) if USE_G: @@ -498,7 +511,10 @@ def chunk_gated_delta_rule_fwd_h( final_state = k.new_zeros(N, H, K, V, dtype=torch.float32) if output_final_state else None v_new = torch.empty_like(u) if save_new_value else None - def grid(meta): return (triton.cdiv(V, meta['BV']), N*H) + + def grid(meta): + return (triton.cdiv(V, meta["BV"]), N * H) + chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( k=k, v=u, @@ -553,7 +569,9 @@ def chunk_gated_delta_rule_bwd_dhu( dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None dv2 = torch.empty_like(dv) - def grid(meta): return (triton.cdiv(V, meta['BV']), N*H) + def grid(meta): + return (triton.cdiv(V, meta["BV"]), N * H) + chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64[grid]( q=q, k=k, diff --git a/src/xorl/ops/linear_attention/ops/common/chunk_o.py b/src/xorl/ops/linear_attention/ops/common/chunk_o.py index ef3b2d75..4a58b2a6 100644 --- a/src/xorl/ops/linear_attention/ops/common/chunk_o.py +++ b/src/xorl/ops/linear_attention/ops/common/chunk_o.py @@ -9,25 +9,28 @@ from xorl.ops.linear_attention.ops.utils.op import exp from xorl.ops.linear_attention.utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem -BKV_LIST = [64, 128] if check_shared_mem() else ([32, 64] if check_shared_mem('ada') else [32]) + +BKV_LIST = [64, 128] if check_shared_mem() else ([32, 64] if check_shared_mem("ada") else [32]) NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_G_GAMMA": lambda args: args["g_gamma"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ - triton.Config({'BK': 128, 'BV': 128}, num_warps=8, num_stages=3), - triton.Config({'BK': 64, 'BV': 64}, num_warps=4, num_stages=3), - triton.Config({'BK': 32, 'BV': 32}, num_warps=2, num_stages=3), + triton.Config({"BK": 128, "BV": 128}, num_warps=8, num_stages=3), + triton.Config({"BK": 64, "BV": 64}, num_warps=4, num_stages=3), + triton.Config({"BK": 32, "BV": 32}, num_warps=2, num_stages=3), ], - key=['H', 'K', 'V', 'BT'], + key=["H", "K", "V", "BT"], **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def chunk_fwd_kernel_o( q, k, @@ -69,14 +72,14 @@ def chunk_fwd_kernel_o( k += (bos * H + i_h) * K v += (bos * H + i_h) * V o += (bos * H + i_h) * V - h += (i_tg * H + i_h).to(tl.int64) * K*V + h += (i_tg * H + i_h).to(tl.int64) * K * V b_o = tl.zeros([BT, BV], dtype=tl.float32) b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (K, T), (1, H * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) # [BT, BK] b_q = tl.load(p_q, boundary_check=(0, 1)) @@ -108,8 +111,8 @@ def chunk_fwd_kernel_o( m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v = tl.make_block_ptr(v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr(o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) b_v = tl.load(p_v, boundary_check=(0, 1)) # to fix mma -> mma layout conversion @@ -118,22 +121,24 @@ def chunk_fwd_kernel_o( tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, - 'USE_DW': lambda args: args['dw'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_G_GAMMA": lambda args: args["g_gamma"] is not None, + "USE_DW": lambda args: args["dw"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in NUM_WARPS for num_stages in [2, 3, 4] ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'USE_G', 'USE_G_GAMMA', 'USE_DW'], + key=["H", "K", "V", "BT", "BK", "BV", "USE_G", "USE_G_GAMMA", "USE_DW"], **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def chunk_bwd_kernel_dqkwg( q, k, @@ -182,8 +187,8 @@ def chunk_bwd_kernel_dqkwg( # offset calculation v += (bos * H + i_h) * V do += (bos * H + i_h) * V - h += (i_tg * H + i_h).to(tl.int64) * K*V - dh += (i_tg * H + i_h).to(tl.int64) * K*V + h += (i_tg * H + i_h).to(tl.int64) * K * V + dh += (i_tg * H + i_h).to(tl.int64) * K * V q += (bos * H + i_h) * K k += (bos * H + i_h) * K dq += (bos * H + i_h) * K @@ -207,8 +212,8 @@ def chunk_bwd_kernel_dqkwg( b_dw = tl.zeros([BT, BK], dtype=tl.float32) if USE_DW else None for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v = tl.make_block_ptr(v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) # [BT, BV] @@ -218,7 +223,7 @@ def chunk_bwd_kernel_dqkwg( b_h = tl.load(p_h, boundary_check=(0, 1)) b_dh = tl.load(p_dh, boundary_check=(0, 1)) if USE_G: - b_dg_last += (tl.sum(b_h * b_dh)) + b_dg_last += tl.sum(b_h * b_dh) # [BT, BV] @ [BV, BT] -> [BT, BT] b_ds += tl.dot(b_do, tl.trans(b_v)) # [BT, BV] @ [BV, BK] -> [BT, BK] @@ -226,22 +231,22 @@ def chunk_bwd_kernel_dqkwg( # [BT, BV] @ [BV, BK] -> [BT, BK] b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) if USE_DW: - p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) b_dv = tl.load(p_dv, boundary_check=(0, 1)) b_dw += tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype)) if USE_DW: - p_dw = tl.make_block_ptr(dw, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dw = tl.make_block_ptr(dw, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) tl.debug_barrier() - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) b_q = tl.load(p_q, boundary_check=(0, 1)) b_k = tl.load(p_k, boundary_check=(0, 1)) - p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dq = tl.make_block_ptr(dq, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T @@ -300,21 +305,23 @@ def chunk_bwd_kernel_dqkwg( tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_G_GAMMA": lambda args: args["g_gamma"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in NUM_WARPS for num_stages in [2, 3, 4] ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'USE_G', 'USE_G_GAMMA'], + key=["H", "K", "V", "BT", "BK", "BV", "USE_G", "USE_G_GAMMA"], **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def chunk_bwd_kernel_dv( q, k, @@ -357,12 +364,12 @@ def chunk_bwd_kernel_dv( k += (bos * H + i_h) * K do += (bos * H + i_h) * V dv += (bos * H + i_h) * V - dh += (i_tg * H + i_h).to(tl.int64) * K*V + dh += (i_tg * H + i_h).to(tl.int64) * K * V b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_q = tl.make_block_ptr(q, (K, T), (1, H * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) b_q = tl.load(p_q, boundary_check=(0, 1)) b_k = tl.load(p_k, boundary_check=(0, 1)) b_A += tl.dot(b_k, b_q) @@ -388,29 +395,31 @@ def chunk_bwd_kernel_dv( b_dv *= tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] else: b_A = tl.where(m_A, b_A * scale, 0).to(do.dtype.element_ty) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) b_do = tl.load(p_do, boundary_check=(0, 1)) b_dv += tl.dot(b_A.to(b_do.dtype), b_do) tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, - 'USE_A': lambda args: args['A'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_G_GAMMA": lambda args: args["g_gamma"] is not None, + "USE_A": lambda args: args["A"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in NUM_WARPS for num_stages in [2, 3, 4] ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'USE_G'], + key=["H", "K", "V", "BT", "BK", "BV", "USE_G"], **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def chunk_bwd_kernel_dv_local( q, k, @@ -450,7 +459,7 @@ def chunk_bwd_kernel_dv_local( dv += (bos * H + i_h) * V if USE_A: - p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1)) b_A = tl.load(p_A, boundary_check=(0, 1)) else: if USE_G: @@ -463,8 +472,8 @@ def chunk_bwd_kernel_dv_local( b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_q = tl.make_block_ptr(q, (K, T), (1, H * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) b_k = tl.load(p_k, boundary_check=(0, 1)) b_q = tl.load(p_q, boundary_check=(0, 1)) @@ -478,8 +487,8 @@ def chunk_bwd_kernel_dv_local( b_A = tl.where(m_A, b_A, 0).to(do.dtype.element_ty) for i_v in range(tl.cdiv(V, BV)): - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) b_do = tl.load(p_do, boundary_check=(0, 1)) b_dv = tl.dot(b_A.to(b_do.dtype), b_do) tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) @@ -504,7 +513,10 @@ def chunk_fwd_o( scale = k.shape[-1] ** -0.5 o = torch.empty_like(v) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) + + def grid(meta): + return (triton.cdiv(V, meta["BV"]), NT, B * H) + chunk_fwd_kernel_o[grid]( q=q, k=k, @@ -540,9 +552,9 @@ def chunk_bwd_dv( BT = chunk_size chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None # H100 can have larger block size - if check_shared_mem('hopper', k.device.index): + if check_shared_mem("hopper", k.device.index): CONST_TILING = 128 - elif check_shared_mem('ada', k.device.index): + elif check_shared_mem("ada", k.device.index): CONST_TILING = 64 else: CONST_TILING = 32 @@ -594,9 +606,9 @@ def chunk_bwd_dv_local( if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) # H100 can have larger block size - if check_shared_mem('hopper', k.device.index): + if check_shared_mem("hopper", k.device.index): CONST_TILING = 128 - elif check_shared_mem('ada', k.device.index): + elif check_shared_mem("ada", k.device.index): CONST_TILING = 64 else: CONST_TILING = 32 @@ -643,15 +655,14 @@ def chunk_bwd_dqkwg( cu_seqlens: torch.LongTensor | None = None, chunk_size: int = 64, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] BT = chunk_size chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - if check_shared_mem('hopper', k.device.index): + if check_shared_mem("hopper", k.device.index): CONST_TILING = 128 - elif check_shared_mem('ada', k.device.index): + elif check_shared_mem("ada", k.device.index): CONST_TILING = 64 else: CONST_TILING = 32 diff --git a/src/xorl/ops/linear_attention/ops/common/chunk_scaled_dot_kkt.py b/src/xorl/ops/linear_attention/ops/common/chunk_scaled_dot_kkt.py index ba7ced91..62cf89e9 100644 --- a/src/xorl/ops/linear_attention/ops/common/chunk_scaled_dot_kkt.py +++ b/src/xorl/ops/linear_attention/ops/common/chunk_scaled_dot_kkt.py @@ -11,21 +11,23 @@ from xorl.ops.linear_attention.utils import autotune_cache_kwargs -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ - triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) + triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages) for BK in [32, 64, 128] for num_warps in [2, 4, 8] for num_stages in [2, 3, 4] ], - key=['H', 'K', 'BT', 'IS_VARLEN'], + key=["H", "K", "BT", "IS_VARLEN"], **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def chunk_scaled_dot_kkt_fwd_kernel( k, g, @@ -52,17 +54,17 @@ def chunk_scaled_dot_kkt_fwd_kernel( o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T - p_b = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) b_b = tl.load(p_b, boundary_check=(0,)) b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) b_k = tl.load(p_k, boundary_check=(0, 1)) b_A += tl.dot(b_k, tl.trans(b_k)) if USE_G: - p_g = tl.make_block_ptr(g + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) b_g = tl.load(p_g, boundary_check=(0,)) b_g_diff = b_g[:, None] - b_g[None, :] b_A *= exp(b_g_diff) @@ -70,7 +72,7 @@ def chunk_scaled_dot_kkt_fwd_kernel( m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) diff --git a/src/xorl/ops/linear_attention/ops/cp/__init__.py b/src/xorl/ops/linear_attention/ops/cp/__init__.py index 3c636879..3a2ef8e4 100644 --- a/src/xorl/ops/linear_attention/ops/cp/__init__.py +++ b/src/xorl/ops/linear_attention/ops/cp/__init__.py @@ -15,6 +15,7 @@ build_linear_attention_cp_context, ) + __all__ = [ "FLACPContext", "LinearAttentionCPContext", diff --git a/src/xorl/ops/linear_attention/ops/cp/chunk_delta_h.py b/src/xorl/ops/linear_attention/ops/cp/chunk_delta_h.py index 0aaf95a9..59bb550c 100644 --- a/src/xorl/ops/linear_attention/ops/cp/chunk_delta_h.py +++ b/src/xorl/ops/linear_attention/ops/cp/chunk_delta_h.py @@ -14,27 +14,30 @@ from xorl.ops.linear_attention.ops.utils.op import exp, exp2 from xorl.ops.linear_attention.utils import USE_CUDA_GRAPH, autotune_cache_kwargs, check_shared_mem + if TYPE_CHECKING: from xorl.ops.linear_attention.ops.cp.context import FLACPContext -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_GK': lambda args: args['gk'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ - triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) + triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4] for num_stages in [2, 3, 4] for BV in [32, 64] ], - key=['H', 'K', 'V', 'BT', 'USE_EXP2', "STAGE"], + key=["H", "K", "V", "BT", "USE_EXP2", "STAGE"], use_cuda_graph=USE_CUDA_GRAPH, **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def pre_process_fwd_kernel_stage1( k, v, @@ -120,28 +123,36 @@ def pre_process_fwd_kernel_stage1( if USE_GK: o_k1 = tl.arange(0, 64) - b_gk_last1 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) + b_gk_last1 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k1, mask=(o_k1 < K), other=0.0).to( + tl.float32 + ) if USE_EXP2: b_h1 *= exp2(b_gk_last1)[:, None] else: b_h1 *= exp(b_gk_last1)[:, None] if K > 64: o_k2 = 64 + o_k1 - b_gk_last2 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) + b_gk_last2 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k2, mask=(o_k2 < K), other=0.0).to( + tl.float32 + ) if USE_EXP2: b_h2 *= exp2(b_gk_last2)[:, None] else: b_h2 *= exp(b_gk_last2)[:, None] if K > 128: o_k3 = 128 + o_k1 - b_gk_last3 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) + b_gk_last3 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k3, mask=(o_k3 < K), other=0.0).to( + tl.float32 + ) if USE_EXP2: b_h3 *= exp2(b_gk_last3)[:, None] else: b_h3 *= exp(b_gk_last3)[:, None] if K > 192: o_k4 = 192 + o_k1 - b_gk_last4 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) + b_gk_last4 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k4, mask=(o_k4 < K), other=0.0).to( + tl.float32 + ) if USE_EXP2: b_h4 *= exp2(b_gk_last4)[:, None] else: @@ -178,23 +189,25 @@ def pre_process_fwd_kernel_stage1( tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_GK': lambda args: args['gk'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ - triton.Config({'BK2': BK2}, num_warps=num_warps, num_stages=num_stages) + triton.Config({"BK2": BK2}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4] for num_stages in [2, 3, 4] for BK2 in [32] ], - key=['H', 'BT', 'USE_EXP2', 'FORWARD'], + key=["H", "BT", "USE_EXP2", "FORWARD"], use_cuda_graph=USE_CUDA_GRAPH, **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def pre_process_fwd_bwd_kernel_stage2( k, w, @@ -257,7 +270,7 @@ def pre_process_fwd_bwd_kernel_stage2( b_g_last = exp(b_g_last) b_diag = tl.where(row[:, None] == row[None, :], b_g_last, 0.0) elif USE_GK: - b_gk_last = tl.load(gk + (bos + last_idx) * H * K + i_h * K + row, mask=(row < K), other=0.).to(tl.float32) + b_gk_last = tl.load(gk + (bos + last_idx) * H * K + i_h * K + row, mask=(row < K), other=0.0).to(tl.float32) if USE_EXP2: b_gk_last = exp2(b_gk_last) else: @@ -275,22 +288,24 @@ def pre_process_fwd_bwd_kernel_stage2( tl.store(p_m, b_m.to(p_m.dtype.element_ty), boundary_check=(0, 1)) -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_GK': lambda args: args['gk'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4] for num_stages in [2, 3, 4] ], - key=['H', 'K', 'V', 'BT', 'USE_EXP2'], + key=["H", "K", "V", "BT", "USE_EXP2"], use_cuda_graph=USE_CUDA_GRAPH, **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def pre_process_fwd_kernel_merged( k, v, @@ -388,28 +403,36 @@ def pre_process_fwd_kernel_merged( if USE_GK: o_k1 = tl.arange(0, 64) - b_gk_last1 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) + b_gk_last1 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k1, mask=(o_k1 < K), other=0.0).to( + tl.float32 + ) if USE_EXP2: b_h1 *= exp2(b_gk_last1)[:, None] else: b_h1 *= exp(b_gk_last1)[:, None] if K > 64: o_k2 = 64 + o_k1 - b_gk_last2 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) + b_gk_last2 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k2, mask=(o_k2 < K), other=0.0).to( + tl.float32 + ) if USE_EXP2: b_h2 *= exp2(b_gk_last2)[:, None] else: b_h2 *= exp(b_gk_last2)[:, None] if K > 128: o_k3 = 128 + o_k1 - b_gk_last3 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) + b_gk_last3 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k3, mask=(o_k3 < K), other=0.0).to( + tl.float32 + ) if USE_EXP2: b_h3 *= exp2(b_gk_last3)[:, None] else: b_h3 *= exp(b_gk_last3)[:, None] if K > 192: o_k4 = 192 + o_k1 - b_gk_last4 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) + b_gk_last4 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k4, mask=(o_k4 < K), other=0.0).to( + tl.float32 + ) if USE_EXP2: b_h4 *= exp2(b_gk_last4)[:, None] else: @@ -474,7 +497,9 @@ def pre_process_fwd_kernel_merged( b_g_last = exp(b_g_last) b_diag = tl.where(row[:, None] == row[None, :], b_g_last, 0.0) elif USE_GK: - b_gk_last = tl.load(gk + (bos + last_idx) * H * K + i_h * K + row, mask=(row < K), other=0.).to(tl.float32) + b_gk_last = tl.load(gk + (bos + last_idx) * H * K + i_h * K + row, mask=(row < K), other=0.0).to( + tl.float32 + ) if USE_EXP2: b_gk_last = exp2(b_gk_last) else: @@ -492,21 +517,23 @@ def pre_process_fwd_kernel_merged( tl.store(p_m, b_m.to(p_m.dtype.element_ty), boundary_check=(0, 1)) -@triton.heuristics({ - 'HAS_H0': lambda args: args['h0'] is not None, -}) +@triton.heuristics( + { + "HAS_H0": lambda args: args["h0"] is not None, + } +) @triton.autotune( configs=[ - triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) + triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4] for num_stages in [2, 3, 4] for BV in [32, 64] ], - key=['H', 'K', 'V', 'BT', 'USE_EXP2'], + key=["H", "K", "V", "BT", "USE_EXP2"], use_cuda_graph=USE_CUDA_GRAPH, **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['pre_or_post_num_ranks', 'rank', 'NUM_SEQ_ENTRIES']) +@triton.jit(do_not_specialize=["pre_or_post_num_ranks", "rank", "NUM_SEQ_ENTRIES"]) def merge_fwd_bwd_kernel( h, ag_hm, @@ -546,7 +573,11 @@ def merge_fwd_bwd_kernel( orig_seq_id = tl.load(h0_seq_ids + i_seq).to(tl.int32) p_h0 = tl.make_block_ptr( h0 + (orig_seq_id * H + i_h) * K * V, - (K, V), (V, 1), (0, i_v * BV), (BK, BV), (1, 0), + (K, V), + (V, 1), + (0, i_v * BV), + (BK, BV), + (1, 0), ) b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) else: @@ -557,11 +588,21 @@ def merge_fwd_bwd_kernel( base = i_ss * stride_hm_s + i_h * stride_hm_h p_he = tl.make_block_ptr( - ag_hm + base, (K, V), (V + K, 1), (0, i_v * BV), (BK, BV), (1, 0), + ag_hm + base, + (K, V), + (V + K, 1), + (0, i_v * BV), + (BK, BV), + (1, 0), ) b_he = tl.load(p_he, boundary_check=(0, 1)).to(tl.float32) p_m = tl.make_block_ptr( - ag_hm + base + V, (K, K), (V + K, 1), (0, 0), (BK, BK), (1, 0), + ag_hm + base + V, + (K, K), + (V + K, 1), + (0, 0), + (BK, BK), + (1, 0), ) b_m = tl.load(p_m, boundary_check=(0, 1)).to(tl.float32) b_h = tl.dot(b_m.to(tl.float32), b_h.to(tl.float32)) + b_he.to(tl.float32) @@ -571,7 +612,11 @@ def merge_fwd_bwd_kernel( stride_init = H * K * V p_out = tl.make_block_ptr( h + init_idx * stride_init + i_h * K * V, - (K, V), (V, 1), (0, i_v * BV), (BK, BV), (1, 0), + (K, V), + (V, 1), + (0, i_v * BV), + (BK, BV), + (1, 0), ) tl.store(p_out, b_h.to(p_out.dtype.element_ty), boundary_check=(0, 1)) else: @@ -595,23 +640,25 @@ def merge_fwd_bwd_kernel( tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_GK': lambda args: args['gk'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ - triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) + triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4] - for num_stages in ([4, 3, 2] if check_shared_mem('ampere') else [1]) + for num_stages in ([4, 3, 2] if check_shared_mem("ampere") else [1]) for BV in [64, 32] ], - key=['H', 'K', 'V', 'BT', 'BV', 'USE_G'], + key=["H", "K", "V", "BT", "BV", "USE_G"], use_cuda_graph=USE_CUDA_GRAPH, **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def pre_process_bwd_kernel_stage1( q, k, @@ -678,7 +725,7 @@ def pre_process_bwd_kernel_stage1( b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_GK: o_k1 = tl.arange(0, 64) - b_gk_last1 = tl.load(gk + last_idx * H * K + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) + b_gk_last1 = tl.load(gk + last_idx * H * K + o_k1, mask=(o_k1 < K), other=0.0).to(tl.float32) b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) if K > 64: @@ -686,7 +733,7 @@ def pre_process_bwd_kernel_stage1( b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_GK: o_k2 = 64 + o_k1 - b_gk_last2 = tl.load(gk + last_idx * H * K + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) + b_gk_last2 = tl.load(gk + last_idx * H * K + o_k2, mask=(o_k2 < K), other=0.0).to(tl.float32) b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype)) if K > 128: @@ -694,7 +741,7 @@ def pre_process_bwd_kernel_stage1( b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_GK: o_k3 = 128 + o_k1 - b_gk_last3 = tl.load(gk + last_idx * H * K + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) + b_gk_last3 = tl.load(gk + last_idx * H * K + o_k3, mask=(o_k3 < K), other=0.0).to(tl.float32) b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype)) if K > 192: @@ -702,7 +749,7 @@ def pre_process_bwd_kernel_stage1( b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_GK: o_k4 = 192 + o_k1 - b_gk_last4 = tl.load(gk + last_idx * H * K + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) + b_gk_last4 = tl.load(gk + last_idx * H * K + o_k4, mask=(o_k4 < K), other=0.0).to(tl.float32) b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) if USE_G: @@ -767,22 +814,24 @@ def pre_process_bwd_kernel_stage1( tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_GK': lambda args: args['gk'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4] - for num_stages in ([4, 3, 2] if check_shared_mem('ampere') else [1]) + for num_stages in ([4, 3, 2] if check_shared_mem("ampere") else [1]) ], - key=['H', 'K', 'V', 'BT', 'USE_EXP2'], + key=["H", "K", "V", "BT", "USE_EXP2"], use_cuda_graph=USE_CUDA_GRAPH, **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def pre_process_bwd_kernel_merged( q, k, @@ -856,7 +905,7 @@ def pre_process_bwd_kernel_merged( b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_GK: o_k1 = tl.arange(0, 64) - b_gk_last1 = tl.load(gk + last_idx * H * K + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) + b_gk_last1 = tl.load(gk + last_idx * H * K + o_k1, mask=(o_k1 < K), other=0.0).to(tl.float32) b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) if K > 64: @@ -864,7 +913,7 @@ def pre_process_bwd_kernel_merged( b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_GK: o_k2 = 64 + o_k1 - b_gk_last2 = tl.load(gk + last_idx * H * K + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) + b_gk_last2 = tl.load(gk + last_idx * H * K + o_k2, mask=(o_k2 < K), other=0.0).to(tl.float32) b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype)) if K > 128: @@ -872,7 +921,7 @@ def pre_process_bwd_kernel_merged( b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_GK: o_k3 = 128 + o_k1 - b_gk_last3 = tl.load(gk + last_idx * H * K + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) + b_gk_last3 = tl.load(gk + last_idx * H * K + o_k3, mask=(o_k3 < K), other=0.0).to(tl.float32) b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype)) if K > 192: @@ -880,7 +929,7 @@ def pre_process_bwd_kernel_merged( b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_GK: o_k4 = 192 + o_k1 - b_gk_last4 = tl.load(gk + last_idx * H * K + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) + b_gk_last4 = tl.load(gk + last_idx * H * K + o_k4, mask=(o_k4 < K), other=0.0).to(tl.float32) b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) if USE_G: @@ -989,7 +1038,9 @@ def pre_process_bwd_kernel_merged( b_g_last = exp(b_g_last) b_diag = tl.where(row[:, None] == row[None, :], b_g_last, 0.0) elif USE_GK: - b_gk_last = tl.load(gk + (bos + last_idx) * H * K + i_h * K + row, mask=(row < K), other=0.).to(tl.float32) + b_gk_last = tl.load(gk + (bos + last_idx) * H * K + i_h * K + row, mask=(row < K), other=0.0).to( + tl.float32 + ) if USE_EXP2: b_gk_last = exp2(b_gk_last) else: @@ -1058,7 +1109,10 @@ def chunk_gated_delta_rule_fwd_h_pre_process( ) ag_hm, _ = all_gather_into_tensor(hm, group=context.group) if not context.is_first_rank: - def grid(meta): return (triton.cdiv(V, meta['BV']), H) + + def grid(meta): + return (triton.cdiv(V, meta["BV"]), H) + merge_fwd_bwd_kernel[grid]( h=initial_state[0], ag_hm=ag_hm, @@ -1139,7 +1193,10 @@ def chunk_gated_delta_rule_bwd_dhu_pre_process( ag_dhm, _ = all_gather_into_tensor(dhm, group=context.group) if not context.is_last_rank: - def grid(meta): return (triton.cdiv(V, meta['BV']), H) + + def grid(meta): + return (triton.cdiv(V, meta["BV"]), H) + merge_fwd_bwd_kernel[grid]( h=dht[-1], ag_hm=ag_dhm, diff --git a/src/xorl/ops/linear_attention/ops/cp/context.py b/src/xorl/ops/linear_attention/ops/cp/context.py index 8b159272..f3e453e4 100644 --- a/src/xorl/ops/linear_attention/ops/cp/context.py +++ b/src/xorl/ops/linear_attention/ops/cp/context.py @@ -84,8 +84,8 @@ def get_cp_cu_seqlens( subset_cu_seqlens = cu_seqlens_cpu[start_seq_idx : end_seq_idx + 1] local_cu_seqlens_cpu = ( - subset_cu_seqlens.clamp(min=rank_start, max=rank_end) - rank_start - ).unique_consecutive().to(torch.int32) + (subset_cu_seqlens.clamp(min=rank_start, max=rank_end) - rank_start).unique_consecutive().to(torch.int32) + ) local_cu_seqlens_gpu = local_cu_seqlens_cpu.to(device=cu_seqlens.device, non_blocking=True) first_seq_global_start = int(cu_seqlens_cpu[start_seq_idx].item()) diff --git a/src/xorl/ops/linear_attention/ops/gated_delta_rule/__init__.py b/src/xorl/ops/linear_attention/ops/gated_delta_rule/__init__.py index b0183a19..d453f22e 100644 --- a/src/xorl/ops/linear_attention/ops/gated_delta_rule/__init__.py +++ b/src/xorl/ops/linear_attention/ops/gated_delta_rule/__init__.py @@ -1,6 +1,7 @@ from .chunk import chunk_gated_delta_rule from .fused_recurrent import fused_recurrent_gated_delta_rule + __all__ = [ "chunk_gated_delta_rule", "fused_recurrent_gated_delta_rule", diff --git a/src/xorl/ops/linear_attention/ops/gated_delta_rule/chunk.py b/src/xorl/ops/linear_attention/ops/gated_delta_rule/chunk.py index 1a143379..747758d7 100644 --- a/src/xorl/ops/linear_attention/ops/gated_delta_rule/chunk.py +++ b/src/xorl/ops/linear_attention/ops/gated_delta_rule/chunk.py @@ -6,7 +6,10 @@ import torch from xorl.ops.linear_attention.modules.l2norm import l2norm_bwd, l2norm_fwd -from xorl.ops.linear_attention.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h +from xorl.ops.linear_attention.ops.common.chunk_delta_h import ( + chunk_gated_delta_rule_bwd_dhu, + chunk_gated_delta_rule_fwd_h, +) from xorl.ops.linear_attention.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o from xorl.ops.linear_attention.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd from xorl.ops.linear_attention.ops.cp import FLACPContext @@ -193,7 +196,6 @@ def chunk_gated_delta_rule_bwd( class ChunkGatedDeltaRuleFunction(torch.autograd.Function): - @staticmethod @input_guard @autocast_custom_fwd @@ -340,7 +342,7 @@ def chunk_gated_delta_rule( cu_seqlens=cu_seqlens ) """ - if 'head_first' in kwargs: + if "head_first" in kwargs: warnings.warn( "head_first is deprecated and will be removed in a future version. " "Please use head_first=False for now instead.", diff --git a/src/xorl/ops/linear_attention/ops/gated_delta_rule/fused_recurrent.py b/src/xorl/ops/linear_attention/ops/gated_delta_rule/fused_recurrent.py index d585c924..d084c54f 100644 --- a/src/xorl/ops/linear_attention/ops/gated_delta_rule/fused_recurrent.py +++ b/src/xorl/ops/linear_attention/ops/gated_delta_rule/fused_recurrent.py @@ -2,7 +2,6 @@ # Adapted from flash-linear-attention/fla/ops/gated_delta_rule/fused_recurrent.py. # Portions of this file are adapted from flash-linear-attention, Copyright (c) 2023-2025 Songlin Yang, licensed under the MIT License. - import torch import triton import triton.language as tl diff --git a/src/xorl/ops/linear_attention/ops/gated_delta_rule/wy_fast.py b/src/xorl/ops/linear_attention/ops/gated_delta_rule/wy_fast.py index cf3ecced..18424ecc 100644 --- a/src/xorl/ops/linear_attention/ops/gated_delta_rule/wy_fast.py +++ b/src/xorl/ops/linear_attention/ops/gated_delta_rule/wy_fast.py @@ -7,7 +7,13 @@ from xorl.ops.linear_attention.ops.utils import prepare_chunk_indices from xorl.ops.linear_attention.ops.utils.op import exp -from xorl.ops.linear_attention.utils import IS_NVIDIA_BLACKWELL, IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem +from xorl.ops.linear_attention.utils import ( + IS_NVIDIA_BLACKWELL, + IS_NVIDIA_HOPPER, + autotune_cache_kwargs, + check_shared_mem, +) + if IS_NVIDIA_BLACKWELL: """ @@ -20,6 +26,7 @@ TODO: Remove this workaround once the Triton compiler bug is fixed. Track upstream issue at: https://github.com/triton-lang/triton/issues/8695 """ + @triton.jit def safe_dot(a, b): return tl.inline_asm_elementwise( @@ -31,6 +38,7 @@ def safe_dot(a, b): pack=1, ) else: + @triton.jit def safe_dot(a, b): return tl.dot(a, b) @@ -47,20 +55,22 @@ def safe_dot(a, b): ] -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4, 8] for num_stages in [2, 3, 4] ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def recompute_w_u_fwd_kernel( k, v, @@ -89,27 +99,27 @@ def recompute_w_u_fwd_kernel( T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) b_b = tl.load(p_b, boundary_check=(0,)) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) b_A = tl.load(p_A, boundary_check=(0, 1)) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_u = tl.make_block_ptr(u + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) b_v = tl.load(p_v, boundary_check=(0, 1)) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, allow_tf32=False) tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) if USE_G: - p_g = tl.make_block_ptr(g + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_g = tl.make_block_ptr(g + (bos * H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) b_g = exp(tl.load(p_g, boundary_check=(0,))) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_w = tl.make_block_ptr(w + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_w = tl.make_block_ptr(w + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) b_k = tl.load(p_k, boundary_check=(0, 1)) b_kb = b_k * b_b[:, None] if USE_G: @@ -118,16 +128,18 @@ def recompute_w_u_fwd_kernel( tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=PREPARE_WY_REPR_BWD_AUTOTUNE_CONFIGS, - key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def prepare_wy_repr_bwd_kernel( k, v, @@ -161,9 +173,9 @@ def prepare_wy_repr_bwd_kernel( else: bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_db = tl.make_block_ptr(db + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + p_b = tl.make_block_ptr(beta + (bos * H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_db = tl.make_block_ptr(db + (bos * H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1)) b_b = tl.load(p_b, boundary_check=(0,)) b_db = tl.zeros([BT], dtype=tl.float32) @@ -171,15 +183,15 @@ def prepare_wy_repr_bwd_kernel( b_dA = tl.zeros([BT, BT], dtype=tl.float32) if USE_G: - p_g = tl.make_block_ptr(g + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_g = tl.make_block_ptr(g + (bos * H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) b_g = tl.load(p_g, boundary_check=(0,)) b_g_exp = tl.exp(b_g) b_dg = tl.zeros([BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dw = tl.make_block_ptr(dw + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) # [BT, BK] b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_G: @@ -197,9 +209,9 @@ def prepare_wy_repr_bwd_kernel( tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_du = tl.make_block_ptr(du + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_du = tl.make_block_ptr(du + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) b_v = tl.load(p_v, boundary_check=(0, 1)) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_du = tl.load(p_du, boundary_check=(0, 1)) @@ -224,8 +236,8 @@ def prepare_wy_repr_bwd_kernel( tl.debug_barrier() for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) b_k = tl.load(p_k, boundary_check=(0, 1)) b_kt = tl.trans(b_k) b_ktb = b_kt * b_b[None, :] @@ -242,7 +254,7 @@ def prepare_wy_repr_bwd_kernel( b_A *= b_b[:, None] if USE_G: b_AdA = b_dA * b_A - p_dg = tl.make_block_ptr(dg + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_dg = tl.make_block_ptr(dg + (bos * H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) b_dg += tl.sum(b_AdA, axis=1) - tl.sum(b_AdA, axis=0) tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) @@ -265,7 +277,7 @@ def recompute_w_u_fwd( w = torch.empty_like(k) u = torch.empty_like(v) - recompute_w_u_fwd_kernel[(NT, B*H)]( + recompute_w_u_fwd_kernel[(NT, B * H)]( k=k, v=v, beta=beta, diff --git a/src/xorl/ops/linear_attention/ops/utils/__init__.py b/src/xorl/ops/linear_attention/ops/utils/__init__.py index e252c7cd..e754859e 100644 --- a/src/xorl/ops/linear_attention/ops/utils/__init__.py +++ b/src/xorl/ops/linear_attention/ops/utils/__init__.py @@ -3,6 +3,7 @@ from .op import exp, exp2, make_tensor_descriptor from .solve_tril import solve_tril + __all__ = [ "chunk_local_cumsum", "prepare_chunk_indices", diff --git a/src/xorl/ops/linear_attention/ops/utils/cumsum.py b/src/xorl/ops/linear_attention/ops/utils/cumsum.py index 670b6b9c..e27658a3 100644 --- a/src/xorl/ops/linear_attention/ops/utils/cumsum.py +++ b/src/xorl/ops/linear_attention/ops/utils/cumsum.py @@ -9,22 +9,22 @@ from xorl.ops.linear_attention.ops.utils.index import prepare_chunk_indices from xorl.ops.linear_attention.utils import autotune_cache_kwargs, check_shared_mem, input_guard + BS_LIST = [32, 64] if check_shared_mem() else [16, 32] -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "HAS_SCALE": lambda args: args["scale"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8] - ], - key=['B', 'H', 'BT', 'IS_VARLEN', 'REVERSE'], + configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]], + key=["B", "H", "BT", "IS_VARLEN", "REVERSE"], **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def chunk_local_cumsum_scalar_kernel( s, o, @@ -50,11 +50,11 @@ def chunk_local_cumsum_scalar_kernel( bos, eos = i_b * T, i_b * T + T if HEAD_FIRST: - p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) + p_s = tl.make_block_ptr(s + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,)) + p_o = tl.make_block_ptr(o + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,)) else: - p_s = tl.make_block_ptr(s + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_s = tl.make_block_ptr(s + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_o = tl.make_block_ptr(o + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) # [BT] b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) b_o = tl.cumsum(b_s, axis=0) @@ -66,20 +66,18 @@ def chunk_local_cumsum_scalar_kernel( tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "HAS_SCALE": lambda args: args["scale"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( - configs=[ - triton.Config({'BS': BS}, num_warps=num_warps) - for BS in BS_LIST - for num_warps in [2, 4, 8] - ], - key=['B', 'H', 'S', 'BT', 'IS_VARLEN', 'REVERSE'], + configs=[triton.Config({"BS": BS}, num_warps=num_warps) for BS in BS_LIST for num_warps in [2, 4, 8]], + key=["B", "H", "S", "BT", "IS_VARLEN", "REVERSE"], **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def chunk_local_cumsum_vector_kernel( s, o, @@ -107,11 +105,11 @@ def chunk_local_cumsum_vector_kernel( bos, eos = i_b * T, i_b * T + T if HEAD_FIRST: - p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + p_s = tl.make_block_ptr(s + (bos * H + i_h * T) * S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + p_o = tl.make_block_ptr(o + (bos * H + i_h * T) * S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) else: - p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H * S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H * S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) # [BT, BS] b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) if REVERSE: @@ -123,22 +121,23 @@ def chunk_local_cumsum_vector_kernel( tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "HAS_SCALE": lambda args: args["scale"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ - triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages) + triton.Config({"BT": BT}, num_warps=num_warps, num_stages=num_stages) for BT in [32, 64, 128, 256] for num_warps in [2, 4, 8] for num_stages in [1, 2, 3, 4] ], - key=['B', 'H', 'IS_VARLEN', 'REVERSE'], + key=["B", "H", "IS_VARLEN", "REVERSE"], **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def chunk_global_cumsum_scalar_kernel( s, o, @@ -166,11 +165,11 @@ def chunk_global_cumsum_scalar_kernel( for i_c in range(NT): i_t = NT - 1 - i_c if REVERSE else i_c if HEAD_FIRST: - p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) + p_s = tl.make_block_ptr(s + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,)) + p_o = tl.make_block_ptr(o + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,)) else: - p_s = tl.make_block_ptr(s + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_s = tl.make_block_ptr(s + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_o = tl.make_block_ptr(o + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) b_o = tl.cumsum(b_s, axis=0) b_ss = tl.sum(b_s, 0) @@ -184,21 +183,23 @@ def chunk_global_cumsum_scalar_kernel( tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "HAS_SCALE": lambda args: args["scale"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ - triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages) + triton.Config({"BT": BT}, num_warps=num_warps, num_stages=num_stages) for BT in [16, 32, 64, 128] for num_warps in [2, 4, 8] for num_stages in [1, 2, 3, 4] ], - key=['B', 'H', 'S', 'IS_VARLEN', 'REVERSE'], + key=["B", "H", "S", "IS_VARLEN", "REVERSE"], **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def chunk_global_cumsum_vector_kernel( s, o, @@ -228,11 +229,11 @@ def chunk_global_cumsum_vector_kernel( for i_c in range(NT): i_t = NT - 1 - i_c if REVERSE else i_c if HEAD_FIRST: - p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + p_s = tl.make_block_ptr(s + (bos * H + i_h * T) * S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + p_o = tl.make_block_ptr(o + (bos * H + i_h * T) * S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) else: - p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H * S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H * S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) # [BT, BS] b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) if REVERSE: @@ -259,7 +260,7 @@ def chunk_local_cumsum_scalar( B, H, T = g.shape else: B, T, H = g.shape - assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" + assert chunk_size == 2 ** (chunk_size.bit_length() - 1), "chunk_size must be a power of 2" BT = chunk_size if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) @@ -300,11 +301,14 @@ def chunk_local_cumsum_vector( if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" + assert chunk_size == 2 ** (chunk_size.bit_length() - 1), "chunk_size must be a power of 2" g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) - def grid(meta): return (triton.cdiv(meta['S'], meta['BS']), NT, B * H) - # keep cummulative normalizer in fp32 + + def grid(meta): + return (triton.cdiv(meta["S"], meta["BS"]), NT, B * H) + + # keep cumulative normalizer in fp32 # this kernel is equivalent to # g = g.view(B, H, NT, BT, -1).cumsum(-2).view(B, H, T, -1) chunk_local_cumsum_vector_kernel[grid]( diff --git a/src/xorl/ops/linear_attention/ops/utils/index.py b/src/xorl/ops/linear_attention/ops/utils/index.py index 8314f83d..bcae2f2a 100644 --- a/src/xorl/ops/linear_attention/ops/utils/index.py +++ b/src/xorl/ops/linear_attention/ops/utils/index.py @@ -1,9 +1,8 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang # Portions of this file are adapted from flash-linear-attention, Copyright (c) 2023-2025 Songlin Yang, licensed under the MIT License. import torch -import triton - import torch.nn.functional as F +import triton from xorl.ops.linear_attention.utils import tensor_cache @@ -20,8 +19,12 @@ def prepare_chunk_indices( cu_seqlens_cpu: torch.LongTensor | None = None, ) -> torch.LongTensor: if cu_seqlens_cpu is not None: - indices = torch.cat([torch.arange(n, device=cu_seqlens.device) - for n in triton.cdiv(prepare_lens(cu_seqlens_cpu), chunk_size).tolist()]) + indices = torch.cat( + [ + torch.arange(n, device=cu_seqlens.device) + for n in triton.cdiv(prepare_lens(cu_seqlens_cpu), chunk_size).tolist() + ] + ) return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens) indices = torch.cat([torch.arange(n) for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()]) return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens) diff --git a/src/xorl/ops/linear_attention/ops/utils/op.py b/src/xorl/ops/linear_attention/ops/utils/op.py index 9be5167f..200575c7 100644 --- a/src/xorl/ops/linear_attention/ops/utils/op.py +++ b/src/xorl/ops/linear_attention/ops/utils/op.py @@ -2,14 +2,15 @@ # Adapted from flash-linear-attention/fla/ops/utils/op.py. # Portions of this file are adapted from flash-linear-attention, Copyright (c) 2023-2025 Songlin Yang, licensed under the MIT License. - import os import triton import triton.language as tl import triton.language.extra.libdevice as tldevice + from xorl.ops.linear_attention.utils import IS_GATHER_SUPPORTED + if os.environ.get("FLA_USE_FAST_OPS", "0") == "1": @triton.jit diff --git a/src/xorl/ops/linear_attention/ops/utils/solve_tril.py b/src/xorl/ops/linear_attention/ops/utils/solve_tril.py index 4c0f4856..d1a34f59 100644 --- a/src/xorl/ops/linear_attention/ops/utils/solve_tril.py +++ b/src/xorl/ops/linear_attention/ops/utils/solve_tril.py @@ -11,25 +11,29 @@ from xorl.ops.linear_attention.ops.utils.op import make_tensor_descriptor from xorl.ops.linear_attention.utils import IS_TMA_SUPPORTED, autotune_cache_kwargs, input_guard -FLA_TRIL_PRECISION = os.environ.get('FLA_TRIL_PRECISION', 'ieee') -assert FLA_TRIL_PRECISION in ['ieee', 'tf32', 'tf32x3'], \ + +FLA_TRIL_PRECISION = os.environ.get("FLA_TRIL_PRECISION", "ieee") +assert FLA_TRIL_PRECISION in ["ieee", "tf32", "tf32x3"], ( f"FLA_TRIL_PRECISION must be one of 'ieee', 'tf32', or 'tf32x3', but got {FLA_TRIL_PRECISION}" +) DOT_PRECISION_AUTOTUNE_LIST = ["ieee"] if not IS_TMA_SUPPORTED else list({"ieee", FLA_TRIL_PRECISION}) -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ - triton.Config({'DOT_PRECISION': 'ieee'}, num_warps=num_warps, num_stages=num_stages) + triton.Config({"DOT_PRECISION": "ieee"}, num_warps=num_warps, num_stages=num_stages) for num_warps in [1, 2, 4, 8] for num_stages in [2, 3, 4, 5] ], - key=['BT'], + key=["BT"], **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def solve_tril_16x16_kernel( A, Ai, @@ -54,50 +58,52 @@ def solve_tril_16x16_kernel( m_A = o_i[:, None] > o_i[None, :] m_I = o_i[:, None] == o_i[None, :] - A = A + (bos*H + i_h) * BT - Ai = Ai + (bos*H + i_h) * 16 + A = A + (bos * H + i_h) * BT + Ai = Ai + (bos * H + i_h) * 16 offset = (i_t * 16) % BT if not USE_TMA: - p_A = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * 16, offset), (16, 16), (1, 0)) + p_A = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * 16, offset), (16, 16), (1, 0)) # [16, 16] b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) b_A = tl.where(m_A, b_A, 0) else: - desc = make_tensor_descriptor(A, [T, BT], [H*BT, 1], [16, 16]) - desc_o = make_tensor_descriptor(Ai, [T, 16], [H*16, 1], [16, 16]) + desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) + desc_o = make_tensor_descriptor(Ai, [T, 16], [H * 16, 1], [16, 16]) b_A = desc.load([i_t * 16, offset]).to(tl.float32) b_A = tl.where(m_A, b_A, 0) b_A = -b_A for i in range(2, min(16, T - i_t * 16)): # [16] - b_a = -tl.load(A + (i_t * 16 + i) * H*BT + o_i + offset) - b_a = tl.where(o_i < i, b_a, 0.) + b_a = -tl.load(A + (i_t * 16 + i) * H * BT + o_i + offset) + b_a = tl.where(o_i < i, b_a, 0.0) b_a = b_a + tl.sum(b_a[:, None] * b_A, 0) b_A = tl.where((o_i == i)[:, None], b_a, b_A) b_A += m_I if not USE_TMA: - p_Ai = tl.make_block_ptr(Ai, (T, 16), (H*16, 1), (i_t * 16, 0), (16, 16), (1, 0)) + p_Ai = tl.make_block_ptr(Ai, (T, 16), (H * 16, 1), (i_t * 16, 0), (16, 16), (1, 0)) tl.store(p_Ai, b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) else: desc_o.store([i_t * 16, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne")) -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ - triton.Config({'DOT_PRECISION': DOT_PRECISION}, num_warps=num_warps, num_stages=num_stages) + triton.Config({"DOT_PRECISION": DOT_PRECISION}, num_warps=num_warps, num_stages=num_stages) for num_warps in [1, 2, 4, 8] for num_stages in [2, 3, 4, 5] for DOT_PRECISION in DOT_PRECISION_AUTOTUNE_LIST ], - key=['H', 'BT', 'IS_VARLEN'], + key=["H", "BT", "IS_VARLEN"], **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def merge_16x16_to_32x32_inverse_kernel( A, Ai, @@ -126,13 +132,13 @@ def merge_16x16_to_32x32_inverse_kernel( Ai += (bos * H + i_h) * BT if not USE_TMA: - p_A_11 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) - p_A_22 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) + p_A_11 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) + p_A_22 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) else: - desc = make_tensor_descriptor(A, [T, BT], [H*BT, 1], [16, 16]) - desc_o = make_tensor_descriptor(Ai, [T, BT], [H*BT, 1], [16, 16]) + desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) + desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32) b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32) @@ -141,11 +147,11 @@ def merge_16x16_to_32x32_inverse_kernel( b_Ai_22 = -tl.where(m_A, b_Ai_22, 0) for i in range(2, min(16, T - i_t * BT)): - b_a_11 = -tl.load(A + (i_t * BT + i) * H*BT + o_i) + b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i) b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) for i in range(16 + 2, min(32, T - i_t * BT)): - b_a_22 = -tl.load(A + (i_t * BT + i) * H*BT + o_i + 16) + b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16) b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) @@ -153,7 +159,7 @@ def merge_16x16_to_32x32_inverse_kernel( b_Ai_22 += m_I if not USE_TMA: - p_A_21 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) + p_A_21 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) else: b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) @@ -161,9 +167,9 @@ def merge_16x16_to_32x32_inverse_kernel( b_Ai_21 = -tl.dot(tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), b_Ai_11, input_precision=DOT_PRECISION) if not USE_TMA: - p_Ai_11 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) - p_Ai_21 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) - p_Ai_22 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) + p_Ai_11 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) + p_Ai_21 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) + p_Ai_22 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) tl.store(p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) tl.store(p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) tl.store(p_Ai_21, b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) @@ -173,20 +179,22 @@ def merge_16x16_to_32x32_inverse_kernel( desc_o.store([i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne")) -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ - triton.Config({'DOT_PRECISION': DOT_PRECISION}, num_warps=num_warps, num_stages=num_stages) + triton.Config({"DOT_PRECISION": DOT_PRECISION}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4, 8] for num_stages in [2, 3, 4, 5] for DOT_PRECISION in DOT_PRECISION_AUTOTUNE_LIST ], - key=['H', 'BT', 'IS_VARLEN'], + key=["H", "BT", "IS_VARLEN"], **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def merge_16x16_to_64x64_inverse_kernel( A, Ai, @@ -215,17 +223,17 @@ def merge_16x16_to_64x64_inverse_kernel( Ai += (bos * H + i_h) * BT if not USE_TMA: - p_A_11 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) - p_A_22 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) - p_A_33 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)) - p_A_44 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)) + p_A_11 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) + p_A_22 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) + p_A_33 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)) + p_A_44 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)) b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32) b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32) else: - desc = make_tensor_descriptor(A, [T, BT], [H*BT, 1], [16, 16]) - desc_o = make_tensor_descriptor(Ai, [T, BT], [H*BT, 1], [16, 16]) + desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) + desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32) b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32) b_Ai_33 = desc.load([i_t * BT + 32, 32]).to(tl.float32) @@ -238,23 +246,23 @@ def merge_16x16_to_64x64_inverse_kernel( b_Ai_44 = -tl.where(m_A, b_Ai_44, 0) for i in range(2, min(16, T - i_t * BT)): - b_a_11 = -tl.load(A + (i_t * BT + i) * H*BT + o_i) - b_a_11 = tl.where(o_i < i, b_a_11, 0.) + b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i) + b_a_11 = tl.where(o_i < i, b_a_11, 0.0) b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) for i in range(16 + 2, min(32, T - i_t * BT)): - b_a_22 = -tl.load(A + (i_t * BT + i) * H*BT + o_i + 16) - b_a_22 = tl.where(o_i < i - 16, b_a_22, 0.) + b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16) + b_a_22 = tl.where(o_i < i - 16, b_a_22, 0.0) b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) for i in range(32 + 2, min(48, T - i_t * BT)): - b_a_33 = -tl.load(A + (i_t * BT + i) * H*BT + o_i + 32) - b_a_33 = tl.where(o_i < i - 32, b_a_33, 0.) + b_a_33 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 32) + b_a_33 = tl.where(o_i < i - 32, b_a_33, 0.0) b_a_33 += tl.sum(b_a_33[:, None] * b_Ai_33, 0) b_Ai_33 = tl.where((o_i == i - 32)[:, None], b_a_33, b_Ai_33) for i in range(48 + 2, min(64, T - i_t * BT)): - b_a_44 = -tl.load(A + (i_t * BT + i) * H*BT + o_i + 48) - b_a_44 = tl.where(o_i < i - 48, b_a_44, 0.) + b_a_44 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 48) + b_a_44 = tl.where(o_i < i - 48, b_a_44, 0.0) b_a_44 += tl.sum(b_a_44[:, None] * b_Ai_44, 0) b_Ai_44 = tl.where((o_i == i - 48)[:, None], b_a_44, b_Ai_44) b_Ai_11 += m_I @@ -263,12 +271,12 @@ def merge_16x16_to_64x64_inverse_kernel( b_Ai_44 += m_I if not USE_TMA: - p_A_21 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) - p_A_31 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)) - p_A_32 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)) - p_A_41 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)) - p_A_42 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)) - p_A_43 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)) + p_A_21 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) + p_A_31 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)) + p_A_32 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)) + p_A_41 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)) + p_A_42 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)) + p_A_43 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)) b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) b_A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32) b_A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32) @@ -289,35 +297,33 @@ def merge_16x16_to_64x64_inverse_kernel( b_Ai_31 = -tl.dot( b_Ai_33, - tl.dot(b_A_31, b_Ai_11, input_precision=DOT_PRECISION) + - tl.dot(b_A_32, b_Ai_21, input_precision=DOT_PRECISION), + tl.dot(b_A_31, b_Ai_11, input_precision=DOT_PRECISION) + tl.dot(b_A_32, b_Ai_21, input_precision=DOT_PRECISION), input_precision=DOT_PRECISION, ) b_Ai_42 = -tl.dot( b_Ai_44, - tl.dot(b_A_42, b_Ai_22, input_precision=DOT_PRECISION) + - tl.dot(b_A_43, b_Ai_32, input_precision=DOT_PRECISION), + tl.dot(b_A_42, b_Ai_22, input_precision=DOT_PRECISION) + tl.dot(b_A_43, b_Ai_32, input_precision=DOT_PRECISION), input_precision=DOT_PRECISION, ) b_Ai_41 = -tl.dot( b_Ai_44, - tl.dot(b_A_41, b_Ai_11, input_precision=DOT_PRECISION) + - tl.dot(b_A_42, b_Ai_21, input_precision=DOT_PRECISION) + - tl.dot(b_A_43, b_Ai_31, input_precision=DOT_PRECISION), + tl.dot(b_A_41, b_Ai_11, input_precision=DOT_PRECISION) + + tl.dot(b_A_42, b_Ai_21, input_precision=DOT_PRECISION) + + tl.dot(b_A_43, b_Ai_31, input_precision=DOT_PRECISION), input_precision=DOT_PRECISION, ) if not USE_TMA: - p_Ai_11 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) - p_Ai_22 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) - p_Ai_33 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)) - p_Ai_44 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)) - p_Ai_21 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) - p_Ai_31 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)) - p_Ai_32 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)) - p_Ai_41 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)) - p_Ai_42 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)) - p_Ai_43 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)) + p_Ai_11 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) + p_Ai_22 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) + p_Ai_33 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)) + p_Ai_44 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)) + p_Ai_21 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) + p_Ai_31 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)) + p_Ai_32 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)) + p_Ai_41 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)) + p_Ai_42 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)) + p_Ai_43 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)) tl.store(p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) tl.store(p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) tl.store(p_Ai_33, b_Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) diff --git a/src/xorl/ops/linear_attention/utils.py b/src/xorl/ops/linear_attention/utils.py index 238c4c86..733f0134 100644 --- a/src/xorl/ops/linear_attention/utils.py +++ b/src/xorl/ops/linear_attention/utils.py @@ -2,7 +2,6 @@ # Adapted from flash-linear-attention/fla/utils.py. # Portions of this file are adapted from flash-linear-attention, Copyright (c) 2023-2025 Songlin Yang, licensed under the MIT License. - import contextlib import functools import inspect @@ -13,6 +12,7 @@ import torch import triton + FLA_CACHE_RESULTS = os.getenv("FLA_CACHE_RESULTS", "1") == "1" SUPPORTS_AUTOTUNE_CACHE = "cache_results" in inspect.signature(triton.autotune).parameters autotune_cache_kwargs = {"cache_results": FLA_CACHE_RESULTS} if SUPPORTS_AUTOTUNE_CACHE else {} diff --git a/src/xorl/ops/loss/causallm_loss.py b/src/xorl/ops/loss/causallm_loss.py index 98723f2f..29bed248 100644 --- a/src/xorl/ops/loss/causallm_loss.py +++ b/src/xorl/ops/loss/causallm_loss.py @@ -48,7 +48,7 @@ def causallm_loss_function( # Flatten the labels and hidden_states labels_flat = labels.view(-1) hidden_states_flat = hidden_states.view(-1, hidden_states.size(-1)) - valid_mask = (labels_flat != ignore_index) + valid_mask = labels_flat != ignore_index # Vocab-parallel cross-entropy for tensor parallelism if tp_group is not None: @@ -56,7 +56,10 @@ def causallm_loss_function( local_weight = weight.to_local() if hasattr(weight, "to_local") else weight per_token_ce = vocab_parallel_cross_entropy( - hidden_states_flat, local_weight, labels_flat, tp_group, + hidden_states_flat, + local_weight, + labels_flat, + tp_group, ignore_index=ignore_index, use_compile=use_compile, ) @@ -73,7 +76,9 @@ def causallm_loss_function( if return_per_token: # Compute cross-entropy based on mode if ce_mode == "compiled": - per_token_ce = compiled_cross_entropy_function(hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, lm_head_fp32=lm_head_fp32) + per_token_ce = compiled_cross_entropy_function( + hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, lm_head_fp32=lm_head_fp32 + ) else: # eager mode if lm_head_fp32: logits_flat = (hidden_states_flat.float() @ weight.float().t()).float() @@ -93,7 +98,9 @@ def causallm_loss_function( # Keeping the autograd graph intact is critical for FSDP2: all ranks must # trigger reduce-scatter for every parameter, including lm_head weight. if ce_mode == "compiled": - per_token_ce = compiled_cross_entropy_function(hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, lm_head_fp32=lm_head_fp32) + per_token_ce = compiled_cross_entropy_function( + hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, lm_head_fp32=lm_head_fp32 + ) else: # eager mode if lm_head_fp32: logits_flat = (hidden_states_flat.float() @ weight.float().t()).float() diff --git a/src/xorl/ops/loss/compiled_cross_entropy.py b/src/xorl/ops/loss/compiled_cross_entropy.py index f5addf79..c4daf237 100644 --- a/src/xorl/ops/loss/compiled_cross_entropy.py +++ b/src/xorl/ops/loss/compiled_cross_entropy.py @@ -17,6 +17,7 @@ # Check if auto_chunker is available _AUTO_CHUNKER_AVAILABLE = None + def _check_auto_chunker_available() -> bool: """Check if torch.compile auto_chunker option is available.""" global _AUTO_CHUNKER_AVAILABLE @@ -24,9 +25,11 @@ def _check_auto_chunker_available() -> bool: return _AUTO_CHUNKER_AVAILABLE try: + @torch.compile(options={"auto_chunker.enable": True, "auto_chunker.num_chunk": 2}) def _test_fn(x): return x * 2 + # Don't actually run, just check if compilation setup works _AUTO_CHUNKER_AVAILABLE = True except (RuntimeError, TypeError): @@ -91,6 +94,7 @@ def _get_compiled_ce_fn(num_chunks: int, reduction: str = "none") -> Callable: """ cache_key = (num_chunks, reduction) if cache_key not in _compiled_ce_cache: + def _compute_ce(hidden_states, weight, labels, ignore_index): logits = (hidden_states @ weight.t()).float() return F.cross_entropy(logits, labels, reduction=reduction, ignore_index=ignore_index) diff --git a/src/xorl/ops/loss/grpo_loss.py b/src/xorl/ops/loss/grpo_loss.py index bca416e5..8628a41f 100644 --- a/src/xorl/ops/loss/grpo_loss.py +++ b/src/xorl/ops/loss/grpo_loss.py @@ -204,8 +204,14 @@ def drgrpo_loss_function( hidden_states_flat = hidden_states.reshape(-1, H) per_token_ce = compute_per_token_ce( - hidden_states_flat, weight, labels_flat, ignore_index, - ce_mode, num_chunks, tp_group=tp_group, lm_head_fp32=lm_head_fp32, + hidden_states_flat, + weight, + labels_flat, + ignore_index, + ce_mode, + num_chunks, + tp_group=tp_group, + lm_head_fp32=lm_head_fp32, ) logprobs = -per_token_ce.view(B, S) diff --git a/src/xorl/ops/loss/importance_sampling_loss.py b/src/xorl/ops/loss/importance_sampling_loss.py index 6e71039c..beb5e274 100644 --- a/src/xorl/ops/loss/importance_sampling_loss.py +++ b/src/xorl/ops/loss/importance_sampling_loss.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Dict, Optional, Tuple +from typing import Optional import torch import torch.distributed as dist @@ -64,20 +64,26 @@ def importance_sampling_loss_function( advantages_flat = advantages.reshape(-1) # Valid/action mask - valid_mask = (labels_flat != ignore_index) + valid_mask = labels_flat != ignore_index n_valid = valid_mask.sum().clamp(min=1).float() # ---- Cross-entropy computation ---- per_token_ce = compute_per_token_ce( - hidden_states_flat, weight, labels_flat, ignore_index, ce_mode, num_chunks, - tp_group=tp_group, lm_head_fp32=lm_head_fp32, + hidden_states_flat, + weight, + labels_flat, + ignore_index, + ce_mode, + num_chunks, + tp_group=tp_group, + lm_head_fp32=lm_head_fp32, ) # new logprobs = log p(target) = -CE new_logprobs_flat = -per_token_ce.detach() # ---- ratio computation (no sanitization) ---- - delta = (new_logprobs_flat - old_logprobs_flat) + delta = new_logprobs_flat - old_logprobs_flat delta = delta.masked_fill(~valid_mask, 0.0) delta = delta.clamp(min=-20.0, max=20.0) ratio = torch.exp(delta) diff --git a/src/xorl/ops/loss/per_token_ce.py b/src/xorl/ops/loss/per_token_ce.py index 2dafc3c4..d28f35ce 100644 --- a/src/xorl/ops/loss/per_token_ce.py +++ b/src/xorl/ops/loss/per_token_ce.py @@ -42,14 +42,21 @@ def compute_per_token_ce( if tp_group is not None: local_weight = weight.to_local() if hasattr(weight, "to_local") else weight return vocab_parallel_cross_entropy( - hidden_states_flat, local_weight, labels_flat, tp_group, + hidden_states_flat, + local_weight, + labels_flat, + tp_group, ignore_index=ignore_index, use_compile=use_compile, ) if ce_mode == "compiled": return compiled_cross_entropy_function( - hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, + hidden_states_flat, + weight, + labels_flat, + ignore_index, + num_chunks, lm_head_fp32=lm_head_fp32, ) else: diff --git a/src/xorl/ops/loss/policy_loss.py b/src/xorl/ops/loss/policy_loss.py index 4c1ce1c5..5616b5e9 100644 --- a/src/xorl/ops/loss/policy_loss.py +++ b/src/xorl/ops/loss/policy_loss.py @@ -14,9 +14,11 @@ import torch import torch.distributed as dist + from xorl.ops.loss.loss_output import LossOutput from xorl.ops.loss.per_token_ce import compute_per_token_ce + logger = logging.getLogger(__name__) @@ -179,13 +181,19 @@ def policy_loss_function( advantages_flat = advantages.view(-1) # Create mask for valid tokens (use labels != ignore_index) - valid_mask = (labels_flat != ignore_index) + valid_mask = labels_flat != ignore_index n_valid = valid_mask.sum().clamp(min=1) # Compute cross-entropy (supports vocab-parallel TP via tp_group) per_token_ce = compute_per_token_ce( - hidden_states_flat, weight, labels_flat, ignore_index, ce_mode, num_chunks, - tp_group=tp_group, lm_head_fp32=lm_head_fp32, + hidden_states_flat, + weight, + labels_flat, + ignore_index, + ce_mode, + num_chunks, + tp_group=tp_group, + lm_head_fp32=lm_head_fp32, ) new_logprobs_flat = -per_token_ce.detach() @@ -239,8 +247,10 @@ def policy_loss_function( icepop_mask = None if icepop_beta is not None: if use_tis: - logger.warning("IcePop and TIS are both enabled. IcePop makes TIS redundant " - "when using inference logprobs as old_logprobs.") + logger.warning( + "IcePop and TIS are both enabled. IcePop makes TIS redundant " + "when using inference logprobs as old_logprobs." + ) ratio_d = ratio.detach() icepop_mask = (ratio_d >= 1.0 / icepop_beta) & (ratio_d <= icepop_beta) diff --git a/src/xorl/ops/loss/vocab_parallel_cross_entropy.py b/src/xorl/ops/loss/vocab_parallel_cross_entropy.py index 0d55a632..494ca491 100644 --- a/src/xorl/ops/loss/vocab_parallel_cross_entropy.py +++ b/src/xorl/ops/loss/vocab_parallel_cross_entropy.py @@ -175,7 +175,13 @@ def forward( ) -> torch.Tensor: kernel = _get_compiled_forward_kernel() if use_compile else _forward_kernel per_token_ce, global_max, global_sumexp, target_in_range, safe_local_target, valid_mask = kernel( - hidden_states, weight, labels, tp_group, vocab_offset, local_vocab_size, ignore_index, + hidden_states, + weight, + labels, + tp_group, + vocab_offset, + local_vocab_size, + ignore_index, ) # Save tiny [BT,1] statistics + inputs; NOT the huge [BT, V/tp] softmax @@ -194,9 +200,16 @@ def backward(ctx, grad_output: torch.Tensor): kernel = _get_compiled_backward_kernel() if ctx.use_compile else _backward_kernel grad_hidden, grad_weight = kernel( - grad_output, hidden_states, weight, global_max, global_sumexp, - ctx.target_in_range, ctx.safe_local_target, ctx.valid_mask, - ctx.tp_group, weight.requires_grad, + grad_output, + hidden_states, + weight, + global_max, + global_sumexp, + ctx.target_in_range, + ctx.safe_local_target, + ctx.valid_mask, + ctx.tp_group, + weight.requires_grad, ) return grad_hidden, grad_weight, None, None, None, None, None, None @@ -239,8 +252,13 @@ def vocab_parallel_cross_entropy( end = min(start + chunk_size, BT) ce_chunks.append( _VocabParallelCrossEntropy.apply( - hidden_states[start:end], weight, labels[start:end], - tp_group, vocab_offset, local_vocab_size, ignore_index, + hidden_states[start:end], + weight, + labels[start:end], + tp_group, + vocab_offset, + local_vocab_size, + ignore_index, use_compile, ) ) diff --git a/src/xorl/ops/moe/__init__.py b/src/xorl/ops/moe/__init__.py index 280b9fa8..da074616 100644 --- a/src/xorl/ops/moe/__init__.py +++ b/src/xorl/ops/moe/__init__.py @@ -1,16 +1,17 @@ """MoE ops subpackage — re-exports for backward compatibility.""" -from .triton import TritonMoeExpertsFunction, triton_moe_forward, TritonEPGroupGemm -from .triton_lora import ( - TritonMoeExpertsLoRAFunction, - triton_moe_lora_forward, - TritonEPGroupGemmWithLoRA, -) from .lora import ( make_ep_lora_compute, make_local_lora_compute, ) -from .quack import quack_moe_forward, QuackEPGroupGemm +from .quack import QuackEPGroupGemm, quack_moe_forward +from .triton import TritonEPGroupGemm, TritonMoeExpertsFunction, triton_moe_forward +from .triton_lora import ( + TritonEPGroupGemmWithLoRA, + TritonMoeExpertsLoRAFunction, + triton_moe_lora_forward, +) + try: from .quack_lora import QuackEPGroupGemmWithLoRA, quack_moe_lora_forward diff --git a/src/xorl/ops/moe/lora.py b/src/xorl/ops/moe/lora.py index a115afc9..0af7b710 100644 --- a/src/xorl/ops/moe/lora.py +++ b/src/xorl/ops/moe/lora.py @@ -49,10 +49,15 @@ def forward( ctx, permute_tokens, cumsum, - gate_proj, up_proj, down_proj, - gate_proj_lora_A, gate_proj_lora_B, - up_proj_lora_A, up_proj_lora_B, - down_proj_lora_A, down_proj_lora_B, + gate_proj, + up_proj, + down_proj, + gate_proj_lora_A, + gate_proj_lora_B, + up_proj_lora_A, + up_proj_lora_B, + down_proj_lora_A, + down_proj_lora_B, scaling, ): max_M = permute_tokens.shape[0] @@ -94,8 +99,12 @@ def forward( # === gate_proj and up_proj with FUSED base GEMM === gate_up_weight = torch.cat([gate_proj, up_proj], dim=2) gate_up_output = gemm_nk( - a=permute_tokens, b=gate_up_weight, - cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, + a=permute_tokens, + b=gate_up_weight, + cumsum_M=cumsum, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) gate_base = gate_up_output[:, :intermediate_size] up_base = gate_up_output[:, intermediate_size:] @@ -103,23 +112,35 @@ def forward( # === gate and up LoRA A: Fused GEMM === gate_up_lora_A = torch.cat([gate_proj_lora_A, up_proj_lora_A], dim=2) gate_up_lora_intermediate = gemm_nk( - a=permute_tokens, b=gate_up_lora_A, - cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, + a=permute_tokens, + b=gate_up_lora_A, + cumsum_M=cumsum, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) gate_lora_intermediate = gate_up_lora_intermediate[:, :lora_r].contiguous() up_lora_intermediate = gate_up_lora_intermediate[:, lora_r:].contiguous() # === gate LoRA B === gate_lora_output = gemm_nk( - a=gate_lora_intermediate, b=gate_proj_lora_B, - cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, + a=gate_lora_intermediate, + b=gate_proj_lora_B, + cumsum_M=cumsum, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) gate_output = (gate_base + gate_lora_output * scaling).contiguous() # === up LoRA B === up_lora_output = gemm_nk( - a=up_lora_intermediate, b=up_proj_lora_B, - cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, + a=up_lora_intermediate, + b=up_proj_lora_B, + cumsum_M=cumsum, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) up_output = (up_base + up_lora_output * scaling).contiguous() @@ -129,16 +150,28 @@ def forward( # === down_proj with LoRA === down_output = gemm_nk( - a=gated_output, b=down_proj, - cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, + a=gated_output, + b=down_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) down_lora_intermediate = gemm_nk( - a=gated_output, b=down_proj_lora_A, - cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, + a=gated_output, + b=down_proj_lora_A, + cumsum_M=cumsum, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) down_lora_output = gemm_nk( - a=down_lora_intermediate, b=down_proj_lora_B, - cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, + a=down_lora_intermediate, + b=down_proj_lora_B, + cumsum_M=cumsum, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) down_output = down_output + down_lora_output * scaling @@ -150,26 +183,46 @@ def forward( ctx.num_local_experts = num_local_experts ctx.compute_dtype = compute_dtype ctx.save_for_backward( - permute_tokens, cumsum, - gate_proj, up_proj, down_proj, - orig_gate_proj_lora_A, orig_gate_proj_lora_B, - orig_up_proj_lora_A, orig_up_proj_lora_B, - orig_down_proj_lora_A, orig_down_proj_lora_B, - gate_output, up_output, gated_output, - gate_lora_intermediate, up_lora_intermediate, down_lora_intermediate, + permute_tokens, + cumsum, + gate_proj, + up_proj, + down_proj, + orig_gate_proj_lora_A, + orig_gate_proj_lora_B, + orig_up_proj_lora_A, + orig_up_proj_lora_B, + orig_down_proj_lora_A, + orig_down_proj_lora_B, + gate_output, + up_output, + gated_output, + gate_lora_intermediate, + up_lora_intermediate, + down_lora_intermediate, ) return down_output @staticmethod def backward(ctx, grad_output): ( - permute_tokens, cumsum, - gate_proj, up_proj, down_proj, - gate_proj_lora_A, gate_proj_lora_B, - up_proj_lora_A, up_proj_lora_B, - down_proj_lora_A, down_proj_lora_B, - gate_output, up_output, gated_output, - gate_lora_intermediate, up_lora_intermediate, down_lora_intermediate, + permute_tokens, + cumsum, + gate_proj, + up_proj, + down_proj, + gate_proj_lora_A, + gate_proj_lora_B, + up_proj_lora_A, + up_proj_lora_B, + down_proj_lora_A, + down_proj_lora_B, + gate_output, + up_output, + gated_output, + gate_lora_intermediate, + up_lora_intermediate, + down_lora_intermediate, ) = ctx.saved_tensors scaling = ctx.scaling @@ -182,12 +235,24 @@ def backward(ctx, grad_output): max_M = grad_output.shape[0] # Cast LoRA weights to compute dtype for backward - gate_proj_lora_A_compute = gate_proj_lora_A.to(compute_dtype) if gate_proj_lora_A.dtype != compute_dtype else gate_proj_lora_A - gate_proj_lora_B_compute = gate_proj_lora_B.to(compute_dtype) if gate_proj_lora_B.dtype != compute_dtype else gate_proj_lora_B - up_proj_lora_A_compute = up_proj_lora_A.to(compute_dtype) if up_proj_lora_A.dtype != compute_dtype else up_proj_lora_A - up_proj_lora_B_compute = up_proj_lora_B.to(compute_dtype) if up_proj_lora_B.dtype != compute_dtype else up_proj_lora_B - down_proj_lora_A_compute = down_proj_lora_A.to(compute_dtype) if down_proj_lora_A.dtype != compute_dtype else down_proj_lora_A - down_proj_lora_B_compute = down_proj_lora_B.to(compute_dtype) if down_proj_lora_B.dtype != compute_dtype else down_proj_lora_B + gate_proj_lora_A_compute = ( + gate_proj_lora_A.to(compute_dtype) if gate_proj_lora_A.dtype != compute_dtype else gate_proj_lora_A + ) + gate_proj_lora_B_compute = ( + gate_proj_lora_B.to(compute_dtype) if gate_proj_lora_B.dtype != compute_dtype else gate_proj_lora_B + ) + up_proj_lora_A_compute = ( + up_proj_lora_A.to(compute_dtype) if up_proj_lora_A.dtype != compute_dtype else up_proj_lora_A + ) + up_proj_lora_B_compute = ( + up_proj_lora_B.to(compute_dtype) if up_proj_lora_B.dtype != compute_dtype else up_proj_lora_B + ) + down_proj_lora_A_compute = ( + down_proj_lora_A.to(compute_dtype) if down_proj_lora_A.dtype != compute_dtype else down_proj_lora_A + ) + down_proj_lora_B_compute = ( + down_proj_lora_B.to(compute_dtype) if down_proj_lora_B.dtype != compute_dtype else down_proj_lora_B + ) if gate_A_shared: gate_proj_lora_A_compute = gate_proj_lora_A_compute.expand(num_local_experts, -1, -1).contiguous() @@ -198,47 +263,75 @@ def backward(ctx, grad_output): # Initialize LoRA gradients grad_gate_proj_lora_A_full = torch.zeros( - num_local_experts, gate_proj_lora_A.shape[1], gate_proj_lora_A.shape[2], - dtype=gate_proj_lora_A.dtype, device=gate_proj_lora_A.device, + num_local_experts, + gate_proj_lora_A.shape[1], + gate_proj_lora_A.shape[2], + dtype=gate_proj_lora_A.dtype, + device=gate_proj_lora_A.device, ) grad_gate_proj_lora_B = torch.zeros_like(gate_proj_lora_B) grad_up_proj_lora_A_full = torch.zeros( - num_local_experts, up_proj_lora_A.shape[1], up_proj_lora_A.shape[2], - dtype=up_proj_lora_A.dtype, device=up_proj_lora_A.device, + num_local_experts, + up_proj_lora_A.shape[1], + up_proj_lora_A.shape[2], + dtype=up_proj_lora_A.dtype, + device=up_proj_lora_A.device, ) grad_up_proj_lora_B = torch.zeros_like(up_proj_lora_B) grad_down_proj_lora_A = torch.zeros_like(down_proj_lora_A) grad_down_proj_lora_B_full = torch.zeros( - num_local_experts, down_proj_lora_B.shape[1], down_proj_lora_B.shape[2], - dtype=down_proj_lora_B.dtype, device=down_proj_lora_B.device, + num_local_experts, + down_proj_lora_B.shape[1], + down_proj_lora_B.shape[2], + dtype=down_proj_lora_B.dtype, + device=down_proj_lora_B.device, ) # === down_proj backward === grad_gated_output = gemm_nk( - a=grad_output, b=down_proj, - cumsum_M=cumsum, max_M=max_M, transpose_b=True, + a=grad_output, + b=down_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_b=True, ) gemm_mn( - a=down_lora_intermediate, b=grad_output, c=grad_down_proj_lora_B_full, - cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, + a=down_lora_intermediate, + b=grad_output, + c=grad_down_proj_lora_B_full, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) grad_down_proj_lora_B_full.mul_(scaling) grad_down_lora_intermediate = gemm_nk( - a=grad_output, b=down_proj_lora_B_compute, - cumsum_M=cumsum, max_M=max_M, transpose_b=True, + a=grad_output, + b=down_proj_lora_B_compute, + cumsum_M=cumsum, + max_M=max_M, + transpose_b=True, ) grad_down_lora_intermediate.mul_(scaling) gemm_mn( - a=gated_output, b=grad_down_lora_intermediate, c=grad_down_proj_lora_A, - cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, + a=gated_output, + b=grad_down_lora_intermediate, + c=grad_down_proj_lora_A, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) grad_gated_from_lora = gemm_nk( - a=grad_down_lora_intermediate, b=down_proj_lora_A_compute, - cumsum_M=cumsum, max_M=max_M, transpose_b=True, + a=grad_down_lora_intermediate, + b=down_proj_lora_A_compute, + cumsum_M=cumsum, + max_M=max_M, + transpose_b=True, ) grad_gated_output = grad_gated_output + grad_gated_from_lora @@ -249,64 +342,104 @@ def backward(ctx, grad_output): # === up_proj backward === grad_permute_tokens_2 = gemm_nk( - a=grad_up_output, b=up_proj, - cumsum_M=cumsum, max_M=max_M, transpose_b=True, + a=grad_up_output, + b=up_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_b=True, ) gemm_mn( - a=up_lora_intermediate, b=grad_up_output, c=grad_up_proj_lora_B, - cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, + a=up_lora_intermediate, + b=grad_up_output, + c=grad_up_proj_lora_B, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) grad_up_proj_lora_B.mul_(scaling) grad_up_lora_intermediate = gemm_nk( - a=grad_up_output, b=up_proj_lora_B_compute, - cumsum_M=cumsum, max_M=max_M, transpose_b=True, + a=grad_up_output, + b=up_proj_lora_B_compute, + cumsum_M=cumsum, + max_M=max_M, + transpose_b=True, ) grad_up_lora_intermediate.mul_(scaling) gemm_mn( - a=permute_tokens, b=grad_up_lora_intermediate, c=grad_up_proj_lora_A_full, - cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, + a=permute_tokens, + b=grad_up_lora_intermediate, + c=grad_up_proj_lora_A_full, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) grad_permute_from_lora_2 = gemm_nk( - a=grad_up_lora_intermediate, b=up_proj_lora_A_compute, - cumsum_M=cumsum, max_M=max_M, transpose_b=True, + a=grad_up_lora_intermediate, + b=up_proj_lora_A_compute, + cumsum_M=cumsum, + max_M=max_M, + transpose_b=True, ) grad_permute_tokens_2 = grad_permute_tokens_2 + grad_permute_from_lora_2 # === SiLU backward === - grad_gate_output = torch.ops.aten.silu_backward( - grad_gate_activation.float(), gate_output.float() - ).to(gate_output.dtype).contiguous() + grad_gate_output = ( + torch.ops.aten.silu_backward(grad_gate_activation.float(), gate_output.float()) + .to(gate_output.dtype) + .contiguous() + ) # === gate_proj backward === grad_permute_tokens_1 = gemm_nk( - a=grad_gate_output, b=gate_proj, - cumsum_M=cumsum, max_M=max_M, transpose_b=True, + a=grad_gate_output, + b=gate_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_b=True, ) gemm_mn( - a=gate_lora_intermediate, b=grad_gate_output, c=grad_gate_proj_lora_B, - cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, + a=gate_lora_intermediate, + b=grad_gate_output, + c=grad_gate_proj_lora_B, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) grad_gate_proj_lora_B.mul_(scaling) grad_gate_lora_intermediate = gemm_nk( - a=grad_gate_output, b=gate_proj_lora_B_compute, - cumsum_M=cumsum, max_M=max_M, transpose_b=True, + a=grad_gate_output, + b=gate_proj_lora_B_compute, + cumsum_M=cumsum, + max_M=max_M, + transpose_b=True, ) grad_gate_lora_intermediate.mul_(scaling) gemm_mn( - a=permute_tokens, b=grad_gate_lora_intermediate, c=grad_gate_proj_lora_A_full, - cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, + a=permute_tokens, + b=grad_gate_lora_intermediate, + c=grad_gate_proj_lora_A_full, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) grad_permute_from_lora_1 = gemm_nk( - a=grad_gate_lora_intermediate, b=gate_proj_lora_A_compute, - cumsum_M=cumsum, max_M=max_M, transpose_b=True, + a=grad_gate_lora_intermediate, + b=gate_proj_lora_A_compute, + cumsum_M=cumsum, + max_M=max_M, + transpose_b=True, ) grad_permute_tokens_1 = grad_permute_tokens_1 + grad_permute_from_lora_1 @@ -314,6 +447,7 @@ def backward(ctx, grad_output): # === Reduce gradients for shared weights === from xorl.distributed.parallel_state import get_parallel_state + ep_group = get_parallel_state().ep_group if gate_A_shared: @@ -337,10 +471,15 @@ def backward(ctx, grad_output): return ( grad_permute_tokens, None, # cumsum - None, None, None, # base weights (frozen) - grad_gate_proj_lora_A, grad_gate_proj_lora_B, - grad_up_proj_lora_A, grad_up_proj_lora_B, - grad_down_proj_lora_A, grad_down_proj_lora_B, + None, + None, + None, # base weights (frozen) + grad_gate_proj_lora_A, + grad_gate_proj_lora_B, + grad_up_proj_lora_A, + grad_up_proj_lora_B, + grad_down_proj_lora_A, + grad_down_proj_lora_B, None, # scaling ) @@ -442,8 +581,12 @@ def forward( # MOE Step 4 & 6: base GEMM (concat on dim=2 for (G,K,2N)) gate_up_weight = torch.cat([gate_proj, up_proj], dim=2) gate_up_output = gemm_nk( - a=scatter_output, b=gate_up_weight, - cumsum_M=cumsum_t, max_M=max_M, transpose_a=False, transpose_b=False, + a=scatter_output, + b=gate_up_weight, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) gate_base = gate_up_output[:, :intermediate_size] up_base = gate_up_output[:, intermediate_size:] @@ -452,23 +595,35 @@ def forward( lora_r = gate_proj_lora_A.shape[2] gate_up_lora_A = torch.cat([gate_proj_lora_A, up_proj_lora_A], dim=2) gate_up_lora_intermediate = gemm_nk( - a=scatter_output, b=gate_up_lora_A, - cumsum_M=cumsum_t, max_M=max_M, transpose_a=False, transpose_b=False, + a=scatter_output, + b=gate_up_lora_A, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) gate_lora_intermediate = gate_up_lora_intermediate[:, :lora_r].contiguous() up_lora_intermediate = gate_up_lora_intermediate[:, lora_r:].contiguous() # LoRA B for gate: z @ B gate_lora_output = gemm_nk( - a=gate_lora_intermediate, b=gate_proj_lora_B, - cumsum_M=cumsum_t, max_M=max_M, transpose_a=False, transpose_b=False, + a=gate_lora_intermediate, + b=gate_proj_lora_B, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) gate_output = (gate_base + gate_lora_output * scaling).contiguous() # LoRA B for up: z @ B up_lora_output = gemm_nk( - a=up_lora_intermediate, b=up_proj_lora_B, - cumsum_M=cumsum_t, max_M=max_M, transpose_a=False, transpose_b=False, + a=up_lora_intermediate, + b=up_proj_lora_B, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) up_output = (up_base + up_lora_output * scaling).contiguous() @@ -484,18 +639,30 @@ def forward( # down_proj base: h @ W down_output = gemm_nk( - a=gated_weighted, b=down_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_a=False, transpose_b=False, + a=gated_weighted, + b=down_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) # down_proj LoRA: (h @ A) @ B * scaling down_lora_intermediate = gemm_nk( - a=gated_weighted, b=down_proj_lora_A, - cumsum_M=cumsum_t, max_M=max_M, transpose_a=False, transpose_b=False, + a=gated_weighted, + b=down_proj_lora_A, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) down_lora_output = gemm_nk( - a=down_lora_intermediate, b=down_proj_lora_B, - cumsum_M=cumsum_t, max_M=max_M, transpose_a=False, transpose_b=False, + a=down_lora_intermediate, + b=down_proj_lora_B, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) down_output = down_output + down_lora_output * scaling @@ -511,14 +678,27 @@ def forward( ctx.down_B_shared = down_B_shared ctx.save_for_backward( gate_weights, - gate_proj, up_proj, down_proj, - orig_gate_proj_lora_A, orig_gate_proj_lora_B, - orig_up_proj_lora_A, orig_up_proj_lora_B, - orig_down_proj_lora_A, orig_down_proj_lora_B, - hidden_states, scatter_index, scatter_output, cumsum_t, - gate_output, up_output, gated_activation, - scattered_gate_weight, gated_weighted, - gate_lora_intermediate, up_lora_intermediate, down_lora_intermediate, + gate_proj, + up_proj, + down_proj, + orig_gate_proj_lora_A, + orig_gate_proj_lora_B, + orig_up_proj_lora_A, + orig_up_proj_lora_B, + orig_down_proj_lora_A, + orig_down_proj_lora_B, + hidden_states, + scatter_index, + scatter_output, + cumsum_t, + gate_output, + up_output, + gated_activation, + scattered_gate_weight, + gated_weighted, + gate_lora_intermediate, + up_lora_intermediate, + down_lora_intermediate, ) return output @@ -527,14 +707,27 @@ def forward( def backward(ctx, grad_output): ( gate_weights, - gate_proj, up_proj, down_proj, - gate_proj_lora_A, gate_proj_lora_B, - up_proj_lora_A, up_proj_lora_B, - down_proj_lora_A, down_proj_lora_B, - hidden_states, scatter_index, scatter_output, cumsum_t, - gate_output, up_output, gated_activation, - scattered_gate_weight, gated_weighted, - gate_lora_intermediate, up_lora_intermediate, down_lora_intermediate, + gate_proj, + up_proj, + down_proj, + gate_proj_lora_A, + gate_proj_lora_B, + up_proj_lora_A, + up_proj_lora_B, + down_proj_lora_A, + down_proj_lora_B, + hidden_states, + scatter_index, + scatter_output, + cumsum_t, + gate_output, + up_output, + gated_activation, + scattered_gate_weight, + gated_weighted, + gate_lora_intermediate, + up_lora_intermediate, + down_lora_intermediate, ) = ctx.saved_tensors num_experts = ctx.num_experts @@ -548,12 +741,24 @@ def backward(ctx, grad_output): max_M = scatter_output.shape[0] # Cast LoRA weights to compute dtype - gate_proj_lora_A_c = gate_proj_lora_A.to(compute_dtype) if gate_proj_lora_A.dtype != compute_dtype else gate_proj_lora_A - gate_proj_lora_B_c = gate_proj_lora_B.to(compute_dtype) if gate_proj_lora_B.dtype != compute_dtype else gate_proj_lora_B - up_proj_lora_A_c = up_proj_lora_A.to(compute_dtype) if up_proj_lora_A.dtype != compute_dtype else up_proj_lora_A - up_proj_lora_B_c = up_proj_lora_B.to(compute_dtype) if up_proj_lora_B.dtype != compute_dtype else up_proj_lora_B - down_proj_lora_A_c = down_proj_lora_A.to(compute_dtype) if down_proj_lora_A.dtype != compute_dtype else down_proj_lora_A - down_proj_lora_B_c = down_proj_lora_B.to(compute_dtype) if down_proj_lora_B.dtype != compute_dtype else down_proj_lora_B + gate_proj_lora_A_c = ( + gate_proj_lora_A.to(compute_dtype) if gate_proj_lora_A.dtype != compute_dtype else gate_proj_lora_A + ) + gate_proj_lora_B_c = ( + gate_proj_lora_B.to(compute_dtype) if gate_proj_lora_B.dtype != compute_dtype else gate_proj_lora_B + ) + up_proj_lora_A_c = ( + up_proj_lora_A.to(compute_dtype) if up_proj_lora_A.dtype != compute_dtype else up_proj_lora_A + ) + up_proj_lora_B_c = ( + up_proj_lora_B.to(compute_dtype) if up_proj_lora_B.dtype != compute_dtype else up_proj_lora_B + ) + down_proj_lora_A_c = ( + down_proj_lora_A.to(compute_dtype) if down_proj_lora_A.dtype != compute_dtype else down_proj_lora_A + ) + down_proj_lora_B_c = ( + down_proj_lora_B.to(compute_dtype) if down_proj_lora_B.dtype != compute_dtype else down_proj_lora_B + ) # Expand shared LoRA weights for backward compute if gate_A_shared: @@ -565,19 +770,28 @@ def backward(ctx, grad_output): # Initialize LoRA gradients — use full-size buffers for shared weights grad_gate_proj_lora_A_full = torch.zeros( - num_experts, gate_proj_lora_A.shape[1], gate_proj_lora_A.shape[2], - dtype=gate_proj_lora_A.dtype, device=gate_proj_lora_A.device, + num_experts, + gate_proj_lora_A.shape[1], + gate_proj_lora_A.shape[2], + dtype=gate_proj_lora_A.dtype, + device=gate_proj_lora_A.device, ) grad_gate_proj_lora_B = torch.zeros_like(gate_proj_lora_B) grad_up_proj_lora_A_full = torch.zeros( - num_experts, up_proj_lora_A.shape[1], up_proj_lora_A.shape[2], - dtype=up_proj_lora_A.dtype, device=up_proj_lora_A.device, + num_experts, + up_proj_lora_A.shape[1], + up_proj_lora_A.shape[2], + dtype=up_proj_lora_A.dtype, + device=up_proj_lora_A.device, ) grad_up_proj_lora_B = torch.zeros_like(up_proj_lora_B) grad_down_proj_lora_A = torch.zeros_like(down_proj_lora_A) grad_down_proj_lora_B_full = torch.zeros( - num_experts, down_proj_lora_B.shape[1], down_proj_lora_B.shape[2], - dtype=down_proj_lora_B.dtype, device=down_proj_lora_B.device, + num_experts, + down_proj_lora_B.shape[1], + down_proj_lora_B.shape[2], + dtype=down_proj_lora_B.dtype, + device=down_proj_lora_B.device, ) # MOE Step 10': scatter grad @@ -585,32 +799,49 @@ def backward(ctx, grad_output): # ====== down_proj backward ====== grad_gated_weighted = gemm_nk( - a=grad_down_output, b=down_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, + a=grad_down_output, + b=down_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_b=True, ) gemm_mn( - a=down_lora_intermediate, b=grad_down_output, - c=grad_down_proj_lora_B_full, cumsum_K=cumsum_t, max_K=max_M, - transpose_a=True, transpose_b=False, + a=down_lora_intermediate, + b=grad_down_output, + c=grad_down_proj_lora_B_full, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) grad_down_proj_lora_B_full.mul_(scaling) grad_down_lora_intermediate = gemm_nk( - a=grad_down_output, b=down_proj_lora_B_c, - cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, + a=grad_down_output, + b=down_proj_lora_B_c, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_b=True, ) grad_down_lora_intermediate.mul_(scaling) gemm_mn( - a=gated_weighted, b=grad_down_lora_intermediate, - c=grad_down_proj_lora_A, cumsum_K=cumsum_t, max_K=max_M, - transpose_a=True, transpose_b=False, + a=gated_weighted, + b=grad_down_lora_intermediate, + c=grad_down_proj_lora_A, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) grad_gated_weighted_from_lora = gemm_nk( - a=grad_down_lora_intermediate, b=down_proj_lora_A_c, - cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, + a=grad_down_lora_intermediate, + b=down_proj_lora_A_c, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_b=True, ) grad_gated_weighted = grad_gated_weighted + grad_gated_weighted_from_lora @@ -626,32 +857,49 @@ def backward(ctx, grad_output): # ====== up_proj backward ====== grad_scatter_output_2 = gemm_nk( - a=grad_up_output, b=up_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, + a=grad_up_output, + b=up_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_b=True, ) gemm_mn( - a=up_lora_intermediate, b=grad_up_output, - c=grad_up_proj_lora_B, cumsum_K=cumsum_t, max_K=max_M, - transpose_a=True, transpose_b=False, + a=up_lora_intermediate, + b=grad_up_output, + c=grad_up_proj_lora_B, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) grad_up_proj_lora_B.mul_(scaling) grad_up_lora_intermediate = gemm_nk( - a=grad_up_output, b=up_proj_lora_B_c, - cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, + a=grad_up_output, + b=up_proj_lora_B_c, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_b=True, ) grad_up_lora_intermediate.mul_(scaling) gemm_mn( - a=scatter_output, b=grad_up_lora_intermediate, - c=grad_up_proj_lora_A_full, cumsum_K=cumsum_t, max_K=max_M, - transpose_a=True, transpose_b=False, + a=scatter_output, + b=grad_up_lora_intermediate, + c=grad_up_proj_lora_A_full, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) grad_scatter_from_lora_2 = gemm_nk( - a=grad_up_lora_intermediate, b=up_proj_lora_A_c, - cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, + a=grad_up_lora_intermediate, + b=up_proj_lora_A_c, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_b=True, ) grad_scatter_output_2 = grad_scatter_output_2 + grad_scatter_from_lora_2 @@ -660,32 +908,49 @@ def backward(ctx, grad_output): # ====== gate_proj backward ====== grad_scatter_output_1 = gemm_nk( - a=grad_gate_output, b=gate_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, + a=grad_gate_output, + b=gate_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_b=True, ) gemm_mn( - a=gate_lora_intermediate, b=grad_gate_output, - c=grad_gate_proj_lora_B, cumsum_K=cumsum_t, max_K=max_M, - transpose_a=True, transpose_b=False, + a=gate_lora_intermediate, + b=grad_gate_output, + c=grad_gate_proj_lora_B, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) grad_gate_proj_lora_B.mul_(scaling) grad_gate_lora_intermediate = gemm_nk( - a=grad_gate_output, b=gate_proj_lora_B_c, - cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, + a=grad_gate_output, + b=gate_proj_lora_B_c, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_b=True, ) grad_gate_lora_intermediate.mul_(scaling) gemm_mn( - a=scatter_output, b=grad_gate_lora_intermediate, - c=grad_gate_proj_lora_A_full, cumsum_K=cumsum_t, max_K=max_M, - transpose_a=True, transpose_b=False, + a=scatter_output, + b=grad_gate_lora_intermediate, + c=grad_gate_proj_lora_A_full, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) grad_scatter_from_lora_1 = gemm_nk( - a=grad_gate_lora_intermediate, b=gate_proj_lora_A_c, - cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, + a=grad_gate_lora_intermediate, + b=gate_proj_lora_A_c, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_b=True, ) grad_scatter_output_1 = grad_scatter_output_1 + grad_scatter_from_lora_1 @@ -714,10 +979,15 @@ def backward(ctx, grad_output): grad_gate_weight, None, # expert_index grad_hidden_states, - None, None, None, # base weights (frozen) - grad_gate_proj_lora_A, grad_gate_proj_lora_B, - grad_up_proj_lora_A, grad_up_proj_lora_B, - grad_down_proj_lora_A, grad_down_proj_lora_B, + None, + None, + None, # base weights (frozen) + grad_gate_proj_lora_A, + grad_gate_proj_lora_B, + grad_up_proj_lora_A, + grad_up_proj_lora_B, + grad_down_proj_lora_A, + grad_down_proj_lora_B, None, # scaling ) diff --git a/src/xorl/ops/moe/quack.py b/src/xorl/ops/moe/quack.py index 8a3d8f57..7128e8f2 100644 --- a/src/xorl/ops/moe/quack.py +++ b/src/xorl/ops/moe/quack.py @@ -81,8 +81,14 @@ def backward(ctx, grad_output): if down_proj.requires_grad: grad_down_proj = torch.empty_like(down_proj) quack_group_gemm_same_mn( - a=gated_output, b=grad_output, c=grad_down_proj, cumsum_K=cumsum, - max_K=max_M, transpose_a=True, transpose_b=False, cu_seqlens_k=cu_seqlens_m + a=gated_output, + b=grad_output, + c=grad_down_proj, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, + cu_seqlens_k=cu_seqlens_m, ) del gated_output @@ -90,9 +96,7 @@ def backward(ctx, grad_output): grad_up_output = gate_activation * grad_gated_output grad_gate_activation = grad_gated_output * up_output del grad_gated_output, gate_activation, up_output - grad_gate_output = torch.ops.aten.silu_backward( - grad_gate_activation, gate_output - ) + grad_gate_output = torch.ops.aten.silu_backward(grad_gate_activation, gate_output) del grad_gate_activation, gate_output # dgrad FC1: in-place add @@ -108,16 +112,28 @@ def backward(ctx, grad_output): if gate_proj.requires_grad: grad_gate_proj = torch.empty_like(gate_proj) quack_group_gemm_same_mn( - a=permute_tokens, b=grad_gate_output, c=grad_gate_proj, cumsum_K=cumsum, - max_K=max_M, transpose_a=True, transpose_b=False, cu_seqlens_k=cu_seqlens_m + a=permute_tokens, + b=grad_gate_output, + c=grad_gate_proj, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, + cu_seqlens_k=cu_seqlens_m, ) del grad_gate_output grad_up_proj = None if up_proj.requires_grad: grad_up_proj = torch.empty_like(up_proj) quack_group_gemm_same_mn( - a=permute_tokens, b=grad_up_output, c=grad_up_proj, cumsum_K=cumsum, - max_K=max_M, transpose_a=True, transpose_b=False, cu_seqlens_k=cu_seqlens_m + a=permute_tokens, + b=grad_up_output, + c=grad_up_proj, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, + cu_seqlens_k=cu_seqlens_m, ) del grad_up_output @@ -129,7 +145,9 @@ class QuackMoeExpertsFunction(torch.autograd.Function): explicit del for dead tensors, in-place add for dgrad.""" @staticmethod - def forward(ctx, num_experts, gate_weights, expert_index, hidden_states, gate_proj, up_proj, down_proj, gate_up_weight=None): + def forward( + ctx, num_experts, gate_weights, expert_index, hidden_states, gate_proj, up_proj, down_proj, gate_up_weight=None + ): scatter_output, scatter_index, cumsum_t = _scatter_and_cumsum(hidden_states, expert_index, num_experts) max_M = scatter_output.shape[0] cu_seqlens = cumsum_to_cu_seqlens(cumsum_t) @@ -159,18 +177,32 @@ def forward(ctx, num_experts, gate_weights, expert_index, hidden_states, gate_pr del down_output ctx.save_for_backward( - gate_weights, gate_proj, up_proj, down_proj, - hidden_states, scatter_index, cumsum_t, - gate_output, up_output, scattered_gate_weight, + gate_weights, + gate_proj, + up_proj, + down_proj, + hidden_states, + scatter_index, + cumsum_t, + gate_output, + up_output, + scattered_gate_weight, ) return output @staticmethod def backward(ctx, grad_output): ( - gate_weights, gate_proj, up_proj, down_proj, - hidden_states, scatter_index, cumsum_t, - gate_output, up_output, scattered_gate_weight, + gate_weights, + gate_proj, + up_proj, + down_proj, + hidden_states, + scatter_index, + cumsum_t, + gate_output, + up_output, + scattered_gate_weight, ) = ctx.saved_tensors grad_output = grad_output.view(-1, grad_output.shape[-1]) max_M = grad_output.shape[0] @@ -194,8 +226,14 @@ def backward(ctx, grad_output): if down_proj.requires_grad: grad_down_proj = torch.empty_like(down_proj) quack_group_gemm_same_mn( - a=gated_weighted, b=grad_down_output, c=grad_down_proj, cumsum_K=cumsum_t, - max_K=max_M, transpose_a=True, transpose_b=False, cu_seqlens_k=cu_seqlens_m + a=gated_weighted, + b=grad_down_output, + c=grad_down_proj, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, + cu_seqlens_k=cu_seqlens_m, ) del grad_down_output, gated_weighted @@ -209,9 +247,7 @@ def backward(ctx, grad_output): grad_up_output = gate_activation * grad_gated_activation grad_gate_activation = grad_gated_activation * up_output del grad_gated_activation, gate_activation, up_output - grad_gate_output = torch.ops.aten.silu_backward( - grad_gate_activation, gate_output - ) + grad_gate_output = torch.ops.aten.silu_backward(grad_gate_activation, gate_output) del grad_gate_activation, gate_output # dgrad FC1: in-place add @@ -227,16 +263,28 @@ def backward(ctx, grad_output): if gate_proj.requires_grad: grad_gate_proj = torch.empty_like(gate_proj) quack_group_gemm_same_mn( - a=scatter_output, b=grad_gate_output, c=grad_gate_proj, cumsum_K=cumsum_t, - max_K=max_M, transpose_a=True, transpose_b=False, cu_seqlens_k=cu_seqlens_m + a=scatter_output, + b=grad_gate_output, + c=grad_gate_proj, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, + cu_seqlens_k=cu_seqlens_m, ) del grad_gate_output grad_up_proj = None if up_proj.requires_grad: grad_up_proj = torch.empty_like(up_proj) quack_group_gemm_same_mn( - a=scatter_output, b=grad_up_output, c=grad_up_proj, cumsum_K=cumsum_t, - max_K=max_M, transpose_a=True, transpose_b=False, cu_seqlens_k=cu_seqlens_m + a=scatter_output, + b=grad_up_output, + c=grad_up_proj, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, + cu_seqlens_k=cu_seqlens_m, ) del grad_up_output, scatter_output @@ -254,8 +302,14 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): if _DEBUG_EP: return QuackEPGroupGemm._forward_debug( - ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, - max_M, cu_seqlens, + ctx, + permute_tokens, + cumsum, + gate_proj, + up_proj, + down_proj, + max_M, + cu_seqlens, ) gate_output = quack_group_gemm_same_nk( @@ -346,8 +400,14 @@ def backward(ctx, grad_output): if down_proj.requires_grad: grad_down_proj = torch.empty_like(down_proj) quack_group_gemm_same_mn( - a=gated_output, b=grad_output, c=grad_down_proj, cumsum_K=cumsum, - max_K=max_M, transpose_a=True, transpose_b=False, cu_seqlens_k=cu_seqlens_m + a=gated_output, + b=grad_output, + c=grad_down_proj, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, + cu_seqlens_k=cu_seqlens_m, ) del gated_output @@ -355,9 +415,7 @@ def backward(ctx, grad_output): grad_up_output = gate_activation * grad_gated_output grad_gate_activation = grad_gated_output * up_output del grad_gated_output, gate_activation, up_output - grad_gate_output = torch.ops.aten.silu_backward( - grad_gate_activation, gate_output - ) + grad_gate_output = torch.ops.aten.silu_backward(grad_gate_activation, gate_output) del grad_gate_activation, gate_output # dgrad FC1: in-place add @@ -373,16 +431,28 @@ def backward(ctx, grad_output): if gate_proj.requires_grad: grad_gate_proj = torch.empty_like(gate_proj) quack_group_gemm_same_mn( - a=permute_tokens, b=grad_gate_output, c=grad_gate_proj, cumsum_K=cumsum, - max_K=max_M, transpose_a=True, transpose_b=False, cu_seqlens_k=cu_seqlens_m + a=permute_tokens, + b=grad_gate_output, + c=grad_gate_proj, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, + cu_seqlens_k=cu_seqlens_m, ) del grad_gate_output grad_up_proj = None if up_proj.requires_grad: grad_up_proj = torch.empty_like(up_proj) quack_group_gemm_same_mn( - a=permute_tokens, b=grad_up_output, c=grad_up_proj, cumsum_K=cumsum, - max_K=max_M, transpose_a=True, transpose_b=False, cu_seqlens_k=cu_seqlens_m + a=permute_tokens, + b=grad_up_output, + c=grad_up_proj, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, + cu_seqlens_k=cu_seqlens_m, ) del grad_up_output @@ -425,18 +495,32 @@ def forward(ctx, num_experts, gate_weights, expert_index, hidden_states, gate_pr ctx.tp_group = tp_group ctx.save_for_backward( - gate_weights, gate_proj, up_proj, down_proj, - hidden_states, scatter_index, cumsum_t, - gate_output, up_output, scattered_gate_weight, + gate_weights, + gate_proj, + up_proj, + down_proj, + hidden_states, + scatter_index, + cumsum_t, + gate_output, + up_output, + scattered_gate_weight, ) return output @staticmethod def backward(ctx, grad_output): ( - gate_weights, gate_proj, up_proj, down_proj, - hidden_states, scatter_index, cumsum_t, - gate_output, up_output, scattered_gate_weight, + gate_weights, + gate_proj, + up_proj, + down_proj, + hidden_states, + scatter_index, + cumsum_t, + gate_output, + up_output, + scattered_gate_weight, ) = ctx.saved_tensors tp_group = ctx.tp_group grad_output = grad_output.view(-1, grad_output.shape[-1]) @@ -461,8 +545,14 @@ def backward(ctx, grad_output): if down_proj.requires_grad: grad_down_proj = torch.empty_like(down_proj) quack_group_gemm_same_mn( - a=gated_weighted, b=grad_down_output, c=grad_down_proj, cumsum_K=cumsum_t, - max_K=max_M, transpose_a=True, transpose_b=False, cu_seqlens_k=cu_seqlens_m + a=gated_weighted, + b=grad_down_output, + c=grad_down_proj, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, + cu_seqlens_k=cu_seqlens_m, ) del grad_down_output, gated_weighted @@ -476,9 +566,7 @@ def backward(ctx, grad_output): grad_up_output = gate_activation * grad_gated_activation grad_gate_activation = grad_gated_activation * up_output del grad_gated_activation, gate_activation, up_output - grad_gate_output = torch.ops.aten.silu_backward( - grad_gate_activation, gate_output - ) + grad_gate_output = torch.ops.aten.silu_backward(grad_gate_activation, gate_output) del grad_gate_activation, gate_output # dgrad FC1: in-place add @@ -495,16 +583,28 @@ def backward(ctx, grad_output): if gate_proj.requires_grad: grad_gate_proj = torch.empty_like(gate_proj) quack_group_gemm_same_mn( - a=scatter_output, b=grad_gate_output, c=grad_gate_proj, cumsum_K=cumsum_t, - max_K=max_M, transpose_a=True, transpose_b=False, cu_seqlens_k=cu_seqlens_m + a=scatter_output, + b=grad_gate_output, + c=grad_gate_proj, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, + cu_seqlens_k=cu_seqlens_m, ) del grad_gate_output grad_up_proj = None if up_proj.requires_grad: grad_up_proj = torch.empty_like(up_proj) quack_group_gemm_same_mn( - a=scatter_output, b=grad_up_output, c=grad_up_proj, cumsum_K=cumsum_t, - max_K=max_M, transpose_a=True, transpose_b=False, cu_seqlens_k=cu_seqlens_m + a=scatter_output, + b=grad_up_output, + c=grad_up_proj, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, + cu_seqlens_k=cu_seqlens_m, ) del grad_up_output, scatter_output @@ -548,7 +648,9 @@ class QuackMoeExpertsFunctionMoeAct(torch.autograd.Function): from save_for_backward and recomputes them in backward.""" @staticmethod - def forward(ctx, num_experts, gate_weights, expert_index, hidden_states, gate_proj, up_proj, down_proj, gate_up_weight=None): + def forward( + ctx, num_experts, gate_weights, expert_index, hidden_states, gate_proj, up_proj, down_proj, gate_up_weight=None + ): scatter_output, scatter_index, cumsum_t = _scatter_and_cumsum(hidden_states, expert_index, num_experts) max_M = scatter_output.shape[0] cu_seqlens = cumsum_to_cu_seqlens(cumsum_t) @@ -579,8 +681,13 @@ def forward(ctx, num_experts, gate_weights, expert_index, hidden_states, gate_pr # moe_act: save 8 tensors (drop gate_output, up_output vs 10 in standard) ctx.save_for_backward( - gate_weights, gate_proj, up_proj, down_proj, - hidden_states, scatter_index, cumsum_t, + gate_weights, + gate_proj, + up_proj, + down_proj, + hidden_states, + scatter_index, + cumsum_t, scattered_gate_weight, ) return output @@ -588,8 +695,13 @@ def forward(ctx, num_experts, gate_weights, expert_index, hidden_states, gate_pr @staticmethod def backward(ctx, grad_output): ( - gate_weights, gate_proj, up_proj, down_proj, - hidden_states, scatter_index, cumsum_t, + gate_weights, + gate_proj, + up_proj, + down_proj, + hidden_states, + scatter_index, + cumsum_t, scattered_gate_weight, ) = ctx.saved_tensors grad_output = grad_output.view(-1, grad_output.shape[-1]) @@ -621,8 +733,14 @@ def backward(ctx, grad_output): if down_proj.requires_grad: grad_down_proj = torch.empty_like(down_proj) quack_group_gemm_same_mn( - a=gated_weighted, b=grad_down_output, c=grad_down_proj, cumsum_K=cumsum_t, - max_K=max_M, transpose_a=True, transpose_b=False, cu_seqlens_k=cu_seqlens_m + a=gated_weighted, + b=grad_down_output, + c=grad_down_proj, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, + cu_seqlens_k=cu_seqlens_m, ) del grad_down_output, gated_weighted @@ -636,9 +754,7 @@ def backward(ctx, grad_output): grad_up_output = gate_activation * grad_gated_activation grad_gate_activation = grad_gated_activation * up_output del grad_gated_activation, gate_activation, up_output - grad_gate_output = torch.ops.aten.silu_backward( - grad_gate_activation, gate_output - ) + grad_gate_output = torch.ops.aten.silu_backward(grad_gate_activation, gate_output) del grad_gate_activation, gate_output # dgrad FC1: in-place add @@ -654,16 +770,28 @@ def backward(ctx, grad_output): if gate_proj.requires_grad: grad_gate_proj = torch.empty_like(gate_proj) quack_group_gemm_same_mn( - a=scatter_output, b=grad_gate_output, c=grad_gate_proj, cumsum_K=cumsum_t, - max_K=max_M, transpose_a=True, transpose_b=False, cu_seqlens_k=cu_seqlens_m + a=scatter_output, + b=grad_gate_output, + c=grad_gate_proj, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, + cu_seqlens_k=cu_seqlens_m, ) del grad_gate_output grad_up_proj = None if up_proj.requires_grad: grad_up_proj = torch.empty_like(up_proj) quack_group_gemm_same_mn( - a=scatter_output, b=grad_up_output, c=grad_up_proj, cumsum_K=cumsum_t, - max_K=max_M, transpose_a=True, transpose_b=False, cu_seqlens_k=cu_seqlens_m + a=scatter_output, + b=grad_up_output, + c=grad_up_proj, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, + cu_seqlens_k=cu_seqlens_m, ) del grad_up_output, scatter_output diff --git a/src/xorl/ops/moe/quack_lora.py b/src/xorl/ops/moe/quack_lora.py index d68ce46c..a112a4c6 100644 --- a/src/xorl/ops/moe/quack_lora.py +++ b/src/xorl/ops/moe/quack_lora.py @@ -6,8 +6,10 @@ """ from xorl.ops.group_gemm.kernel.quack import quack_group_gemm_same_mn, quack_group_gemm_same_nk + from .lora import make_ep_lora_compute, make_local_lora_compute + # EP LoRA compute (Expert Parallelism) QuackEPGroupGemmWithLoRA = make_ep_lora_compute(quack_group_gemm_same_nk, quack_group_gemm_same_mn) QuackEPGroupGemmWithLoRA.__name__ = QuackEPGroupGemmWithLoRA.__qualname__ = "QuackEPGroupGemmWithLoRA" @@ -35,8 +37,18 @@ def quack_moe_lora_forward( ): """MoE + LoRA forward pass using quack group GEMM kernels (local single-GPU).""" return QuackMoeExpertsLoRAFunction.apply( - num_experts, routing_weights.to(hidden_states.dtype), selected_experts, hidden_states, - gate_proj, up_proj, down_proj, - gate_proj_lora_A, gate_proj_lora_B, up_proj_lora_A, up_proj_lora_B, - down_proj_lora_A, down_proj_lora_B, scaling, + num_experts, + routing_weights.to(hidden_states.dtype), + selected_experts, + hidden_states, + gate_proj, + up_proj, + down_proj, + gate_proj_lora_A, + gate_proj_lora_B, + up_proj_lora_A, + up_proj_lora_B, + down_proj_lora_A, + down_proj_lora_B, + scaling, ) diff --git a/src/xorl/ops/moe/triton.py b/src/xorl/ops/moe/triton.py index e3dcd8e3..6db8c8a5 100644 --- a/src/xorl/ops/moe/triton.py +++ b/src/xorl/ops/moe/triton.py @@ -1,5 +1,3 @@ -import os - import torch from xorl.utils.import_utils import is_fused_moe_available @@ -7,8 +5,14 @@ if is_fused_moe_available(): from xorl.ops.group_gemm.kernel.group_gemm import group_gemm_same_mn, group_gemm_same_nk - from xorl.ops.group_gemm.kernel.moe import expert_histogram, moe_add_gather, moe_gather, moe_index_compute, moe_scatter - from xorl.ops.fused_silu_and_mul import fused_silu_and_mul, silu_and_mul_backward + from xorl.ops.group_gemm.kernel.moe import ( + expert_histogram, + moe_gather, + moe_index_compute, + moe_scatter, + ) + + class TritonEPGroupGemm(torch.autograd.Function): """EP expert MLP computation using Triton group GEMM. @@ -25,12 +29,20 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): max_M = permute_tokens.shape[0] gate_output = group_gemm_same_nk( - a=permute_tokens, b=gate_proj, - cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, + a=permute_tokens, + b=gate_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) up_output = group_gemm_same_nk( - a=permute_tokens, b=up_proj, - cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, + a=permute_tokens, + b=up_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) gate_activation = torch.ops.aten.silu(gate_output) @@ -38,8 +50,12 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): del gate_activation down_output = group_gemm_same_nk( - a=gated_output, b=down_proj, - cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, + a=gated_output, + b=down_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) del gated_output @@ -49,8 +65,13 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): @staticmethod def backward(ctx, grad_output): ( - permute_tokens, cumsum, gate_proj, up_proj, down_proj, - gate_output, up_output, + permute_tokens, + cumsum, + gate_proj, + up_proj, + down_proj, + gate_output, + up_output, ) = ctx.saved_tensors max_M = grad_output.shape[0] @@ -60,8 +81,11 @@ def backward(ctx, grad_output): # dgrad FC2 grad_gated_output = group_gemm_same_nk( - a=grad_output, b=down_proj, - cumsum_M=cumsum, max_M=max_M, transpose_b=True, + a=grad_output, + b=down_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_b=True, ) # wgrad FC2 @@ -69,8 +93,13 @@ def backward(ctx, grad_output): if down_proj.requires_grad: grad_down_proj = torch.empty_like(down_proj) group_gemm_same_mn( - a=gated_output, b=grad_output, c=grad_down_proj, - cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, + a=gated_output, + b=grad_output, + c=grad_down_proj, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) del gated_output @@ -80,26 +109,35 @@ def backward(ctx, grad_output): del grad_gated_output, gate_activation, up_output grad_scatter_output_2 = group_gemm_same_nk( - a=grad_up_output, b=up_proj, - cumsum_M=cumsum, max_M=max_M, transpose_b=True, + a=grad_up_output, + b=up_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_b=True, ) grad_up_proj = None if up_proj.requires_grad: grad_up_proj = torch.empty_like(up_proj) group_gemm_same_mn( - a=permute_tokens, b=grad_up_output, c=grad_up_proj, - cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, + a=permute_tokens, + b=grad_up_output, + c=grad_up_proj, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) - grad_gate_output = torch.ops.aten.silu_backward( - grad_gate_activation, gate_output - ) + grad_gate_output = torch.ops.aten.silu_backward(grad_gate_activation, gate_output) del grad_gate_activation, gate_output grad_permute_tokens = group_gemm_same_nk( - a=grad_gate_output, b=gate_proj, - cumsum_M=cumsum, max_M=max_M, transpose_b=True, + a=grad_gate_output, + b=gate_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_b=True, ) grad_permute_tokens += grad_scatter_output_2 del grad_scatter_output_2 @@ -108,8 +146,13 @@ def backward(ctx, grad_output): if gate_proj.requires_grad: grad_gate_proj = torch.empty_like(gate_proj) group_gemm_same_mn( - a=permute_tokens, b=grad_gate_output, c=grad_gate_proj, - cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, + a=permute_tokens, + b=grad_gate_output, + c=grad_gate_proj, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) del grad_gate_output @@ -134,12 +177,20 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): max_M = permute_tokens.shape[0] gate_output = group_gemm_same_nk( - a=permute_tokens, b=gate_proj, - cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, + a=permute_tokens, + b=gate_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) up_output = group_gemm_same_nk( - a=permute_tokens, b=up_proj, - cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, + a=permute_tokens, + b=up_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) gate_activation = torch.ops.aten.silu(gate_output) @@ -147,8 +198,12 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): del gate_activation down_output = group_gemm_same_nk( - a=gated_output, b=down_proj, - cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, + a=gated_output, + b=down_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) del gated_output @@ -163,12 +218,20 @@ def backward(ctx, grad_output): # Recompute gate_output and up_output from saved inputs + weights gate_output = group_gemm_same_nk( - a=permute_tokens, b=gate_proj, - cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, + a=permute_tokens, + b=gate_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) up_output = group_gemm_same_nk( - a=permute_tokens, b=up_proj, - cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, + a=permute_tokens, + b=up_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) # Rest identical to TritonEPGroupGemm.backward @@ -177,8 +240,11 @@ def backward(ctx, grad_output): # dgrad FC2 grad_gated_output = group_gemm_same_nk( - a=grad_output, b=down_proj, - cumsum_M=cumsum, max_M=max_M, transpose_b=True, + a=grad_output, + b=down_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_b=True, ) # wgrad FC2 @@ -186,8 +252,13 @@ def backward(ctx, grad_output): if down_proj.requires_grad: grad_down_proj = torch.empty_like(down_proj) group_gemm_same_mn( - a=gated_output, b=grad_output, c=grad_down_proj, - cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, + a=gated_output, + b=grad_output, + c=grad_down_proj, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) del gated_output @@ -197,26 +268,35 @@ def backward(ctx, grad_output): del grad_gated_output, gate_activation, up_output grad_scatter_output_2 = group_gemm_same_nk( - a=grad_up_output, b=up_proj, - cumsum_M=cumsum, max_M=max_M, transpose_b=True, + a=grad_up_output, + b=up_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_b=True, ) grad_up_proj = None if up_proj.requires_grad: grad_up_proj = torch.empty_like(up_proj) group_gemm_same_mn( - a=permute_tokens, b=grad_up_output, c=grad_up_proj, - cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, + a=permute_tokens, + b=grad_up_output, + c=grad_up_proj, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) - grad_gate_output = torch.ops.aten.silu_backward( - grad_gate_activation, gate_output - ) + grad_gate_output = torch.ops.aten.silu_backward(grad_gate_activation, gate_output) del grad_gate_activation, gate_output grad_permute_tokens = group_gemm_same_nk( - a=grad_gate_output, b=gate_proj, - cumsum_M=cumsum, max_M=max_M, transpose_b=True, + a=grad_gate_output, + b=gate_proj, + cumsum_M=cumsum, + max_M=max_M, + transpose_b=True, ) grad_permute_tokens += grad_scatter_output_2 del grad_scatter_output_2 @@ -225,8 +305,13 @@ def backward(ctx, grad_output): if gate_proj.requires_grad: grad_gate_proj = torch.empty_like(gate_proj) group_gemm_same_mn( - a=permute_tokens, b=grad_gate_output, c=grad_gate_proj, - cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, + a=permute_tokens, + b=grad_gate_output, + c=grad_gate_proj, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) del grad_gate_output @@ -268,12 +353,20 @@ def forward( # Separate gate and up GEMMs (avoids allocating concatenated weight tensor) gate_output = group_gemm_same_nk( - a=scatter_output, b=gate_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_a=False, transpose_b=False, + a=scatter_output, + b=gate_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) up_output = group_gemm_same_nk( - a=scatter_output, b=up_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_a=False, transpose_b=False, + a=scatter_output, + b=up_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) del scatter_output @@ -291,8 +384,12 @@ def forward( # Down projection down_output = group_gemm_same_nk( - a=gated_weighted, b=down_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_a=False, transpose_b=False, + a=gated_weighted, + b=down_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) del gated_weighted @@ -303,9 +400,16 @@ def forward( # Save gate_output + up_output for backward (cheap intermediates like # scatter_output, gate_activation, gated_weighted are recomputed instead). ctx.save_for_backward( - gate_weights, gate_proj, up_proj, down_proj, - hidden_states, scatter_index, cumsum_t, - gate_output, up_output, scattered_gate_weight, + gate_weights, + gate_proj, + up_proj, + down_proj, + hidden_states, + scatter_index, + cumsum_t, + gate_output, + up_output, + scattered_gate_weight, ) return output @@ -313,9 +417,16 @@ def forward( @staticmethod def backward(ctx, grad_output): ( - gate_weights, gate_proj, up_proj, down_proj, - hidden_states, scatter_index, cumsum_t, - gate_output, up_output, scattered_gate_weight, + gate_weights, + gate_proj, + up_proj, + down_proj, + hidden_states, + scatter_index, + cumsum_t, + gate_output, + up_output, + scattered_gate_weight, ) = ctx.saved_tensors grad_output = grad_output.view(-1, grad_output.shape[-1]) max_M = grad_output.shape[0] @@ -331,8 +442,11 @@ def backward(ctx, grad_output): # FC2 dgrad: grad @ down_proj^T grad_gated_weighted = group_gemm_same_nk( - a=grad_down_output, b=down_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, + a=grad_down_output, + b=down_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_b=True, ) # FC2 wgrad: gated_weighted^T @ grad @@ -340,8 +454,13 @@ def backward(ctx, grad_output): if down_proj.requires_grad: grad_down_proj = torch.empty_like(down_proj) group_gemm_same_mn( - a=gated_weighted, b=grad_down_output, c=grad_down_proj, - cumsum_K=cumsum_t, max_K=max_M, transpose_a=True, transpose_b=False, + a=gated_weighted, + b=grad_down_output, + c=grad_down_proj, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) del grad_down_output, gated_weighted @@ -355,19 +474,23 @@ def backward(ctx, grad_output): grad_up_output = gate_activation * grad_gated_activation grad_gate_activation = grad_gated_activation * up_output del grad_gated_activation, gate_activation, up_output - grad_gate_output = torch.ops.aten.silu_backward( - grad_gate_activation, gate_output - ) + grad_gate_output = torch.ops.aten.silu_backward(grad_gate_activation, gate_output) del grad_gate_activation, gate_output # FC1 dgrad: separate GEMMs, in-place add (avoids allocating sum tensor) grad_scatter_output = group_gemm_same_nk( - a=grad_gate_output, b=gate_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, + a=grad_gate_output, + b=gate_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_b=True, ) grad_scatter_output += group_gemm_same_nk( - a=grad_up_output, b=up_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, + a=grad_up_output, + b=up_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_b=True, ) # FC1 wgrad: separate GEMMs (avoids concatenated grad weight alloc + .contiguous() copies) @@ -375,16 +498,26 @@ def backward(ctx, grad_output): if gate_proj.requires_grad: grad_gate_proj = torch.empty_like(gate_proj) group_gemm_same_mn( - a=scatter_output, b=grad_gate_output, c=grad_gate_proj, - cumsum_K=cumsum_t, max_K=max_M, transpose_a=True, transpose_b=False, + a=scatter_output, + b=grad_gate_output, + c=grad_gate_proj, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) del grad_gate_output grad_up_proj = None if up_proj.requires_grad: grad_up_proj = torch.empty_like(up_proj) group_gemm_same_mn( - a=scatter_output, b=grad_up_output, c=grad_up_proj, - cumsum_K=cumsum_t, max_K=max_M, transpose_a=True, transpose_b=False, + a=scatter_output, + b=grad_up_output, + c=grad_up_proj, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) del grad_up_output, scatter_output @@ -468,12 +601,20 @@ def forward( max_M = scatter_output.shape[0] gate_output = group_gemm_same_nk( - a=scatter_output, b=gate_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_a=False, transpose_b=False, + a=scatter_output, + b=gate_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) up_output = group_gemm_same_nk( - a=scatter_output, b=up_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_a=False, transpose_b=False, + a=scatter_output, + b=up_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) del scatter_output @@ -488,8 +629,12 @@ def forward( del gated_activation down_output = group_gemm_same_nk( - a=gated_weighted, b=down_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_a=False, transpose_b=False, + a=gated_weighted, + b=down_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) del gated_weighted @@ -498,8 +643,13 @@ def forward( # moe_act: save 8 tensors (drop gate_output, up_output vs 10 in standard) ctx.save_for_backward( - gate_weights, gate_proj, up_proj, down_proj, - hidden_states, scatter_index, cumsum_t, + gate_weights, + gate_proj, + up_proj, + down_proj, + hidden_states, + scatter_index, + cumsum_t, scattered_gate_weight, ) @@ -508,8 +658,13 @@ def forward( @staticmethod def backward(ctx, grad_output): ( - gate_weights, gate_proj, up_proj, down_proj, - hidden_states, scatter_index, cumsum_t, + gate_weights, + gate_proj, + up_proj, + down_proj, + hidden_states, + scatter_index, + cumsum_t, scattered_gate_weight, ) = ctx.saved_tensors grad_output = grad_output.view(-1, grad_output.shape[-1]) @@ -518,12 +673,20 @@ def backward(ctx, grad_output): # Recompute scatter_output, gate_output, up_output from saved inputs + weights scatter_output = moe_scatter(hidden_states, scatter_index) gate_output = group_gemm_same_nk( - a=scatter_output, b=gate_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_a=False, transpose_b=False, + a=scatter_output, + b=gate_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) up_output = group_gemm_same_nk( - a=scatter_output, b=up_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_a=False, transpose_b=False, + a=scatter_output, + b=up_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_a=False, + transpose_b=False, ) gate_activation = torch.ops.aten.silu(gate_output) @@ -535,8 +698,11 @@ def backward(ctx, grad_output): # FC2 dgrad grad_gated_weighted = group_gemm_same_nk( - a=grad_down_output, b=down_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, + a=grad_down_output, + b=down_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_b=True, ) # FC2 wgrad @@ -544,8 +710,13 @@ def backward(ctx, grad_output): if down_proj.requires_grad: grad_down_proj = torch.empty_like(down_proj) group_gemm_same_mn( - a=gated_weighted, b=grad_down_output, c=grad_down_proj, - cumsum_K=cumsum_t, max_K=max_M, transpose_a=True, transpose_b=False, + a=gated_weighted, + b=grad_down_output, + c=grad_down_proj, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) del grad_down_output, gated_weighted @@ -559,19 +730,23 @@ def backward(ctx, grad_output): grad_up_output = gate_activation * grad_gated_activation grad_gate_activation = grad_gated_activation * up_output del grad_gated_activation, gate_activation, up_output - grad_gate_output = torch.ops.aten.silu_backward( - grad_gate_activation, gate_output - ) + grad_gate_output = torch.ops.aten.silu_backward(grad_gate_activation, gate_output) del grad_gate_activation, gate_output # FC1 dgrad grad_scatter_output = group_gemm_same_nk( - a=grad_gate_output, b=gate_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, + a=grad_gate_output, + b=gate_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_b=True, ) grad_scatter_output += group_gemm_same_nk( - a=grad_up_output, b=up_proj, - cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, + a=grad_up_output, + b=up_proj, + cumsum_M=cumsum_t, + max_M=max_M, + transpose_b=True, ) # FC1 wgrad @@ -579,16 +754,26 @@ def backward(ctx, grad_output): if gate_proj.requires_grad: grad_gate_proj = torch.empty_like(gate_proj) group_gemm_same_mn( - a=scatter_output, b=grad_gate_output, c=grad_gate_proj, - cumsum_K=cumsum_t, max_K=max_M, transpose_a=True, transpose_b=False, + a=scatter_output, + b=grad_gate_output, + c=grad_gate_proj, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) del grad_gate_output grad_up_proj = None if up_proj.requires_grad: grad_up_proj = torch.empty_like(up_proj) group_gemm_same_mn( - a=scatter_output, b=grad_up_output, c=grad_up_proj, - cumsum_K=cumsum_t, max_K=max_M, transpose_a=True, transpose_b=False, + a=scatter_output, + b=grad_up_output, + c=grad_up_proj, + cumsum_K=cumsum_t, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) del grad_up_output, scatter_output diff --git a/src/xorl/ops/moe/triton_lora.py b/src/xorl/ops/moe/triton_lora.py index 2d82fea7..ad6f22b7 100644 --- a/src/xorl/ops/moe/triton_lora.py +++ b/src/xorl/ops/moe/triton_lora.py @@ -7,8 +7,10 @@ import torch from xorl.ops.group_gemm.kernel import group_gemm_same_mn, group_gemm_same_nk + from .lora import make_ep_lora_compute, make_local_lora_compute + # EP LoRA compute (Expert Parallelism) TritonEPGroupGemmWithLoRA = make_ep_lora_compute(group_gemm_same_nk, group_gemm_same_mn) TritonEPGroupGemmWithLoRA.__name__ = TritonEPGroupGemmWithLoRA.__qualname__ = "TritonEPGroupGemmWithLoRA" @@ -39,7 +41,18 @@ def triton_moe_lora_forward( EP is handled centrally by ``MoEExpertsLoRA._ep_forward()``. """ return TritonMoeExpertsLoRAFunction.apply( - num_experts, routing_weights.to(hidden_states.dtype), selected_experts, hidden_states, - gate_proj, up_proj, down_proj, - gate_proj_lora_A, gate_proj_lora_B, up_proj_lora_A, up_proj_lora_B, down_proj_lora_A, down_proj_lora_B, scaling, + num_experts, + routing_weights.to(hidden_states.dtype), + selected_experts, + hidden_states, + gate_proj, + up_proj, + down_proj, + gate_proj_lora_A, + gate_proj_lora_B, + up_proj_lora_A, + up_proj_lora_B, + down_proj_lora_A, + down_proj_lora_B, + scaling, ) diff --git a/src/xorl/ops/quack/__init__.py b/src/xorl/ops/quack/__init__.py index 06c39487..0c4f0c24 100644 --- a/src/xorl/ops/quack/__init__.py +++ b/src/xorl/ops/quack/__init__.py @@ -2,9 +2,9 @@ import os +from .cross_entropy import cross_entropy from .rmsnorm import rmsnorm from .softmax import softmax -from .cross_entropy import cross_entropy if os.environ.get("CUTE_DSL_PTXAS_PATH", None) is not None: diff --git a/src/xorl/ops/quack/activation.py b/src/xorl/ops/quack/activation.py index e8cbeb29..6ebe6917 100644 --- a/src/xorl/ops/quack/activation.py +++ b/src/xorl/ops/quack/activation.py @@ -1,13 +1,13 @@ # Copyright (c) 2025, Tri Dao. import math -from typing import Tuple from functools import partial +from typing import Tuple import cutlass.cute as cute -from cutlass import Float32, Boolean, const_expr -from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass import Boolean, Float32, const_expr from cutlass._mlir.dialects import llvm, nvvm +from cutlass.cutlass_dsl import T, dsl_user_op F32_or_F32x2 = Float32 | Tuple[Float32, Float32] @@ -62,9 +62,7 @@ def relu(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2: @dsl_user_op @cute.jit -def drelu( - x: F32_or_F32x2, dout: F32_or_F32x2, *, loc=None, ip=None -) -> Tuple[F32_or_F32x2, F32_or_F32x2]: +def drelu(x: F32_or_F32x2, dout: F32_or_F32x2, *, loc=None, ip=None) -> Tuple[F32_or_F32x2, F32_or_F32x2]: if const_expr(not isinstance(x, tuple)): x_pos = Boolean(x > 0) return dout if x_pos else Float32(0.0), cute.arch.fmax(x, Float32(0.0)) @@ -86,9 +84,7 @@ def relu_sq(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2: @dsl_user_op @cute.jit -def drelu_sq( - x: F32_or_F32x2, dout: F32_or_F32x2, *, loc=None, ip=None -) -> Tuple[F32_or_F32x2, F32_or_F32x2]: +def drelu_sq(x: F32_or_F32x2, dout: F32_or_F32x2, *, loc=None, ip=None) -> Tuple[F32_or_F32x2, F32_or_F32x2]: """ ReLU squared backward pass: computes gradient w.r.t. x and recomputes forward Given: relu_sq_out = max(x, 0) * x, and dout = grad w.r.t. relu_sq_out @@ -136,9 +132,7 @@ def gelu_tanh_approx(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2: @dsl_user_op -def dgelu_tanh_approx( - x: F32_or_F32x2, dout: F32_or_F32x2, *, loc=None, ip=None -) -> Tuple[F32_or_F32x2, F32_or_F32x2]: +def dgelu_tanh_approx(x: F32_or_F32x2, dout: F32_or_F32x2, *, loc=None, ip=None) -> Tuple[F32_or_F32x2, F32_or_F32x2]: """ GELU tanh approximation backward pass: computes gradient w.r.t. x and recomputes forward Given: gelu_out = 0.5 * x * (1 + tanh(x * (c1 + c2 * x^2))), and dout = grad w.r.t. gelu_out @@ -203,11 +197,7 @@ def dgelu_tanh_approx( def softplus(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2: if const_expr(not isinstance(x, tuple)): use_linear = Boolean(x > 20.0) - return ( - cute.math.log(Float32(cute.math.exp(x, fastmath=True)) + 1.0, fastmath=True) - if not use_linear - else x - ) + return cute.math.log(Float32(cute.math.exp(x, fastmath=True)) + 1.0, fastmath=True) if not use_linear else x else: log2_e = math.log2(math.e) x_log2e = cute.arch.mul_packed_f32x2(x, (log2_e, log2_e)) @@ -314,12 +304,8 @@ def dswiglu( silu_x = cute.arch.fma_packed_f32x2(x, tanh_x, x) silu_x_dout = cute.arch.mul_packed_f32x2(silu_x, dout) # d_silu(x) * dout = (sigmoid_x - silu_x * sigmoid_x) * dout + silu_x * dout - sigmoid_x_minus_silu_x_sigmoid_x = cute.arch.fma_packed_f32x2( - sigmoid_x, (-silu_x[0], -silu_x[1]), sigmoid_x - ) - d_silu_x_dout = cute.arch.fma_packed_f32x2( - sigmoid_x_minus_silu_x_sigmoid_x, dout, silu_x_dout - ) + sigmoid_x_minus_silu_x_sigmoid_x = cute.arch.fma_packed_f32x2(sigmoid_x, (-silu_x[0], -silu_x[1]), sigmoid_x) + d_silu_x_dout = cute.arch.fma_packed_f32x2(sigmoid_x_minus_silu_x_sigmoid_x, dout, silu_x_dout) dx = cute.arch.mul_packed_f32x2(d_silu_x_dout, y) dy = silu_x_dout swiglu_out = cute.arch.mul_packed_f32x2(silu_x, y) @@ -327,9 +313,7 @@ def dswiglu( @dsl_user_op -def swiglu_oai( - x: F32_or_F32x2, y: F32_or_F32x2, alpha: float = 1.702, *, loc=None, ip=None -) -> F32_or_F32x2: +def swiglu_oai(x: F32_or_F32x2, y: F32_or_F32x2, alpha: float = 1.702, *, loc=None, ip=None) -> F32_or_F32x2: """The swiglu variant used in gpt-oss, which has a scaling factor on x and bias of 1 to y. https://github.com/openai/gpt-oss/blob/7be9334950053a888e24887a57dac797a17d6e00/gpt_oss/torch/model.py#L249 x * sigmoid(alpha * x) * (y + 1) @@ -384,12 +368,8 @@ def dswiglu_oai( silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_alpha_x) silu_x_dout = cute.arch.mul_packed_f32x2(silu_x, dout) # d_silu_x_dout = (sigmoid_alpha_x + alpha * (silu_x - silu_x * sigmoid_alpha_x)) * dout - silu_x_minus_product = cute.arch.fma_packed_f32x2( - silu_x, (-sigmoid_alpha_x[0], -sigmoid_alpha_x[1]), silu_x - ) - sigmoid_plus_alpha_diff = cute.arch.fma_packed_f32x2( - (alpha, alpha), silu_x_minus_product, sigmoid_alpha_x - ) + silu_x_minus_product = cute.arch.fma_packed_f32x2(silu_x, (-sigmoid_alpha_x[0], -sigmoid_alpha_x[1]), silu_x) + sigmoid_plus_alpha_diff = cute.arch.fma_packed_f32x2((alpha, alpha), silu_x_minus_product, sigmoid_alpha_x) d_silu_x_dout = cute.arch.mul_packed_f32x2(sigmoid_plus_alpha_diff, dout) dx = cute.arch.fma_packed_f32x2(d_silu_x_dout, y, d_silu_x_dout) dy = silu_x_dout diff --git a/src/xorl/ops/quack/autotuner.py b/src/xorl/ops/quack/autotuner.py index de1f63a1..dc5b4c52 100644 --- a/src/xorl/ops/quack/autotuner.py +++ b/src/xorl/ops/quack/autotuner.py @@ -2,21 +2,20 @@ # Copyright (C) 2025, Tri Dao. from __future__ import annotations -import builtins -import os -import time -import inspect import base64 +import builtins import hashlib +import inspect import json -from pathlib import Path +import os +import time from functools import cached_property, partial -from typing import Dict, Tuple, List, Optional, Any +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple import torch -from torch import Tensor - import triton +from torch import Tensor from . import __version__ @@ -36,9 +35,7 @@ def default_cache_dir(): class FileCacheManager(triton.runtime.cache.FileCacheManager): def __init__(self, key): super().__init__(key) - self.cache_dir = ( - os.getenv(f"{PACKAGE_NAME.upper()}_CACHE_DIR", "").strip() or default_cache_dir() - ) + self.cache_dir = os.getenv(f"{PACKAGE_NAME.upper()}_CACHE_DIR", "").strip() or default_cache_dir() if self.cache_dir: self.cache_dir = os.path.join(self.cache_dir, self.key) self.lock_path = os.path.join(self.cache_dir, "lock") @@ -77,9 +74,7 @@ def __init__( self.keys = key self.cache: Dict[Tuple, AutotuneConfig] = {} self.arg_names = list(signature.parameters.keys()) - self.cache_results = ( - cache_results or os.getenv(f"{PACKAGE_NAME.upper()}_CACHE_AUTOTUNING", None) == "1" - ) + self.cache_results = cache_results or os.getenv(f"{PACKAGE_NAME.upper()}_CACHE_AUTOTUNING", None) == "1" self.restore_value = [] if restore_value is not None: @@ -111,9 +106,7 @@ def _post_hook(kwargs, exception): if prune_configs_by: self.perf_model = prune_configs_by.get("perf_model", self.perf_model) self.configs_top_k = prune_configs_by.get("top_k", self.configs_top_k) - self.early_config_prune = prune_configs_by.get( - "early_config_prune", self.early_config_prune - ) + self.early_config_prune = prune_configs_by.get("early_config_prune", self.early_config_prune) self.fn = fn self._do_bench = do_bench @@ -197,9 +190,7 @@ def check_disk_cache(self, tuning_key, configs, bench_fn): json.dumps( { "key": tuning_key, - "configs_timings": [ - (str(config), timings) for config, timings in self.configs_timings.items() - ], + "configs_timings": [(str(config), timings) for config, timings in self.configs_timings.items()], } ), file_name, @@ -228,10 +219,7 @@ def __call__(self, *args, **kwargs): @torch.compiler.disable # Don't want any tracing here def benchmark(): bench_start = time.time() - timings = { - config: self._bench(*args, config=config, **kwargs) - for config in pruned_configs - } + timings = {config: self._bench(*args, config=config, **kwargs) for config in pruned_configs} bench_end = time.time() if os.getenv(f"{PACKAGE_NAME.upper()}_PRINT_AUTOTUNING", None) == "1": for config, time_ in timings.items(): @@ -249,10 +237,7 @@ def benchmark(): else: config = self.configs[0] self.best_config = config - if ( - os.getenv(f"{PACKAGE_NAME.upper()}_PRINT_AUTOTUNING", None) == "1" - and not used_cached_result - ): + if os.getenv(f"{PACKAGE_NAME.upper()}_PRINT_AUTOTUNING", None) == "1" and not used_cached_result: print( f"{PACKAGE_NAME} autotuning for function {self.fn.__name__} finished after " f"{self.bench_time:.2f}s; best config selected: {self.best_config};" @@ -275,9 +260,7 @@ def prune_configs(self, kwargs: Dict) -> List[Any]: top_k = int(len(self.configs) * top_k) elif not isinstance(top_k, int): # Slice index must be an integer - raise TypeError( - "Error while pruning configs, top_k must be either 1) a float <= 1.0 or 2) an int" - ) + raise TypeError("Error while pruning configs, top_k must be either 1) a float <= 1.0 or 2) an int") if len(pruned_configs) > top_k: est_timing = { @@ -324,9 +307,7 @@ def __eq__(self, other): return self_tuple == other_tuple -def autotune( - configs, key=None, prune_configs_by=None, restore_value=None, do_bench=None, cache_results=True -): +def autotune(configs, key=None, prune_configs_by=None, restore_value=None, do_bench=None, cache_results=True): f""" Decorator for auto-tuning a function function. diff --git a/src/xorl/ops/quack/compile_utils.py b/src/xorl/ops/quack/compile_utils.py index 43755946..d2ed1998 100644 --- a/src/xorl/ops/quack/compile_utils.py +++ b/src/xorl/ops/quack/compile_utils.py @@ -10,10 +10,5 @@ def make_fake_tensor(dtype, shape, divisibility=1, leading_dim=-1) -> Optional[c leading_dim = len(shape) + leading_dim if dtype is None: return None - stride = tuple( - cute.sym_int64(divisibility=divisibility) if i != leading_dim else 1 - for i in range(len(shape)) - ) - return cute.runtime.make_fake_tensor( - dtype, shape, stride=stride, assumed_align=divisibility * dtype.width // 8 - ) + stride = tuple(cute.sym_int64(divisibility=divisibility) if i != leading_dim else 1 for i in range(len(shape))) + return cute.runtime.make_fake_tensor(dtype, shape, stride=stride, assumed_align=divisibility * dtype.width // 8) diff --git a/src/xorl/ops/quack/copy_utils.py b/src/xorl/ops/quack/copy_utils.py index 7ad98955..0a08f580 100644 --- a/src/xorl/ops/quack/copy_utils.py +++ b/src/xorl/ops/quack/copy_utils.py @@ -1,20 +1,19 @@ # Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao. import re -from typing import Optional, Type, Tuple, Callable, Sequence from functools import partial +from typing import Callable, Optional, Sequence, Tuple, Type import cutlass import cutlass.cute as cute - -from cutlass import Int32, Int16, Boolean, const_expr -from cutlass.cute.nvgpu import cpasync, warp, warpgroup -from cutlass.cute.nvgpu.tcgen05.mma import CtaGroup # noqa -from cutlass.cutlass_dsl import dsl_user_op import cutlass.pipeline -from cutlass._mlir.dialects import llvm +from cutlass import Boolean, Int16, Int32, const_expr from cutlass._mlir import ir from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, warp, warpgroup +from cutlass.cute.nvgpu.tcgen05.mma import CtaGroup # noqa +from cutlass.cutlass_dsl import dsl_user_op Sm100MmaPeerBitMask = 0xFEFFFFFF @@ -197,25 +196,19 @@ def as_position_independent_swizzle_tensor(tensor: cute.Tensor) -> cute.Tensor: inner = cute.make_swizzle(*parse_swizzle_from_pointer(tensor.iterator)) # Need to recast the swizzle from byte (e.g. <3, 4, 3> to element units (e.g. <3, 3, 3> for # for 16 bits and <3, 2, 3> for 32 bits) - new_layout = cute.recast_layout( - width, 8, cute.make_composed_layout(inner, 0, cute.recast_layout(8, width, outer)) - ) + new_layout = cute.recast_layout(width, 8, cute.make_composed_layout(inner, 0, cute.recast_layout(8, width, outer))) # recast_ptr to remove the pointer swizzle return cute.make_tensor(cute.recast_ptr(tensor.iterator, dtype=tensor.element_type), new_layout) -def partition_D_position_independent( - thr_copy: cute.core.ThrCopy, tensor: cute.Tensor -) -> cute.Tensor: +def partition_D_position_independent(thr_copy: cute.core.ThrCopy, tensor: cute.Tensor) -> cute.Tensor: return cute.make_tensor( swizzle_ptr(thr_copy.partition_D(tensor).iterator), thr_copy.partition_D(as_position_independent_swizzle_tensor(tensor)).layout, ) -def partition_S_position_independent( - thr_copy: cute.core.ThrCopy, tensor: cute.Tensor -) -> cute.Tensor: +def partition_S_position_independent(thr_copy: cute.core.ThrCopy, tensor: cute.Tensor) -> cute.Tensor: return cute.make_tensor( swizzle_ptr(thr_copy.partition_S(tensor).iterator), thr_copy.partition_S(as_position_independent_swizzle_tensor(tensor)).layout, @@ -338,9 +331,7 @@ def copy_fn(src_idx: Optional[Int32] = None, **new_kwargs): return copy_fn, thr_copy, tSR_sC -def epilog_smem_copy_atom( - tiled_mma: cute.TiledMma, epi_tile: cute.Shape, transpose: bool = False -) -> cute.TiledCopy: +def epilog_smem_copy_atom(tiled_mma: cute.TiledMma, epi_tile: cute.Shape, transpose: bool = False) -> cute.TiledCopy: copy_atom_C = cute.make_copy_atom( warp.StMatrix8x8x16bOp(transpose, num_matrices=4 if epi_tile[1] % 16 == 0 else 2), cutlass.Float16, # this is just to get the right source layout @@ -419,9 +410,7 @@ def get_smem_load_A( tRS_shape = tiled_mma.partition_shape_A(sA.shape[:2]) def copy_fn(src_idx: Int32, **new_kwargs): - return load_s2r_retile( - tiled_copy, tSR_sA[None, None, None, src_idx], dst_shape=tRS_shape, **new_kwargs - ) + return load_s2r_retile(tiled_copy, tSR_sA[None, None, None, src_idx], dst_shape=tRS_shape, **new_kwargs) def copy_fn_w_dst_tensor(src_idx: Int32, dst: cute.Tensor, **new_kwargs): return load_s2r_retile(tiled_copy, tSR_sA[None, None, None, src_idx], dst, **new_kwargs) @@ -469,9 +458,7 @@ def get_tma_desc_addr(tma_atom: cute.CopyAtom, *, loc=None, ip=None) -> cute.Poi >>> desc_ptr = get_tma_descriptor_address(tma_atom) """ exec_atom = _cute_nvgpu_ir.atom_make_exec_tma(tma_atom._trait.value, loc=loc, ip=ip) - tma_desc_ptr_type = ir.Type.parse( - "!cute.ptr>" - ) + tma_desc_ptr_type = ir.Type.parse("!cute.ptr>") return _cute_nvgpu_ir.get_tma_desc_addr(tma_desc_ptr_type, exec_atom, loc=loc, ip=ip) @@ -630,8 +617,7 @@ def tma_get_copy_fn( **kwargs, ) -> Callable: src_is_smem = const_expr( - isinstance(src_tensor.iterator, cute.Pointer) - and src_tensor.memspace == cute.AddressSpace.smem + isinstance(src_tensor.iterator, cute.Pointer) and src_tensor.memspace == cute.AddressSpace.smem ) smem_tensor, gmem_tensor = (src_tensor, dst_tensor) if src_is_smem else (dst_tensor, src_tensor) group_rank_smem = const_expr(cute.rank(smem_tensor) - (1 if not single_stage else 0)) @@ -726,9 +712,9 @@ def copy_fn(src_idx, dst_idx, pred: bool = False): # ((elems_per_load), thread_per_row) # But we actually want shape ((elems_per_load, 1), thread_per_row) to match tAsA # So we append 1s to the last dimension and then do tiled_divide, then slice. - mA_row = cute.tiled_divide( - cute.append_ones(mA_cur[m_idx[m], None], up_to_rank=2), (elems_per_load, 1) - )[None, None, 0] + mA_row = cute.tiled_divide(cute.append_ones(mA_cur[m_idx[m], None], up_to_rank=2), (elems_per_load, 1))[ + None, None, 0 + ] if const_expr(is_even_m_smem) or tApA_m[m]: # There's only 1 load per row assert cute.size(tAcA.shape, mode=[2]) == 1 @@ -782,9 +768,9 @@ def gather_k_get_copy_fn( # for tile_M=128, flat_divide gives (8, 16, K), # then logical_divide gives ((8, 1), (8, 2), K). tidx = thr_copy_A.thr_idx - tAmA = cute.logical_divide( - cute.flat_divide(mA, (elems_per_load,)), (elems_per_load, threads_per_col) - )[None, (tidx % threads_per_col, None), None] # ((8, 1), 2, K) + tAmA = cute.logical_divide(cute.flat_divide(mA, (elems_per_load,)), (elems_per_load, threads_per_col))[ + None, (tidx % threads_per_col, None), None + ] # ((8, 1), 2, K) def prefetch_from_gmem_fn(src_idx, pred: bool = False) -> Tuple[cute.Tensor, cute.Tensor]: # Prefetch mAIdx early, even before smem is free @@ -827,9 +813,7 @@ def prefetch_from_smem_fn( a_prefetch_pipeline.consumer_release(a_prefetch_consumer_state) return k_idx, tApA_k - def copy_fn( - src_idx, dst_idx, k_idx_tApA_k: Tuple[cute.Tensor, cute.Tensor], pred: bool = False - ): + def copy_fn(src_idx, dst_idx, k_idx_tApA_k: Tuple[cute.Tensor, cute.Tensor], pred: bool = False): k_idx, tApA_k = k_idx_tApA_k tApA_k_pred = None if const_expr(pred): @@ -845,9 +829,7 @@ def copy_fn( pred=None if const_expr(tApA_k_pred is None) else tApA_k_pred[None, k], ) - return copy_fn, prefetch_from_gmem_fn if const_expr( - gAIdx is not None - ) else prefetch_from_smem_fn + return copy_fn, prefetch_from_gmem_fn if const_expr(gAIdx is not None) else prefetch_from_smem_fn @cute.jit diff --git a/src/xorl/ops/quack/cross_entropy.py b/src/xorl/ops/quack/cross_entropy.py index 788e7dca..bbc6de34 100644 --- a/src/xorl/ops/quack/cross_entropy.py +++ b/src/xorl/ops/quack/cross_entropy.py @@ -2,24 +2,20 @@ import math from functools import partial -from typing import Optional, Type, Literal - -import torch -from torch import Tensor +from typing import Literal, Optional, Type import cuda.bindings.driver as cuda - import cutlass import cutlass.cute as cute -from cutlass import Int32, Int64, Float32, Boolean, const_expr +import torch +from cutlass import Boolean, Float32, Int32, Int64, const_expr +from torch import Tensor -from . import utils -from . import copy_utils -from . import layout_utils +from . import copy_utils, layout_utils, utils from .compile_utils import make_fake_tensor as fake_tensor -from .reduce import row_reduce, online_softmax_reduce -from .reduction_base import ReductionBase from .cute_dsl_utils import torch2cute_dtype_map +from .reduce import online_softmax_reduce, row_reduce +from .reduction_base import ReductionBase class CrossEntropy(ReductionBase): @@ -120,9 +116,7 @@ def kernel( gX, cX = [cute.local_tile(mT, tiler_mn, (bidx, cluster_y)) for mT in (mX, idX)] smem = cutlass.utils.SmemAllocator() - sX = smem.allocate_tensor( - mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16 - ) + sX = smem.allocate_tensor(mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16) reduction_buffer, mbar_ptr = self._allocate_reduction_buffer_and_mbar(smem, tv_layout) thr_copy = tiled_copy.get_slice(tidx) @@ -133,9 +127,7 @@ def kernel( tXrX = cute.make_fragment_like(tXgX) is_even_N = const_expr(shape[1] == tiler_mn[1] * self.cluster_n) - tXpX = ( - None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) - ) + tXpX = None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) copy = partial(copy_utils.copy, pred=tXpX) num_warps = cute.size(tiled_copy) // cute.arch.WARP_SIZE @@ -201,11 +193,7 @@ def kernel( ) # Write loss and lse to gmem - if ( - tXcX[0][1] == 0 - and row < shape[0] - and (self.cluster_n == 1 or cute.arch.block_idx_in_cluster() == 0) - ): + if tXcX[0][1] == 0 and row < shape[0] and (self.cluster_n == 1 or cute.arch.block_idx_in_cluster() == 0): lse = max_x + cute.math.log(denom, fastmath=True) # Set loss to 0 if this index should be ignored, otherwise compute normally loss_val = (lse - target_logit) if not should_ignore else Float32.zero @@ -219,9 +207,7 @@ def kernel( # If ignored, gradient should be zero denom_inv = ( # 1.0 / denom - cute.arch.rcp_approx(denom) - if not (denom == 0.0 or denom != denom or should_ignore) - else Float32.zero + cute.arch.rcp_approx(denom) if not (denom == 0.0 or denom != denom or should_ignore) else Float32.zero ) probs = exp_x * denom_inv gdX = cute.local_tile(mdX, tiler_mn, (bidx, cluster_y)) @@ -278,9 +264,7 @@ def cross_entropy_fwd_out( N = x.size(1) dtype = torch2cute_dtype_map[x.dtype] target_dtype = torch2cute_dtype_map[target.dtype] - target_logit_dtype = ( - torch2cute_dtype_map[target_logit.dtype] if target_logit is not None else None - ) + target_logit_dtype = torch2cute_dtype_map[target_logit.dtype] if target_logit is not None else None compile_key = ( dtype, target_dtype, @@ -297,9 +281,7 @@ def cross_entropy_fwd_out( target_cute = fake_tensor(target_dtype, (batch_sym,)) if target_logit is not None: if target_logit.ndim == 2: - target_logit_cute = fake_tensor( - target_logit_dtype, (batch_sym, cute.sym_int()), div - ) + target_logit_cute = fake_tensor(target_logit_dtype, (batch_sym, cute.sym_int()), div) else: target_logit_cute = fake_tensor(target_logit_dtype, (batch_sym,)) else: @@ -320,9 +302,7 @@ def cross_entropy_fwd_out( cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), options="--enable-tvm-ffi", ) - cross_entropy_fwd_out.compile_cache[compile_key]( - x, target, target_logit, loss, lse, dx, Int32(ignore_index) - ) + cross_entropy_fwd_out.compile_cache[compile_key](x, target, target_logit, loss, lse, dx, Int32(ignore_index)) cross_entropy_fwd_out.compile_cache = {} @@ -374,9 +354,7 @@ def _get_tiled_copy(self, vecsize: int): cols_per_block = num_threads // threads_per_row num_blocks_N = cute.ceil_div(N // vecsize, threads_per_row) tiler_mn = (cols_per_block, vecsize * num_blocks_N * threads_per_row) - tiled_copy = copy_utils.tiled_copy_2d( - self.dtype, threads_per_row, num_threads, num_copy_elems=vecsize - ) + tiled_copy = copy_utils.tiled_copy_2d(self.dtype, threads_per_row, num_threads, num_copy_elems=vecsize) return tiled_copy, tiler_mn, threads_per_row @cute.jit @@ -397,9 +375,7 @@ def __call__( tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=vecsize) num_threads = tiled_copy.size # (M,) -> (M, N) with stride 0 in the N dimension - mDLoss, mTarget, mLSE = [ - layout_utils.expand(X, dim=1, size=self.N) for X in (mDLoss, mTarget, mLSE) - ] + mDLoss, mTarget, mLSE = [layout_utils.expand(X, dim=1, size=self.N) for X in (mDLoss, mTarget, mLSE)] self.kernel( mX, mTarget, @@ -439,9 +415,7 @@ def kernel( bidx, bidy, _ = cute.arch.block_idx() smem = cutlass.utils.SmemAllocator() - sX = smem.allocate_tensor( - mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16 - ) + sX = smem.allocate_tensor(mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16) idX = cute.make_identity_tensor(shape) gX, gdX, cX = [cute.local_tile(mT, tiler_mn, (bidx, bidy)) for mT in (mX, mdX, idX)] @@ -456,9 +430,7 @@ def kernel( tXrX, tXrdX = [cute.make_fragment_like(thr) for thr in (tXgX, tXgdX)] is_even_N = const_expr(shape[1] % tiler_mn[1] == 0) - tXpX = ( - None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) - ) + tXpX = None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) copy = partial(copy_utils.copy, pred=tXpX) row = tXcX[0][0] @@ -520,9 +492,7 @@ def _cross_entropy_backward( assert x.shape[0] == target.shape[0], "Batch dimensions must match" assert x.shape[0] == dloss.shape[0], "Batch dimensions must match" assert x.shape[0] == lse.shape[0], "Batch dimensions must match" - assert x.is_cuda and target.is_cuda and dloss.is_cuda and lse.is_cuda, ( - "Tensors must be on CUDA device" - ) + assert x.is_cuda and target.is_cuda and dloss.is_cuda and lse.is_cuda, "Tensors must be on CUDA device" assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported input dtype" assert target.dtype in [torch.int32, torch.int64], "Target must be int32 or int64" @@ -548,9 +518,7 @@ def _cross_entropy_backward( cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), options="--enable-tvm-ffi", ) - _cross_entropy_backward.compile_cache[compile_key]( - x, target, dloss, dx, lse, Int32(ignore_index) - ) + _cross_entropy_backward.compile_cache[compile_key](x, target, dloss, dx, lse, Int32(ignore_index)) _cross_entropy_backward.compile_cache = {} @@ -578,14 +546,10 @@ def cross_entropy_bwd( ) -> None: if inplace_backward and not torch.compiler.is_compiling(): dx = x - _cross_entropy_backward( - x=x, target=target, dloss=dloss, lse=lse, dx=x, ignore_index=ignore_index - ) + _cross_entropy_backward(x=x, target=target, dloss=dloss, lse=lse, dx=x, ignore_index=ignore_index) else: dx = torch.empty_like(x) - cross_entropy_bwd_out( - x=x, target=target, dloss=dloss, lse=lse, dx=dx, ignore_index=ignore_index - ) + cross_entropy_bwd_out(x=x, target=target, dloss=dloss, lse=lse, dx=dx, ignore_index=ignore_index) return dx @@ -608,9 +572,7 @@ def forward(ctx, x, target, lse_partial=None, ignore_index=-100, inplace_backwar @staticmethod def backward(ctx, dloss): x, target, lse = ctx.saved_tensors - dx = cross_entropy_bwd( - x, target, dloss, lse, ctx.ignore_index, inplace_backward=ctx.inplace_backward - ) + dx = cross_entropy_bwd(x, target, dloss, lse, ctx.ignore_index, inplace_backward=ctx.inplace_backward) return dx, None, None, None, None @@ -649,6 +611,4 @@ def cross_entropy( elif reduction == "none": return loss else: - raise ValueError( - f"Invalid reduction mode: {reduction}. Expected one of 'none', 'mean', or 'sum'" - ) + raise ValueError(f"Invalid reduction mode: {reduction}. Expected one of 'none', 'mean', or 'sum'") diff --git a/src/xorl/ops/quack/cute_dsl_ptxas.py b/src/xorl/ops/quack/cute_dsl_ptxas.py index 4e00f3f0..6091aa49 100644 --- a/src/xorl/ops/quack/cute_dsl_ptxas.py +++ b/src/xorl/ops/quack/cute_dsl_ptxas.py @@ -5,11 +5,11 @@ CUTE_DSL_PTXAS_VERBOSE - Set to 1 for verbose output """ +import ctypes import os -import sys import re -import ctypes import subprocess +import sys from pathlib import Path import cutlass @@ -140,9 +140,7 @@ def patch(): # Track if user originally wanted PTX kept _user_wanted_ptx = os.environ.get("CUTE_DSL_KEEP_PTX", "0") == "1" # os.environ['CUTE_DSL_KEEP_PTX'] = '1' - assert os.environ.get("CUTE_DSL_KEEP_PTX", "0") == "1", ( - "Require CUTE_DSL_KEEP_PTX=1 to use system's ptxas" - ) + assert os.environ.get("CUTE_DSL_KEEP_PTX", "0") == "1", "Require CUTE_DSL_KEEP_PTX=1 to use system's ptxas" cls = cutlass.cutlass_dsl.cuda_jit_executor.CudaDialectJitCompiledFunction _original_load_cuda_library = cls._load_cuda_library diff --git a/src/xorl/ops/quack/cute_dsl_utils.py b/src/xorl/ops/quack/cute_dsl_utils.py index 9c92cf39..2ceb256b 100644 --- a/src/xorl/ops/quack/cute_dsl_utils.py +++ b/src/xorl/ops/quack/cute_dsl_utils.py @@ -1,11 +1,12 @@ # Copyright (c) 2025, Tri Dao. -from typing import Tuple -from functools import lru_cache from dataclasses import dataclass, fields +from functools import lru_cache +from typing import Tuple import torch + try: from triton.tools.disasm import extract except ImportError: @@ -13,7 +14,7 @@ import cutlass import cutlass.cute as cute -from cutlass import Int32, Int64, Float16, BFloat16, Float32 +from cutlass import BFloat16, Float16, Float32, Int32, Int64 from cutlass.base_dsl.typing import JitArgument from cutlass.cutlass_dsl import NumericMeta @@ -59,9 +60,7 @@ def __extract_mlir_values__(self): def __new_from_mlir_values__(self, values): all_fields = {field.name: getattr(self, field.name) for field in fields(self)} constexpr_fields = {n: f for n, f in all_fields.items() if isinstance(f, StaticTypes)} - non_constexpr_fields = { - n: f for n, f in all_fields.items() if not isinstance(f, StaticTypes) - } + non_constexpr_fields = {n: f for n, f in all_fields.items() if not isinstance(f, StaticTypes)} for (name, field), n_items in zip(non_constexpr_fields.items(), self._values_pos): non_constexpr_fields[name] = cutlass.new_from_mlir_values(field, values[:n_items]) values = values[n_items:] @@ -95,9 +94,7 @@ def __get_mlir_types__(self): def __new_from_mlir_values__(self, values): all_fields = {field.name: getattr(self, field.name) for field in fields(self)} constexpr_fields = {n: f for n, f in all_fields.items() if isinstance(f, StaticTypes)} - non_constexpr_fields = { - n: f for n, f in all_fields.items() if not isinstance(f, StaticTypes) - } + non_constexpr_fields = {n: f for n, f in all_fields.items() if not isinstance(f, StaticTypes)} for (name, field), n_items in zip(non_constexpr_fields.items(), self._values_pos): non_constexpr_fields[name] = cutlass.new_from_mlir_values(field, values[:n_items]) values = values[n_items:] diff --git a/src/xorl/ops/quack/gemm.py b/src/xorl/ops/quack/gemm.py index d3d3f1af..c5b20077 100644 --- a/src/xorl/ops/quack/gemm.py +++ b/src/xorl/ops/quack/gemm.py @@ -1,16 +1,15 @@ -from typing import Optional from functools import partial - -from torch import Tensor +from typing import Optional import cutlass.cute as cute import cutlass.torch as cutlass_torch from cutlass import Float32 from cutlass.cute.runtime import from_dlpack, make_ptr +from torch import Tensor from .cute_dsl_utils import get_device_capacity, get_max_active_clusters -from .gemm_wrapper_utils import GemmWrapperBase from .gemm_default_epi import GemmDefaultSm90, GemmDefaultSm100 +from .gemm_wrapper_utils import GemmWrapperBase def gemm( @@ -59,9 +58,7 @@ def gemm( L, M, K, N, tensor_infos = GemmWrapperBase.validate_and_prepare_tensors( A, B, D, C, cu_seqlens_m=cu_seqlens_m, cu_seqlens_k=cu_seqlens_k, A_idx=A_idx ) - GemmWrapperBase.permute_tensors( - tensor_infos, varlen_m=cu_seqlens_m is not None, varlen_k=cu_seqlens_k is not None - ) + GemmWrapperBase.permute_tensors(tensor_infos, varlen_m=cu_seqlens_m is not None, varlen_k=cu_seqlens_k is not None) GemmWrapperBase.extract_dtypes(tensor_infos) major_configs = { "A": ("m", "k", "l"), @@ -101,9 +98,7 @@ def scalar_arg(scalar: float | Tensor): epi_args = GemmCls.EpilogueArguments( scalar_arg(alpha), scalar_arg(beta), - mRowVecBroadcast=from_dlpack(rowvec_bias.detach(), assumed_align=4).mark_layout_dynamic( - leading_dim=1 - ) + mRowVecBroadcast=from_dlpack(rowvec_bias.detach(), assumed_align=4).mark_layout_dynamic(leading_dim=1) if rowvec_bias is not None else None, mColVecBroadcast=from_dlpack(colvec_bias.detach(), assumed_align=4).mark_layout_dynamic( diff --git a/src/xorl/ops/quack/gemm_act.py b/src/xorl/ops/quack/gemm_act.py index 2d18ca6b..7cf208d7 100644 --- a/src/xorl/ops/quack/gemm_act.py +++ b/src/xorl/ops/quack/gemm_act.py @@ -1,29 +1,25 @@ # Copyright (c) 2025, Wentao Guo, Tri Dao. -from typing import Tuple, Optional, Callable -from functools import partial from dataclasses import dataclass - -from torch import Tensor +from functools import partial +from typing import Callable, Optional, Tuple import cutlass import cutlass.cute as cute -import cutlass.utils.hopper_helpers as sm90_utils_og -import cutlass.utils.blackwell_helpers as sm100_utils -from cutlass import Int32, Float32, Boolean, const_expr import cutlass.torch as cutlass_torch +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.hopper_helpers as sm90_utils_og +from cutlass import Boolean, Float32, Int32, const_expr from cutlass.cute.runtime import from_dlpack +from torch import Tensor -from .cute_dsl_utils import ArgumentsBase, ParamsBase -from .varlen_utils import VarlenManager +from . import activation, copy_utils, sm90_utils +from .cute_dsl_utils import ArgumentsBase, ParamsBase, get_device_capacity, get_max_active_clusters +from .gemm_default_epi import GemmDefaultEpiMixin from .gemm_sm90 import GemmSm90 from .gemm_sm100 import GemmSm100 -from .gemm_default_epi import GemmDefaultEpiMixin -from .cute_dsl_utils import get_device_capacity, get_max_active_clusters from .gemm_wrapper_utils import GemmTensorInfo, GemmWrapperBase from .layout_utils import permute_gated_Cregs_b16 -from . import sm90_utils -from . import copy_utils -from . import activation +from .varlen_utils import VarlenManager class GemmActMixin(GemmDefaultEpiMixin): @@ -50,9 +46,7 @@ class EpilogueParams(ParamsBase): mRowVecBroadcast: Optional[cute.Tensor] = None mColVecBroadcast: Optional[cute.Tensor] = None - def epi_to_underlying_arguments( - self, args: EpilogueArguments, *, loc=None, ip=None - ) -> EpilogueParams: + def epi_to_underlying_arguments(self, args: EpilogueArguments, *, loc=None, ip=None) -> EpilogueParams: self.postact_dtype = args.mPostAct.element_type self.postact_layout = cutlass.utils.LayoutEnum.from_tensor(args.mPostAct) @@ -70,13 +64,10 @@ def epi_to_underlying_arguments( ) # Assume all strides are divisible by 32 bits except the last stride new_stride = lambda t: tuple( - cute.assume(s, divby=32 // t.element_type.width) if not cute.is_static(s) else s - for s in t.stride + cute.assume(s, divby=32 // t.element_type.width) if not cute.is_static(s) else s for s in t.stride ) mRowVecBroadcast, mColVecBroadcast = [ - cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) - if t is not None - else None + cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) if t is not None else None for t in (args.mRowVecBroadcast, args.mColVecBroadcast) ] return self.EpilogueParams( @@ -91,9 +82,7 @@ def epi_to_underlying_arguments( mColVecBroadcast=mColVecBroadcast, ) - def epi_get_tma_atoms( - self, params: EpilogueParams, *, loc=None, ip=None - ) -> list[cute.CopyAtom]: + def epi_get_tma_atoms(self, params: EpilogueParams, *, loc=None, ip=None) -> list[cute.CopyAtom]: return [params.tma_atom_postact] def epi_get_tensormap_update_shapes_orders( @@ -115,29 +104,21 @@ def epi_smem_bytes_per_stage( ) -> int: postact_dtype = args.mPostAct.element_type postact_bytes_per_stage = cute.size(cute.shape(epi_tile)) * (postact_dtype.width // 8) - rowvec_colvec_bytes = GemmDefaultEpiMixin.epi_smem_bytes_per_stage( - args, cta_tile_shape_mnk, epi_tile - ) + rowvec_colvec_bytes = GemmDefaultEpiMixin.epi_smem_bytes_per_stage(args, cta_tile_shape_mnk, epi_tile) return postact_bytes_per_stage + rowvec_colvec_bytes def epi_get_smem_struct(self, params: EpilogueParams): row_vec_smem_size = 0 if params.mRowVecBroadcast is None else self.cta_tile_shape_mnk[1] col_vec_smem_size = 0 if params.mColVecBroadcast is None else self.cta_tile_shape_mnk[0] - row_vec_dtype = ( - params.mRowVecBroadcast.element_type if params.mRowVecBroadcast is not None else Float32 - ) - col_vec_dtype = ( - params.mColVecBroadcast.element_type if params.mColVecBroadcast is not None else Float32 - ) + row_vec_dtype = params.mRowVecBroadcast.element_type if params.mRowVecBroadcast is not None else Float32 + col_vec_dtype = params.mColVecBroadcast.element_type if params.mColVecBroadcast is not None else Float32 @cute.struct class EpiSharedStorage: sRowVec: cute.struct.Align[cute.struct.MemRange[row_vec_dtype, row_vec_smem_size], 16] sColVec: cute.struct.Align[cute.struct.MemRange[col_vec_dtype, col_vec_smem_size], 16] sPostAct: cute.struct.Align[ - cute.struct.MemRange[ - self.postact_dtype, cute.cosize(params.epi_postact_smem_layout_staged) - ], + cute.struct.MemRange[self.postact_dtype, cute.cosize(params.epi_postact_smem_layout_staged)], self.buffer_align_bytes, ] @@ -191,9 +172,7 @@ def epilogue( if self.arch == 100 else sm90_utils_og.sm90_get_smem_store_op ) - copy_atom_postact_r2s = get_smem_store_op( - self.postact_layout, self.postact_dtype, self.acc_dtype - ) + copy_atom_postact_r2s = get_smem_store_op(self.postact_layout, self.postact_dtype, self.acc_dtype) # tiled_copy_C_atom = self.epilog_smem_copy_atom(tiled_mma) # tiled_copy_postact_r2s = cute.make_tiled_copy_S(copy_atom_postact_r2s, tiled_copy_C_atom) tiled_copy_postact_r2s = cute.make_tiled_copy_S(copy_atom_postact_r2s, tiled_copy_r2s) @@ -211,9 +190,7 @@ def epilogue( ) # We iterate over epi tiles in the N dimension first before the M dimension - epi_tile_shape = cute.zipped_divide( - cute.make_layout(self.cta_tile_shape_mnk[:2]), epi_tile - ).shape[1] + epi_tile_shape = cute.zipped_divide(cute.make_layout(self.cta_tile_shape_mnk[:2]), epi_tile).shape[1] epi_tile_layout = cute.make_layout(epi_tile_shape, stride=(epi_tile_shape[1], 1)) epi_tile_num = cute.size(epi_tile_shape) num_prev_subtiles = tile_scheduler.num_tiles_executed * epi_tile_num @@ -315,9 +292,7 @@ def epi_visit_subtile( tRS_rPostAct[i] = params.act_fn(tRS_rD[i]) else: for i in cutlass.range(cute.size(tRS_rPostAct) // 2, unroll_full=True): - tRS_rPostAct[2 * i], tRS_rPostAct[2 * i + 1] = params.act_fn( - (tRS_rD[2 * i], tRS_rD[2 * i + 1]) - ) + tRS_rPostAct[2 * i], tRS_rPostAct[2 * i + 1] = params.act_fn((tRS_rD[2 * i], tRS_rD[2 * i + 1])) else: tRS_rPostAct = tRS_rD # Type conversion @@ -411,9 +386,7 @@ def gemm_act( epi_args = GemmCls.EpilogueArguments( tensor_infos["PostAct"].cute_tensor, act_fn, - mRowVecBroadcast=from_dlpack(rowvec_bias.detach(), assumed_align=4).mark_layout_dynamic( - leading_dim=1 - ) + mRowVecBroadcast=from_dlpack(rowvec_bias.detach(), assumed_align=4).mark_layout_dynamic(leading_dim=1) if rowvec_bias is not None else None, mColVecBroadcast=from_dlpack(colvec_bias.detach(), assumed_align=4).mark_layout_dynamic( @@ -502,9 +475,7 @@ def epi_to_underlying_arguments( assert self.d_layout is None or self.d_layout.is_n_major_c() assert self.postact_layout.is_n_major_c() if self.arch == 90: - assert self.cta_tile_shape_mnk[1] % 32 == 0, ( - "GemmGatedSm90 requires tileN to be divisible by 32" - ) + assert self.cta_tile_shape_mnk[1] % 32 == 0, "GemmGatedSm90 requires tileN to be divisible by 32" self.cta_tile_shape_postact_mn = ( self.cta_tile_shape_mnk[0], @@ -527,13 +498,10 @@ def epi_to_underlying_arguments( ) # Assume all strides are divisible by 32 bits except the last stride new_stride = lambda t: tuple( - cute.assume(s, divby=32 // t.element_type.width) if not cute.is_static(s) else s - for s in t.stride + cute.assume(s, divby=32 // t.element_type.width) if not cute.is_static(s) else s for s in t.stride ) mRowVecBroadcast, mColVecBroadcast = [ - cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) - if t is not None - else None + cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) if t is not None else None for t in (args.mRowVecBroadcast, args.mColVecBroadcast) ] return self.EpilogueParams( @@ -555,12 +523,8 @@ def epi_smem_bytes_per_stage( epi_tile: cute.Tile, ) -> int: postact_dtype = args.mPostAct.element_type - postact_bytes_per_stage = (cute.size(cute.shape(epi_tile)) // 2) * ( - postact_dtype.width // 8 - ) - rowvec_colvec_bytes = GemmDefaultEpiMixin.epi_smem_bytes_per_stage( - args, cta_tile_shape_mnk, epi_tile - ) + postact_bytes_per_stage = (cute.size(cute.shape(epi_tile)) // 2) * (postact_dtype.width // 8) + rowvec_colvec_bytes = GemmDefaultEpiMixin.epi_smem_bytes_per_stage(args, cta_tile_shape_mnk, epi_tile) return postact_bytes_per_stage + rowvec_colvec_bytes @cute.jit @@ -649,9 +613,7 @@ def gemm_gated( # PostAct shape validation depends on varlen_m if cu_seqlens_m is not None: # varlen_m case: PostAct is 2D (total_m, n//2) - assert PostAct.dim() == 2 and PostAct.is_cuda, ( - "PostAct must be a 2D CUDA tensor for varlen_m" - ) + assert PostAct.dim() == 2 and PostAct.is_cuda, "PostAct must be a 2D CUDA tensor for varlen_m" assert PostAct.shape == ( M, N // 2, diff --git a/src/xorl/ops/quack/gemm_config.py b/src/xorl/ops/quack/gemm_config.py index fa19a28b..e43c544f 100644 --- a/src/xorl/ops/quack/gemm_config.py +++ b/src/xorl/ops/quack/gemm_config.py @@ -1,8 +1,8 @@ # Copyright (C) 2025, Fri Dao. import itertools -from typing import Optional, List, Literal -from functools import partial from dataclasses import dataclass +from functools import partial +from typing import List, Literal, Optional @dataclass(frozen=True) @@ -86,9 +86,7 @@ def get_all_configs( max_swizzle_size_vals = [4, 8, 16] GemmConfigCls = partial(GemmConfig, pingpong=False) # There's no pingpong on Sm100 return [ - GemmConfigCls( - tile_m=m, tile_n=n, cluster_m=cm, cluster_n=cn, swap_ab=sab, max_swizzle_size=ms - ) + GemmConfigCls(tile_m=m, tile_n=n, cluster_m=cm, cluster_n=cn, swap_ab=sab, max_swizzle_size=ms) for (m, n, (cm, cn)), sab, ms in itertools.product( tile_mn_cluster_vals, swap_ab_vals, max_swizzle_size_vals ) diff --git a/src/xorl/ops/quack/gemm_dact.py b/src/xorl/ops/quack/gemm_dact.py index 3812b0a3..caf2d190 100644 --- a/src/xorl/ops/quack/gemm_dact.py +++ b/src/xorl/ops/quack/gemm_dact.py @@ -1,36 +1,33 @@ # Copyright (c) 2025, Tri Dao. -from typing import Optional, Tuple, Callable, Type -from functools import partial -from dataclasses import dataclass import operator - -import torch -from torch import Tensor +from dataclasses import dataclass +from functools import partial +from typing import Callable, Optional, Tuple, Type import cutlass import cutlass.cute as cute -from cutlass import Int32, Float32, const_expr import cutlass.torch as cutlass_torch -from cutlass.cute.runtime import from_dlpack import cutlass.utils.blackwell_helpers as sm100_utils +import torch +from cutlass import Float32, Int32, const_expr +from cutlass.cute.runtime import from_dlpack +from torch import Tensor -from . import sm90_utils -from .sm90_utils import partition_for_epilogue -from .gemm_sm90 import GemmSm90 -from .gemm_sm100 import GemmSm100 -from .gemm_default_epi import GemmDefaultEpiMixin -from .gemm_act import GemmActMixin +from . import activation, layout_utils, sm90_utils from .cute_dsl_utils import ( ArgumentsBase, ParamsBase, - torch2cute_dtype_map, get_device_capacity, get_max_active_clusters, + torch2cute_dtype_map, ) +from .gemm_act import GemmActMixin +from .gemm_default_epi import GemmDefaultEpiMixin +from .gemm_sm90 import GemmSm90 +from .gemm_sm100 import GemmSm100 from .gemm_wrapper_utils import GemmWrapperBase +from .sm90_utils import partition_for_epilogue from .varlen_utils import VarlenManager -from . import layout_utils -from . import activation class GemmDActMixin(GemmActMixin): @@ -259,9 +256,7 @@ class EpilogueParams(ParamsBase): mColVecBroadcast: Optional[cute.Tensor] = None mColVecReduce: Optional[cute.Tensor] = None - def epi_to_underlying_arguments( - self, args: EpilogueArguments, *, loc=None, ip=None - ) -> EpilogueParams: + def epi_to_underlying_arguments(self, args: EpilogueArguments, *, loc=None, ip=None) -> EpilogueParams: self.postact_dtype = args.mPostAct.element_type self.postact_layout = cutlass.utils.LayoutEnum.from_tensor(args.mPostAct) # C and D are implicitly 2 16-bit elements packed into 32 bits, simply for the purpose @@ -284,13 +279,10 @@ def epi_to_underlying_arguments( ) # Assume all strides are divisible by 32 bits except the last stride new_stride = lambda t: tuple( - cute.assume(s, divby=32 // t.element_type.width) if not cute.is_static(s) else s - for s in t.stride + cute.assume(s, divby=32 // t.element_type.width) if not cute.is_static(s) else s for s in t.stride ) mRowVecBroadcast, mColVecBroadcast, mColVecReduce = [ - cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) - if t is not None - else None + cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) if t is not None else None for t in (args.mRowVecBroadcast, args.mColVecBroadcast, args.mColVecReduce) ] return self.EpilogueParams( @@ -342,9 +334,7 @@ def epi_begin( tDrColVecReduce = None if const_expr(params.mColVecReduce is not None): colvec_mma_layout = cute.make_layout(self.cta_tile_shape_mnk[:2], stride=(1, 0)) - tDrColVec_layout = partition_for_epilogue_fn( - cute.make_rmem_tensor(colvec_mma_layout, Float32) - ).layout + tDrColVec_layout = partition_for_epilogue_fn(cute.make_rmem_tensor(colvec_mma_layout, Float32)).layout tDrColVecReduce = cute.make_rmem_tensor(tDrColVec_layout, Float32) cute.filter_zeros(tDrColVecReduce).fill(0.0) return (*epi_tensors, tDrColVecReduce) @@ -384,13 +374,9 @@ def epi_visit_subtile( else: tDrColVec_mn = layout_utils.convert_layout_zero_stride(tDrColVec, tDrColVec.layout) tRS_rD_mn = layout_utils.convert_layout_zero_stride(tRS_rD, tDrColVec.layout) - tRS_rD_scaled_mn = layout_utils.convert_layout_zero_stride( - tRS_rD_scaled, tDrColVec.layout - ) + tRS_rD_scaled_mn = layout_utils.convert_layout_zero_stride(tRS_rD_scaled, tDrColVec.layout) for m in cutlass.range(cute.size(tDrColVec_mn, mode=[0]), unroll_full=True): - for n in cutlass.range( - cute.size(tDrColVec_mn, mode=[1]) // 2, unroll_full=True - ): + for n in cutlass.range(cute.size(tDrColVec_mn, mode=[1]) // 2, unroll_full=True): ( tRS_rD_scaled_mn[m, 2 * n], tRS_rD_scaled_mn[m, 2 * n + 1], @@ -406,9 +392,7 @@ def epi_visit_subtile( tRS_rdXY_f32x2[2 * i], tRS_rdXY_f32x2[2 * i + 1], tRS_rOut[i], - ) = params.act_bwd_fn( - tRS_rXY_f32x2[2 * i], tRS_rXY_f32x2[2 * i + 1], tRS_rD_scaled[i] - ) + ) = params.act_bwd_fn(tRS_rXY_f32x2[2 * i], tRS_rXY_f32x2[2 * i + 1], tRS_rD_scaled[i]) else: for i in cutlass.range(cute.size(tRS_rD) // 2): ( @@ -426,20 +410,14 @@ def epi_visit_subtile( for i in cutlass.range(cute.size(tDrColVecReduce), unroll_full=True): tDrColVecReduce[i] += tRS_rOut[i] * tRS_rD[i] else: - tDrColVecReduce_mn = layout_utils.convert_layout_zero_stride( - tDrColVecReduce, tDrColVecReduce.layout - ) + tDrColVecReduce_mn = layout_utils.convert_layout_zero_stride(tDrColVecReduce, tDrColVecReduce.layout) tRS_rD_mn = layout_utils.convert_layout_zero_stride(tRS_rD, tDrColVecReduce.layout) - tRS_rOut_mn = layout_utils.convert_layout_zero_stride( - tRS_rOut, tDrColVecReduce.layout - ) + tRS_rOut_mn = layout_utils.convert_layout_zero_stride(tRS_rOut, tDrColVecReduce.layout) for m in cutlass.range(cute.size(tDrColVecReduce_mn, mode=[0]), unroll_full=True): row_sum = cute.arch.mul_packed_f32x2( (tRS_rD_mn[m, 0], tRS_rD_mn[m, 1]), (tRS_rOut_mn[m, 0], tRS_rOut_mn[m, 1]) ) - for n in cutlass.range( - 1, cute.size(tDrColVecReduce_mn, mode=[1]) // 2, unroll_full=True - ): + for n in cutlass.range(1, cute.size(tDrColVecReduce_mn, mode=[1]) // 2, unroll_full=True): row_sum = cute.arch.fma_packed_f32x2( (tRS_rD_mn[m, 2 * n], tRS_rD_mn[m, 2 * n + 1]), (tRS_rOut_mn[m, 2 * n], tRS_rOut_mn[m, 2 * n + 1]), @@ -454,14 +432,10 @@ def epi_visit_subtile( tDrColVec_mn = layout_utils.convert_layout_zero_stride(tDrColVec, tDrColVec.layout) tRS_rOut_mn = layout_utils.convert_layout_zero_stride(tRS_rOut, tDrColVec.layout) for m in cutlass.range(cute.size(tDrColVec_mn, mode=[0]), unroll_full=True): - for n in cutlass.range( - cute.size(tDrColVec_mn, mode=[1]) // 2, unroll_full=True - ): - tRS_rOut_mn[m, 2 * n], tRS_rOut_mn[m, 2 * n + 1] = ( - cute.arch.mul_packed_f32x2( - (tRS_rOut_mn[m, 2 * n], tRS_rOut_mn[m, 2 * n + 1]), - (tDrColVec_mn[m, 0], tDrColVec_mn[m, 0]), - ) + for n in cutlass.range(cute.size(tDrColVec_mn, mode=[1]) // 2, unroll_full=True): + tRS_rOut_mn[m, 2 * n], tRS_rOut_mn[m, 2 * n + 1] = cute.arch.mul_packed_f32x2( + (tRS_rOut_mn[m, 2 * n], tRS_rOut_mn[m, 2 * n + 1]), + (tDrColVec_mn[m, 0], tDrColVec_mn[m, 0]), ) # Type conversion tRS_rdXY_f16x2 = cute.make_rmem_tensor(tRS_rdXY_f32x2.layout, implicit_dtype) @@ -496,20 +470,12 @@ def epi_end( tDrCVR_flt = cute.filter_zeros(tDrColVecReduce) if const_expr(self.arch != 100): for i in cutlass.range(cute.size(tDrCVR_flt), unroll_full=True): - tDrCVR_flt[i] = cute.arch.warp_reduction( - tDrCVR_flt[i], operator.add, threads_in_group=4 - ) + tDrCVR_flt[i] = cute.arch.warp_reduction(tDrCVR_flt[i], operator.add, threads_in_group=4) else: # Don't need warp_reduce since we load from tmem with one thread per row - assert self.d_layout.is_n_major_c(), ( - "GemmDGated only supports n-major output for now" - ) + assert self.d_layout.is_n_major_c(), "GemmDGated only supports n-major output for now" batch_idx = tile_coord_mnkl[3] - limit_n = ( - params.mColVecReduce.shape[2] - if not varlen_manager.varlen_m - else params.mColVecReduce.shape[1] - ) + limit_n = params.mColVecReduce.shape[2] if not varlen_manager.varlen_m else params.mColVecReduce.shape[1] if tile_coord_mnkl[1] < limit_n: if const_expr(not varlen_manager.varlen_m): mColVec = params.mColVecReduce[batch_idx, None, tile_coord_mnkl[1]] @@ -521,12 +487,10 @@ def epi_end( gColVec = cute.local_tile(mColVec, (tile_M,), (tile_coord_mnkl[0],)) limit_m = min(varlen_manager.len_m(batch_idx) - tile_coord_mnkl[0] * tile_M, tile_M) tDcCV = partition_for_epilogue_fn(cute.make_identity_tensor((tile_M, tile_N))) - tDrColVecReduce_m = layout_utils.convert_layout_zero_stride( - tDrColVecReduce, tDrColVecReduce.layout - )[None, 0] - tDcCV_m = layout_utils.convert_layout_zero_stride(tDcCV, tDrColVecReduce.layout)[ + tDrColVecReduce_m = layout_utils.convert_layout_zero_stride(tDrColVecReduce, tDrColVecReduce.layout)[ None, 0 ] + tDcCV_m = layout_utils.convert_layout_zero_stride(tDcCV, tDrColVecReduce.layout)[None, 0] if tDcCV_m[0][1] == 0: for m in cutlass.range(cute.size(tDcCV_m, mode=[0])): row_idx = tDcCV_m[m][0] @@ -661,9 +625,7 @@ def gemm_dgated( else None ), ) - scheduler_args = GemmWrapperBase.create_scheduler_args( - max_active_clusters, tile_count_semaphore - ) + scheduler_args = GemmWrapperBase.create_scheduler_args(max_active_clusters, tile_count_semaphore) # Create varlen arguments if needed (assumes persistent=True when varlen_m) varlen_args = GemmWrapperBase.create_varlen_args( diff --git a/src/xorl/ops/quack/gemm_default_epi.py b/src/xorl/ops/quack/gemm_default_epi.py index c5ad17ad..36b82497 100644 --- a/src/xorl/ops/quack/gemm_default_epi.py +++ b/src/xorl/ops/quack/gemm_default_epi.py @@ -1,19 +1,17 @@ # Copyright (c) 2025, Wentao Guo, Tri Dao. -from typing import Optional, Tuple -from functools import partial from dataclasses import dataclass - +from functools import partial +from typing import Optional, Tuple import cutlass import cutlass.cute as cute -from cutlass import Int32, Float32, Boolean, const_expr +from cutlass import Boolean, Float32, Int32, const_expr +from . import copy_utils, utils from .cute_dsl_utils import ArgumentsBase, ParamsBase from .gemm_sm90 import GemmSm90 from .gemm_sm100 import GemmSm100 from .sm90_utils import partition_for_epilogue -from . import utils -from . import copy_utils from .varlen_utils import VarlenManager @@ -35,18 +33,13 @@ class EpilogueParams(ParamsBase): mRowVecBroadcast: Optional[cute.Tensor] = None mColVecBroadcast: Optional[cute.Tensor] = None - def epi_to_underlying_arguments( - self, args: EpilogueArguments, *, loc=None, ip=None - ) -> EpilogueParams: + def epi_to_underlying_arguments(self, args: EpilogueArguments, *, loc=None, ip=None) -> EpilogueParams: # Assume all strides are divisible by 32 bits except the last stride new_stride = lambda t: tuple( - cute.assume(s, divby=32 // t.element_type.width) if not cute.is_static(s) else s - for s in t.stride + cute.assume(s, divby=32 // t.element_type.width) if not cute.is_static(s) else s for s in t.stride ) mRowVecBroadcast, mColVecBroadcast = [ - cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) - if t is not None - else None + cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) if t is not None else None for t in (args.mRowVecBroadcast, args.mColVecBroadcast) ] return self.EpilogueParams( @@ -107,9 +100,7 @@ def epi_begin( cute.copy(thr_copy_RV, tRVgRV, tRVsRV, pred=tRVpRV) # (CPY, CPY_M, CPY_N, EPI_M, EPI_N) tDsRowVec = partition_for_epilogue_fn( - cute.make_tensor( - sRowVec.iterator, cute.make_layout((tile_M, tile_N), stride=(0, 1)) - ) + cute.make_tensor(sRowVec.iterator, cute.make_layout((tile_M, tile_N), stride=(0, 1))) ) if const_expr(tiled_copy_t2r is not None): tDsRowVec = tiled_copy_r2s.retile(tDsRowVec) @@ -124,9 +115,7 @@ def epi_begin( if const_expr(not varlen_manager.varlen_m): mColVec = params.mColVecBroadcast[batch_idx, None] else: - mColVec = cute.domain_offset( - (varlen_manager.params.cu_seqlens_m[batch_idx],), params.mColVecBroadcast - ) + mColVec = cute.domain_offset((varlen_manager.params.cu_seqlens_m[batch_idx],), params.mColVecBroadcast) gColVec = cute.local_tile(mColVec, (tile_M,), (tile_coord_mnkl[0],)) tCVgCV = thr_copy_CV.partition_S(gColVec) tCVsCV = thr_copy_CV.partition_D(sColVec) @@ -137,9 +126,7 @@ def epi_begin( tCVpCV[0, m] = tCVcCV[0, m] < limit_m cute.copy(thr_copy_CV, tCVgCV, tCVsCV, pred=tCVpCV) tDsColVec = partition_for_epilogue_fn( - cute.make_tensor( - sColVec.iterator, cute.make_layout((tile_M, tile_N), stride=(1, 0)) - ) + cute.make_tensor(sColVec.iterator, cute.make_layout((tile_M, tile_N), stride=(1, 0))) ) if const_expr(tiled_copy_t2r is not None): tDsColVec = tiled_copy_r2s.retile(tDsColVec) @@ -154,9 +141,7 @@ def epi_begin_loop(self, params: EpilogueParams, epi_tensors, epi_coord: cute.Co alpha, beta, tDsRowVec, tDsColVec = epi_tensors tDrRowVec_cvt = None if const_expr(tDsRowVec is not None): - tDsRowVec_cur = cute.group_modes(tDsRowVec, 3, cute.rank(tDsRowVec))[ - None, None, None, epi_coord - ] + tDsRowVec_cur = cute.group_modes(tDsRowVec, 3, cute.rank(tDsRowVec))[None, None, None, epi_coord] # tDrRowVec = cute.make_fragment_like(tDsRowVec_cur) tDrRowVec = cute.make_rmem_tensor(tDsRowVec_cur.layout, tDsRowVec_cur.element_type) cute.autovec_copy(cute.filter_zeros(tDsRowVec_cur), cute.filter_zeros(tDrRowVec)) @@ -164,9 +149,7 @@ def epi_begin_loop(self, params: EpilogueParams, epi_tensors, epi_coord: cute.Co tDrRowVec_cvt.store(tDrRowVec.load().to(self.acc_dtype)) tDrColVec_cvt = None if const_expr(tDsColVec is not None): - tDsColVec_cur = cute.group_modes(tDsColVec, 3, cute.rank(tDsColVec))[ - None, None, None, epi_coord - ] + tDsColVec_cur = cute.group_modes(tDsColVec, 3, cute.rank(tDsColVec))[None, None, None, epi_coord] # This somehow doesn't work, some dim with stride 0 turns to non-zero stride # tDrRowVec = cute.make_fragment_like(tDsRowVec_cur) tDrColVec = cute.make_rmem_tensor(tDsColVec_cur.layout, tDsColVec_cur.element_type) @@ -214,25 +197,15 @@ def epi_smem_bytes_per_stage( ) -> int: row_vec_smem_size = 0 if args.mRowVecBroadcast is None else cta_tile_shape_mnk[1] col_vec_smem_size = 0 if args.mColVecBroadcast is None else cta_tile_shape_mnk[0] - row_vec_dtype = ( - args.mRowVecBroadcast.element_type if args.mRowVecBroadcast is not None else Float32 - ) - col_vec_dtype = ( - args.mColVecBroadcast.element_type if args.mColVecBroadcast is not None else Float32 - ) - return ( - row_vec_smem_size * row_vec_dtype.width + col_vec_smem_size * col_vec_dtype.width - ) // 8 + row_vec_dtype = args.mRowVecBroadcast.element_type if args.mRowVecBroadcast is not None else Float32 + col_vec_dtype = args.mColVecBroadcast.element_type if args.mColVecBroadcast is not None else Float32 + return (row_vec_smem_size * row_vec_dtype.width + col_vec_smem_size * col_vec_dtype.width) // 8 def epi_get_smem_struct(self, params: EpilogueParams): row_vec_smem_size = 0 if params.mRowVecBroadcast is None else self.cta_tile_shape_mnk[1] col_vec_smem_size = 0 if params.mColVecBroadcast is None else self.cta_tile_shape_mnk[0] - row_vec_dtype = ( - params.mRowVecBroadcast.element_type if params.mRowVecBroadcast is not None else Float32 - ) - col_vec_dtype = ( - params.mColVecBroadcast.element_type if params.mColVecBroadcast is not None else Float32 - ) + row_vec_dtype = params.mRowVecBroadcast.element_type if params.mRowVecBroadcast is not None else Float32 + col_vec_dtype = params.mColVecBroadcast.element_type if params.mColVecBroadcast is not None else Float32 @cute.struct class EpiSharedStorage: diff --git a/src/xorl/ops/quack/gemm_interface.py b/src/xorl/ops/quack/gemm_interface.py index 1b33101c..441331c3 100644 --- a/src/xorl/ops/quack/gemm_interface.py +++ b/src/xorl/ops/quack/gemm_interface.py @@ -1,18 +1,17 @@ # Copyright (c) 2025, Tri Dao -from typing import Optional, Tuple, Literal from functools import partial +from typing import Literal, Optional, Tuple import torch import torch.nn.functional as F from torch import Tensor -from .gemm_config import GemmConfig, get_all_configs - -from .autotuner import autotune, AutotuneConfig +from .autotuner import AutotuneConfig, autotune from .cute_dsl_utils import get_device_capacity from .gemm import gemm as gemm_sm90_sm100 from .gemm_act import gemm_act as gemm_act_sm90_sm100 from .gemm_act import gemm_gated as gemm_gated_sm90_sm100 +from .gemm_config import GemmConfig, get_all_configs from .gemm_dact import gemm_dact as gemm_dact_sm90_sm100 from .gemm_dact import gemm_dgated as gemm_dgated_sm90_sm100 from .gemm_symmetric import gemm_symmetric as gemm_symmetric_sm90_sm100 @@ -55,18 +54,10 @@ def prune_invalid_gemm_configs(configs, named_args: dict, **kwargs): if varlen_m or gather_A: # Doesn't support swap_ab configs = [conf for conf in configs if not conf.kwargs["config"].swap_ab] if gather_A: - configs = [ - conf - for conf in configs - if conf.kwargs["config"].cluster_n == 1 - ] + configs = [conf for conf in configs if conf.kwargs["config"].cluster_n == 1] if get_device_capacity(kwargs["A"].device)[0] == 9: # tile_n == 208 causes register spills, as gather_A requires more registers for the producer - configs = [ - conf - for conf in configs - if conf.kwargs["config"].tile_n != 208 - ] + configs = [conf for conf in configs if conf.kwargs["config"].tile_n != 208] return configs @@ -122,9 +113,7 @@ def gemm_tuned( else: out_shape = (batch_size, A.shape[-2], B.shape[-2]) assert out.shape == out_shape, f"out shape mismatch: {out.shape} vs {out_shape}" - tile_count_semaphore = ( - torch.zeros(1, dtype=torch.int32, device=A.device) if dynamic_scheduler else None - ) + tile_count_semaphore = torch.zeros(1, dtype=torch.int32, device=A.device) if dynamic_scheduler else None gemm_sm90_sm100( A if not config.swap_ab else B, B if not config.swap_ab else A, @@ -192,9 +181,7 @@ def gemm_act_tuned( PostAct = postact_out if bias is not None and bias.ndim == 1: bias = bias.unsqueeze(0) # (L, N) - tile_count_semaphore = ( - torch.zeros(1, dtype=torch.int32, device=A.device) if dynamic_scheduler else None - ) + tile_count_semaphore = torch.zeros(1, dtype=torch.int32, device=A.device) if dynamic_scheduler else None gemm_act_sm90_sm100( A if not config.swap_ab else B, B if not config.swap_ab else A, @@ -255,9 +242,7 @@ def gemm_dact_tuned( PostAct = postact_out.unsqueeze(0) else: PostAct = postact_out - tile_count_semaphore = ( - torch.zeros(1, dtype=torch.int32, device=A.device) if dynamic_scheduler else None - ) + tile_count_semaphore = torch.zeros(1, dtype=torch.int32, device=A.device) if dynamic_scheduler else None gemm_dact_sm90_sm100( A if not config.swap_ab else B, B if not config.swap_ab else A, @@ -306,9 +291,7 @@ def gemm( # For varlen_k, the first dimension is always A.shape[0] (M dimension) out_shape = (L, A.shape[0], B.shape[-1]) else: - out_shape = ( - (A.shape[0], B.shape[-1]) if A.ndim == 2 else (A.shape[0], A.shape[-2], B.shape[-1]) - ) + out_shape = (A.shape[0], B.shape[-1]) if A.ndim == 2 else (A.shape[0], A.shape[-2], B.shape[-1]) out = torch.empty(out_shape, dtype=out_dtype, device=A.device) alpha_tensor = alpha if not isinstance(alpha, float) else None alpha = alpha if isinstance(alpha, float) else 1.0 @@ -458,9 +441,7 @@ def gemm_add( # For varlen_k, the first dimension is always A.shape[0] (M dimension) out_shape = (L, A.shape[0], B.shape[-1]) else: - out_shape = ( - (A.shape[0], B.shape[-1]) if A.ndim == 2 else (A.shape[0], A.shape[-2], B.shape[-1]) - ) + out_shape = (A.shape[0], B.shape[-1]) if A.ndim == 2 else (A.shape[0], A.shape[-2], B.shape[-1]) out = torch.empty(out_shape, dtype=out_dtype, device=A.device) add_to_output = C is out and isinstance(beta, float) and beta == 1.0 and cu_seqlens_m is None alpha_tensor = alpha if not isinstance(alpha, float) else None @@ -552,9 +533,7 @@ def gemm_add_ref( if isinstance(alpha, float) and isinstance(beta, float): out = torch.addmm(C, A, B, out_dtype=out_dtype, alpha=alpha, beta=beta, out=out) else: - out_dtype = ( - out.dtype if out is not None else (out_dtype if out_dtype is not None else A.dtype) - ) + out_dtype = out.dtype if out is not None else (out_dtype if out_dtype is not None else A.dtype) result = (alpha * (A @ B) + beta * C).to(out_dtype) if out is not None: out.copy_(result) @@ -816,9 +795,7 @@ def gemm_dact( dx_out = torch.empty(out_shape, dtype=out_dtype, device=A.device) if postact_out is None: postact_out = torch.empty(out_shape, dtype=postact_dtype, device=A.device) - gemm_dact_out( - A, B, PreAct, dx_out, postact_out, activation, cu_seqlens_m, A_idx, dynamic_scheduler, tuned - ) + gemm_dact_out(A, B, PreAct, dx_out, postact_out, activation, cu_seqlens_m, A_idx, dynamic_scheduler, tuned) return dx_out, postact_out @@ -984,9 +961,7 @@ def gemm_symmetric_out( out = out.unsqueeze(0) else: out = out - tile_count_semaphore = ( - torch.zeros(1, dtype=torch.int32, device=A.device) if dynamic_scheduler else None - ) + tile_count_semaphore = torch.zeros(1, dtype=torch.int32, device=A.device) if dynamic_scheduler else None gemm_symmetric_sm90_sm100( A, B, @@ -1028,16 +1003,12 @@ def gemm_symmetric( alpha_val = alpha if isinstance(alpha, float) else 1.0 beta_val = beta if isinstance(beta, float) else 1.0 - gemm_symmetric_out( - A, B, out, C, dynamic_scheduler=dynamic_scheduler, alpha=alpha_val, beta=beta_val - ) + gemm_symmetric_out(A, B, out, C, dynamic_scheduler=dynamic_scheduler, alpha=alpha_val, beta=beta_val) return out @autotune( - configs=[ - AutotuneConfig(config=c) for c in get_all_configs(default_device_capacity[0], "gated") - ], + configs=[AutotuneConfig(config=c) for c in get_all_configs(default_device_capacity[0], "gated")], key=["activation", "dynamic_scheduler"], prune_configs_by={"early_config_prune": prune_invalid_gemm_configs}, ) @@ -1078,9 +1049,7 @@ def gemm_gated_tuned( PostAct = postact_out if bias is not None and bias.ndim == 1: bias = bias.unsqueeze(0) # (L, N) - tile_count_semaphore = ( - torch.zeros(1, dtype=torch.int32, device=A.device) if dynamic_scheduler else None - ) + tile_count_semaphore = torch.zeros(1, dtype=torch.int32, device=A.device) if dynamic_scheduler else None gemm_gated_sm90_sm100( A if not config.swap_ab else B, B if not config.swap_ab else A, @@ -1112,9 +1081,7 @@ def prune_invalid_gemm_dgated_configs(configs, named_args: dict, **kwargs): @autotune( - configs=[ - AutotuneConfig(config=c) for c in get_all_configs(default_device_capacity[0], "dgated") - ], + configs=[AutotuneConfig(config=c) for c in get_all_configs(default_device_capacity[0], "dgated")], key=["activation", "colvec_reduce", "dynamic_scheduler"], prune_configs_by={"early_config_prune": prune_invalid_gemm_dgated_configs}, ) @@ -1170,9 +1137,7 @@ def gemm_dgated_tuned( colvec_reduce_partial = torch.empty(colvec_shape, dtype=torch.float32, device=A.device) else: colvec_reduce_partial = None - tile_count_semaphore = ( - torch.zeros(1, dtype=torch.int32, device=A.device) if dynamic_scheduler else None - ) + tile_count_semaphore = torch.zeros(1, dtype=torch.int32, device=A.device) if dynamic_scheduler else None gemm_dgated_sm90_sm100( A if not config.swap_ab else B, B if not config.swap_ab else A, @@ -1209,9 +1174,7 @@ def gemm_gated( bias: Optional[Tensor] = None, # (N,) or (L, N) activation: Literal["swiglu", "swiglu_oai", "reglu", "geglu", "glu"] = "swiglu", preact_out: Optional[Tensor] = None, # (M, N) or (L, M, N) or (total_M, N) if varlen_m - postact_out: Optional[ - Tensor - ] = None, # (M, N//2) or (L, M, N//2) or (total_M, N//2) if varlen_m + postact_out: Optional[Tensor] = None, # (M, N//2) or (L, M, N//2) or (total_M, N//2) if varlen_m out_dtype: Optional[torch.dtype] = None, postact_dtype: Optional[torch.dtype] = None, cu_seqlens_m: Optional[Tensor] = None, diff --git a/src/xorl/ops/quack/gemm_sm100.py b/src/xorl/ops/quack/gemm_sm100.py index 97798bed..dac96acc 100644 --- a/src/xorl/ops/quack/gemm_sm100.py +++ b/src/xorl/ops/quack/gemm_sm100.py @@ -2,37 +2,37 @@ # https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py import argparse -from typing import Optional, Type, Tuple, Union, Callable, Literal from functools import partial +from typing import Callable, Literal, Optional, Tuple, Type, Union import cuda.bindings.driver as cuda -import torch - import cutlass import cutlass.cute as cute -from cutlass.cute.nvgpu import cpasync, tcgen05 -import cutlass.torch as cutlass_torch import cutlass.pipeline as pipeline -from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.torch as cutlass_torch import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.utils.blockscaled_layout as blockscaled_utils +import torch +from cutlass import Boolean, Float32, Int32, const_expr +from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.nvgpu.warp import ( LdMatrix8x8x16bOp, LdMatrix16x16x8bOp, StMatrix8x8x16bOp, StMatrix16x8x8bOp, ) -from cutlass import Int32, Float32, Boolean, const_expr -from cutlass.utils import LayoutEnum from cutlass.cute.runtime import from_dlpack, make_ptr +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from cutlass.utils import LayoutEnum +from . import copy_utils +from . import sm100_utils as quack_sm100_utils +from .cute_dsl_utils import ArgumentsBase, ParamsBase +from .gemm_sm90 import GemmSm90, NamedBarrierGemm from .pipeline import PipelineTmaCpAsyncUmma -from .cute_dsl_utils import ParamsBase, ArgumentsBase from .tile_scheduler import TileSchedulerOptions from .varlen_utils import VarlenArguments, VarlenManager -from .gemm_sm90 import GemmSm90, NamedBarrierGemm -from . import copy_utils -from . import sm100_utils as quack_sm100_utils + # return PipelineStateWAdvance instead of PipelineState @@ -339,9 +339,7 @@ def _setup_attributes(self, epilogue_args: EpilogueArguments, varlen_args: Varle # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory prefetch_A_idx = ( - None - if not self.gather_A - else ("varlen_m" if varlen_args.mCuSeqlensM is not None else "varlen_k") + None if not self.gather_A else ("varlen_m" if varlen_args.mCuSeqlensM is not None else "varlen_k") ) ( self.num_acc_stage, @@ -368,9 +366,7 @@ def _setup_attributes(self, epilogue_args: EpilogueArguments, varlen_args: Varle ) self.sched_stage = 1 self.a_prefetch_stage = ( - 0 - if not self.gather_A - else (2 if varlen_args.mCuSeqlensM is not None else self.ab_stage) + 0 if not self.gather_A else (2 if varlen_args.mCuSeqlensM is not None else self.ab_stage) ) # Compute A/B/SFA/SFB/C shared memory layout @@ -461,9 +457,7 @@ def __call__( self.b_dtype = mB.element_type self.d_dtype = mD.element_type if mD is not None else None self.c_dtype = mC.element_type if mC is not None else None - self.sf_dtype: Optional[Type[cutlass.Numeric]] = ( - mSFA.element_type if mSFA is not None else None - ) + self.sf_dtype: Optional[Type[cutlass.Numeric]] = mSFA.element_type if mSFA is not None else None self.a_layout = LayoutEnum.from_tensor(mA) self.b_layout = LayoutEnum.from_tensor(mB) self.d_layout = LayoutEnum.from_tensor(mD) if mD is not None else None @@ -482,14 +476,11 @@ def __call__( # Assume all strides are divisible by 128 bits except the last stride def new_stride(t: cute.Tensor): return tuple( - cute.assume(s, divby=128 // t.element_type.width) if not cute.is_static(s) else s - for s in t.stride + cute.assume(s, divby=128 // t.element_type.width) if not cute.is_static(s) else s for s in t.stride ) mA, mD = [ - cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) - if t is not None - else None + cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) if t is not None else None for t in (mA, mD) ] @@ -512,9 +503,7 @@ def new_stride(t: cute.Tensor): b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) tma_atom_a, tma_tensor_a = None, None if const_expr(not self.gather_A): - a_op = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mnk, self.tiled_mma.thr_id - ) + a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mnk, self.tiled_mma.thr_id) tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( a_op, mA, @@ -524,9 +513,7 @@ def new_stride(t: cute.Tensor): self.cluster_layout_vmnk.shape, internal_type=(cutlass.TFloat32 if mA.element_type is Float32 else None), ) - b_op = sm100_utils.cluster_shape_to_tma_atom_B( - self.cluster_shape_mnk, self.tiled_mma.thr_id - ) + b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mnk, self.tiled_mma.thr_id) tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( b_op, mB, @@ -541,9 +528,7 @@ def new_stride(t: cute.Tensor): tma_atom_sfb, tma_tensor_sfb = None, None if const_expr(self.blockscaled): # Setup TMA load for SFA - sfa_op = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mnk, self.tiled_mma.thr_id - ) + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mnk, self.tiled_mma.thr_id) sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)) tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( sfa_op, @@ -555,9 +540,7 @@ def new_stride(t: cute.Tensor): internal_type=cutlass.Int16, ) # Setup TMA load for SFB - sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB( - self.cluster_shape_mnk, self.tiled_mma.thr_id - ) + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mnk, self.tiled_mma.thr_id) sfb_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( sfb_op, @@ -601,27 +584,19 @@ def new_stride(t: cute.Tensor): TileSchedulerCls = self.get_scheduler_class(varlen_m=varlen_args.mCuSeqlensM is not None) tile_sched_args = self.get_scheduler_arguments(mA, mB, mD, scheduler_args, varlen_args) tile_sched_params = TileSchedulerCls.to_underlying_arguments(tile_sched_args) - grid = TileSchedulerCls.get_grid_shape( - tile_sched_params, scheduler_args.max_active_clusters - ) + grid = TileSchedulerCls.get_grid_shape(tile_sched_params, scheduler_args.max_active_clusters) self.buffer_align_bytes = 1024 epi_smem_size = cute.cosize(self.epi_smem_layout_staged) if mD is not None else 0 epi_c_smem_size = cute.cosize(self.epi_c_smem_layout_staged) if mC is not None else 0 sf_dtype = self.sf_dtype if const_expr(self.blockscaled) else cutlass.Float8E8M0FNU - sfa_smem_size = ( - cute.cosize(self.sfa_smem_layout_staged) if const_expr(self.blockscaled) else 0 - ) - sfb_smem_size = ( - cute.cosize(self.sfb_smem_layout_staged) if const_expr(self.blockscaled) else 0 - ) + sfa_smem_size = cute.cosize(self.sfa_smem_layout_staged) if const_expr(self.blockscaled) else 0 + sfb_smem_size = cute.cosize(self.sfb_smem_layout_staged) if const_expr(self.blockscaled) else 0 a_idx_smem_size = 0 if const_expr(self.gather_A): a_idx_smem_size = self.a_prefetch_stage * ( - self.cta_tile_shape_mnk[0] - if varlen_args.mCuSeqlensM is not None - else self.cta_tile_shape_mnk[2] + self.cta_tile_shape_mnk[0] if varlen_args.mCuSeqlensM is not None else self.cta_tile_shape_mnk[2] ) # Define shared storage for kernel @@ -631,24 +606,18 @@ class SharedStorage: epi_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.epi_c_stage * 2] acc_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] sched_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.sched_stage * 2] - a_prefetch_pipeline_array_ptr: cute.struct.MemRange[ - cutlass.Int64, self.a_prefetch_stage * 2 - ] + a_prefetch_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.a_prefetch_stage * 2] sched_data: cute.struct.MemRange[Int32, self.sched_stage * 12] tmem_dealloc_mbar_ptr: cutlass.Int64 tmem_holding_buf: Int32 sAIdx: cute.struct.Align[cute.struct.MemRange[Int32, a_idx_smem_size], 16] # (EPI_TILE_M, EPI_TILE_N, STAGE) sD: cute.struct.Align[ - cute.struct.MemRange[ - self.d_dtype if self.d_dtype is not None else Int32, epi_smem_size - ], + cute.struct.MemRange[self.d_dtype if self.d_dtype is not None else Int32, epi_smem_size], self.buffer_align_bytes, ] sC: cute.struct.Align[ - cute.struct.MemRange[ - self.c_dtype if self.c_dtype is not None else Int32, epi_c_smem_size - ], + cute.struct.MemRange[self.c_dtype if self.c_dtype is not None else Int32, epi_c_smem_size], self.buffer_align_bytes, ] epi: self.epi_get_smem_struct(epilogue_params) @@ -864,9 +833,7 @@ def kernel( epi_smem_tensors = self.epi_get_smem_tensors(epilogue_params, storage) thr_mma = tiled_mma.get_slice(mma_tile_coord_v) - thr_mma_sfb = ( - tiled_mma_sfb.get_slice(mma_tile_coord_v) if const_expr(self.blockscaled) else None - ) + thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v) if const_expr(self.blockscaled) else None # (MMA, MMA_M, MMA_N) acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) @@ -879,16 +846,12 @@ def kernel( self.num_epi_tensormaps, # Only used if not varlen_m len_m_static=Int32( - mA_mkl.shape[0] - if varlen_k or varlen_params.mAIdx is None - else varlen_params.mAIdx.shape[0] + mA_mkl.shape[0] if varlen_k or varlen_params.mAIdx is None else varlen_params.mAIdx.shape[0] ), len_k_static=Int32(mA_mkl.shape[1]), ) - TileSchedulerCls = partial( - TileSchedulerCls.create, tile_sched_params, sched_data, sched_pipeline - ) + TileSchedulerCls = partial(TileSchedulerCls.create, tile_sched_params, sched_data, sched_pipeline) epi_load_barrier = None if const_expr(has_C): @@ -910,9 +873,7 @@ def kernel( block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) block_in_cluster_coord_sfb_vmnk = None if const_expr(self.blockscaled): - block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord( - cta_rank_in_cluster - ) + block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord(cta_rank_in_cluster) a_mcast_mask, b_mcast_mask = None, None sfa_mcast_mask, sfb_mcast_mask = None, None if const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): @@ -933,9 +894,7 @@ def kernel( # Persistent tile scheduling loop tile_scheduler = TileSchedulerCls() work_tile = tile_scheduler.initial_work_tile_info() - ab_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self.ab_stage - ) + ab_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.ab_stage) if const_expr(varlen_k): # wait tensormap initialization complete before update varlen_manager.fence_tensormap_init() @@ -989,9 +948,7 @@ def kernel( varlen_manager.fence_tensormap_update_AB(is_tma_warp) len_k = varlen_manager.len_k(batch_idx) # TMA load A partition_S/D - a_cta_layout = cute.make_layout( - cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape - ) + a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) copy_A = None if const_expr(not self.gather_A): # (MMA, MMA_M, MMA_K, RestK) @@ -1016,9 +973,7 @@ def kernel( copy_B, _, _ = copy_utils.tma_get_copy_fn( tma_atom_b, cta_coord=block_in_cluster_coord_vmnk[1], - cta_layout=cute.make_layout( - cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape - ), + cta_layout=cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape), src_tensor=tCgB, dst_tensor=sB, mcast_mask=b_mcast_mask, @@ -1038,9 +993,7 @@ def kernel( # tma_desc_ptr=tma_desc_sfa_ptr, ) # TMA load SFB partition_S/D - sfb_cta_layout = cute.make_layout( - cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape - ) + sfb_cta_layout = cute.make_layout(cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape) copy_SFB, _, _ = copy_utils.tma_get_copy_fn( tma_atom_sfb, cta_coord=block_in_cluster_coord_sfb_vmnk[1], @@ -1075,16 +1028,11 @@ def kernel( ab_pipeline.producer_tail(ab_producer_state) if const_expr(self.gather_A): - if ( - warp_idx >= self.ab_load_warp_id + 1 - and warp_idx < self.ab_load_warp_id + self.num_ab_load_warps - ): + if warp_idx >= self.ab_load_warp_id + 1 and warp_idx < self.ab_load_warp_id + self.num_ab_load_warps: # Persistent tile scheduling loop tile_scheduler = TileSchedulerCls() work_tile = tile_scheduler.initial_work_tile_info() - ab_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self.ab_stage - ) + ab_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.ab_stage) a_prefetch_consumer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Consumer, self.a_prefetch_stage ) @@ -1099,9 +1047,7 @@ def kernel( else: assert varlen_k # (tile_M, K) - mA_mk = cute.local_tile( - mA_mkl, (self.cta_tile_shape_mnk[0],), (tile_coord_mnkl[0], None) - ) + mA_mk = cute.local_tile(mA_mkl, (self.cta_tile_shape_mnk[0],), (tile_coord_mnkl[0], None)) # Partition global tensor for TiledMMA_A/B/D len_m = varlen_manager.len_m(batch_idx) len_k = varlen_manager.len_k(batch_idx) @@ -1162,9 +1108,7 @@ def kernel( tiled_copy_AIdx = copy_utils.tiled_copy_1d(Int32, num_threads=32, is_async=True) thr_copy_AIdx = tiled_copy_AIdx.get_slice(cute.arch.lane_idx()) tAsAIdx = thr_copy_AIdx.partition_D(sAIdx) - tAcAIdx = thr_copy_AIdx.partition_S( - cute.make_identity_tensor(tile_M if varlen_m else tile_K) - ) + tAcAIdx = thr_copy_AIdx.partition_S(cute.make_identity_tensor(tile_M if varlen_m else tile_K)) # Persistent tile scheduling loop tile_scheduler = TileSchedulerCls(is_scheduler_warp=is_scheduler_warp) work_tile = tile_scheduler.initial_work_tile_info() @@ -1236,9 +1180,7 @@ def kernel( # Specialized TMA epi load warp if const_expr(mC_mnl is not None): if warp_idx == self.epi_load_warp_id: - epi_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self.epi_c_stage - ) + epi_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.epi_c_stage) do_epi_load_barrier_wait = Boolean(True) # Persistent tile scheduling loop tile_scheduler = TileSchedulerCls() @@ -1333,12 +1275,8 @@ def kernel( # Persistent tile scheduling loop tile_scheduler = TileSchedulerCls() work_tile = tile_scheduler.initial_work_tile_info() - ab_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.ab_stage - ) - acc_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self.num_acc_stage - ) + ab_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.ab_stage) + acc_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) while work_tile.is_valid_tile: # Get tile coord from tile scheduler tile_coord_mnkl = work_tile.tile_idx @@ -1385,9 +1323,7 @@ def kernel( tmem.wait_for_alloc() is_tma_warp = Boolean(warp_idx == self.epilog_warp_id[0]) - varlen_manager.init_tensormap_epi( - tma_atom_d, self.epi_get_tma_atoms(epilogue_params), is_tma_warp - ) + varlen_manager.init_tensormap_epi(tma_atom_d, self.epi_get_tma_atoms(epilogue_params), is_tma_warp) tma_desc_d_ptr = varlen_manager.get_tma_desc_d_ptr() tma_desc_epi_ptrs = varlen_manager.get_tma_desc_epi_ptrs() @@ -1421,13 +1357,9 @@ def kernel( # Persistent tile scheduling loop tile_scheduler = TileSchedulerCls() work_tile = tile_scheduler.initial_work_tile_info() - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage - ) + acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) epi_store_pipeline = self.make_epi_store_pipeline() - epi_read_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.epi_c_stage - ) + epi_read_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.epi_c_stage) if const_expr(varlen_m): # wait tensormap initialization complete before update varlen_manager.fence_tensormap_init() @@ -1813,9 +1745,7 @@ def epilog_smem_load_and_partition( tRS_rD_layout: cutlass.Layout, tidx: Int32, ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: - copy_atom_r2s = sm100_utils.get_smem_store_op( - c_layout, dtype, self.acc_dtype, tiled_copy_t2r - ) + copy_atom_r2s = sm100_utils.get_smem_store_op(c_layout, dtype, self.acc_dtype, tiled_copy_t2r) store_op = copy_atom_r2s.op # m8n8 16-bit path if isinstance(store_op, StMatrix8x8x16bOp): @@ -1854,16 +1784,12 @@ def make_ab_pipeline( if const_expr(not self.gather_A): producer_cnt = 1 else: - producer_cnt = (self.num_ab_load_warps - 1) * 32 + ( - 1 if const_expr(not self.use_2cta_instrs) else 2 - ) + producer_cnt = (self.num_ab_load_warps - 1) * 32 + (1 if const_expr(not self.use_2cta_instrs) else 2) ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, producer_cnt) # Each warp will contribute to the arrive count with the number of mcast size mcast_size = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 consumer_arrive_cnt = mcast_size - ab_pipeline_consumer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, consumer_arrive_cnt - ) + ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_arrive_cnt) if const_expr(not self.gather_A): pipeline_ab = pipeline.PipelineTmaUmma.create( barrier_storage=ab_pipeline_mbar_ptr, @@ -1882,9 +1808,7 @@ def make_ab_pipeline( consumer_group=ab_pipeline_consumer_group, tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cluster_layout_vmnk, - producer_drop_count=None - if not self.use_2cta_instrs - else (2 if not is_leader_cta else 0), + producer_drop_count=None if not self.use_2cta_instrs else (2 if not is_leader_cta else 0), defer_sync=True, ) return pipeline_ab @@ -1894,9 +1818,7 @@ def make_acc_pipeline( ) -> pipeline.PipelineAsync: acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) num_acc_consumer_threads = self.num_epi_warps * (2 if self.use_2cta_instrs else 1) - acc_pipeline_consumer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, num_acc_consumer_threads - ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_consumer_threads) return pipeline.PipelineUmmaAsync.create( barrier_storage=acc_pipeline_mbar_ptr, num_stages=self.num_acc_stage, @@ -1916,15 +1838,11 @@ def make_sched_pipeline( sched_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) cluster_size = cute.size(cluster_layout_mnk) # Each warp will contribute 1 to the arrive count - warps_per_cta = self.num_ab_load_warps + len( - (self.mma_warp_id, *self.epilog_warp_id, self.scheduler_warp_id) - ) + warps_per_cta = self.num_ab_load_warps + len((self.mma_warp_id, *self.epilog_warp_id, self.scheduler_warp_id)) if has_C: warps_per_cta += 1 consumer_arrive_cnt = warps_per_cta * cluster_size - sched_pipeline_consumer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, consumer_arrive_cnt - ) + sched_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_arrive_cnt) return pipeline.PipelineAsync.create( barrier_storage=sched_pipeline_mbar_ptr, num_stages=self.sched_stage, @@ -1936,15 +1854,11 @@ def make_sched_pipeline( ) @cute.jit - def make_a_prefetch_pipeline( - self, a_prefetch_pipeline_mbar_ptr: cute.Pointer - ) -> pipeline.PipelineAsync: + def make_a_prefetch_pipeline(self, a_prefetch_pipeline_mbar_ptr: cute.Pointer) -> pipeline.PipelineAsync: producer_cnt = 32 a_prefetch_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, producer_cnt) consumer_arrive_cnt = self.num_ab_load_warps - 1 - a_prefetch_consumer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, consumer_arrive_cnt - ) + a_prefetch_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_arrive_cnt) return pipeline.PipelineCpAsync.create( barrier_storage=a_prefetch_pipeline_mbar_ptr, num_stages=self.a_prefetch_stage, @@ -2023,14 +1937,10 @@ def _compute_stages( 1, # a tmp 1 stage is provided ) d_smem_layout_staged_one = ( - sm100_utils.make_smem_layout_epi(d_dtype, d_layout, epi_tile, 1) - if d_dtype is not None - else None + sm100_utils.make_smem_layout_epi(d_dtype, d_layout, epi_tile, 1) if d_dtype is not None else None ) c_smem_layout_staged_one = ( - sm100_utils.make_smem_layout_epi(c_dtype, c_layout, epi_tile, 1) - if c_dtype is not None - else None + sm100_utils.make_smem_layout_epi(c_dtype, c_layout, epi_tile, 1) if c_dtype is not None else None ) if const_expr(blockscaled): sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( @@ -2046,21 +1956,19 @@ def _compute_stages( 1, # a tmp 1 stage is provided ) - ab_bytes_per_stage = cute.size_in_bytes( - a_dtype, a_smem_layout_staged_one - ) + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + ab_bytes_per_stage = cute.size_in_bytes(a_dtype, a_smem_layout_staged_one) + cute.size_in_bytes( + b_dtype, b_smem_layout_staged_one + ) if const_expr(prefetch_A_idx == "varlen_k"): # Need smem to prefetch A indices ab_bytes_per_stage += Int32.width // 8 * cta_tile_shape_mnk[2] if const_expr(blockscaled): - ab_bytes_per_stage += cute.size_in_bytes( - sf_dtype, sfa_smem_layout_staged_one - ) + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ab_bytes_per_stage += cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) + cute.size_in_bytes( + sf_dtype, sfb_smem_layout_staged_one + ) mbar_helpers_bytes = 1024 if const_expr(prefetch_A_idx == "varlen_m"): mbar_helpers_bytes += Int32.width // 8 * cta_tile_shape_mnk[0] * 2 - d_bytes_per_stage = ( - cute.size_in_bytes(d_dtype, d_smem_layout_staged_one) if d_dtype is not None else 0 - ) + d_bytes_per_stage = cute.size_in_bytes(d_dtype, d_smem_layout_staged_one) if d_dtype is not None else 0 epi_bytes_per_stage = d_bytes_per_stage + cls.epi_smem_bytes_per_stage( epilogue_args, cta_tile_shape_mnk, epi_tile ) @@ -2387,14 +2295,10 @@ def can_implement( if not GemmSm100.is_valid_dtypes(ab_dtype, ab_dtype, acc_dtype, d_dtype, a_major, b_major): can_implement = False # Skip invalid mma tile shape and cluster shape - if not GemmSm100.is_valid_mma_tiler_and_cluster_shape( - mma_tiler_mn, cluster_shape_mn, blockscaled=False - ): + if not GemmSm100.is_valid_mma_tiler_and_cluster_shape(mma_tiler_mn, cluster_shape_mn, blockscaled=False): can_implement = False # Skip illegal problem shape for load/store alignment - if not GemmSm100.is_valid_tensor_alignment( - m, n, k, l, ab_dtype, d_dtype, a_major, b_major, d_major - ): + if not GemmSm100.is_valid_tensor_alignment(m, n, k, l, ab_dtype, d_dtype, a_major, b_major, d_major): can_implement = False return can_implement @@ -2500,11 +2404,7 @@ def create_and_permute_tensor(l, mode0, mode1, is_mode0_major, dtype, is_dynamic is_unsigned = dtype in {cutlass.Uint8} # Temporarily use uint8 as torch does not support fp8 type torch_dtype = cutlass_torch.dtype(dtype) - gen_dtype = ( - torch_dtype - if dtype not in {cutlass.Float8E5M2, cutlass.Float8E4M3FN} - else torch.bfloat16 - ) + gen_dtype = torch_dtype if dtype not in {cutlass.Float8E5M2, cutlass.Float8E4M3FN} else torch.bfloat16 # Create dtype torch tensor (cpu) torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( @@ -2526,9 +2426,7 @@ def create_and_permute_tensor(l, mode0, mode1, is_mode0_major, dtype, is_dynamic # Create dtype cute tensor (gpu) torch_tensor_view = ( - torch_tensor - if dtype not in {cutlass.Float8E5M2, cutlass.Float8E4M3FN} - else torch_tensor.view(torch.uint8) + torch_tensor if dtype not in {cutlass.Float8E5M2, cutlass.Float8E4M3FN} else torch_tensor.view(torch.uint8) ) cute_tensor = from_dlpack(torch_tensor_view, assumed_align=16) cute_tensor.element_type = dtype @@ -2549,9 +2447,7 @@ def create_and_permute_tensor(l, mode0, mode1, is_mode0_major, dtype, is_dynamic b_ref, mB, b_torch, b_torch_cpu = create_and_permute_tensor( l, n, k, b_major == "n", ab_dtype, is_dynamic_layout=True ) - _, mD, d_torch, d_torch_cpu = create_and_permute_tensor( - l, m, n, d_major == "m", d_dtype, is_dynamic_layout=True - ) + _, mD, d_torch, d_torch_cpu = create_and_permute_tensor(l, m, n, d_major == "m", d_dtype, is_dynamic_layout=True) if c_dtype is not None: c, mC, c_torch, d_torch_cpu = create_and_permute_tensor(l, m, n, c_major == "m", c_dtype) else: @@ -2563,9 +2459,7 @@ def create_and_permute_tensor(l, mode0, mode1, is_mode0_major, dtype, is_dynamic # Compute max active clusters on current device hardware_info = cutlass.utils.HardwareInfo() - max_active_clusters = hardware_info.get_max_active_clusters( - cluster_shape_mn[0] * cluster_shape_mn[1] - ) + max_active_clusters = hardware_info.get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) if dynamic_persistent: tile_count_semaphore = torch.zeros(1, dtype=torch.int32, device="cuda") else: @@ -2573,9 +2467,7 @@ def create_and_permute_tensor(l, mode0, mode1, is_mode0_major, dtype, is_dynamic scheduler_args = TileSchedulerOptions( Int32(max_active_clusters), - tile_count_semaphore=make_ptr( - Int32, tile_count_semaphore.data_ptr(), cute.AddressSpace.gmem, assumed_align=4 - ) + tile_count_semaphore=make_ptr(Int32, tile_count_semaphore.data_ptr(), cute.AddressSpace.gmem, assumed_align=4) if tile_count_semaphore is not None else None, ) @@ -2688,9 +2580,7 @@ def create_and_permute_tensor(l, mode0, mode1, is_mode0_major, dtype, is_dynamic print(f"CuBLAS Average time: {timing_cublas:.3f} ms, TFLOPS: {tflops_cublas:.1f}") time.sleep(0.5) - fn = lambda: compiled_gemm( - mA, mB, mD, mC, epi_args, scheduler_args, varlen_args, current_stream - ) + fn = lambda: compiled_gemm(mA, mB, mD, mC, epi_args, scheduler_args, varlen_args, current_stream) timing = do_bench(fn, warmup=warmup, rep=repeats) tflops = flops / (timing * 1e9) # Convert to TFlops print(f"Cute-DSL Average time: {timing:.3f} ms, TFLOPS: {tflops:.1f}") @@ -2747,9 +2637,7 @@ def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: help="Number of iterations to run the kernel", ) parser.add_argument("--skip_ref_check", action="store_true", help="Skip reference checking") - parser.add_argument( - "--dynamic_persistent", action="store_true", help="Dynamic persistent kernel" - ) + parser.add_argument("--dynamic_persistent", action="store_true", help="Dynamic persistent kernel") args = parser.parse_args() diff --git a/src/xorl/ops/quack/gemm_sm90.py b/src/xorl/ops/quack/gemm_sm90.py index aa32fa6e..ac3d1337 100644 --- a/src/xorl/ops/quack/gemm_sm90.py +++ b/src/xorl/ops/quack/gemm_sm90.py @@ -2,38 +2,36 @@ # https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/hopper/dense_gemm.py import enum -from typing import Tuple, Type, Callable, Optional, Union, Literal -from functools import partial import math - +from functools import partial +from typing import Callable, Literal, Optional, Tuple, Type, Union import cuda.bindings.driver as cuda - import cutlass import cutlass.cute as cute import cutlass.pipeline as pipeline -from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait -from cutlass.cute.nvgpu import cpasync, warp, warpgroup import cutlass.utils.hopper_helpers as sm90_utils -from cutlass import Int32, Float32, Float16, Boolean, const_expr +from cutlass import Boolean, Float16, Float32, Int32, const_expr +from cutlass.cute.nvgpu import cpasync, warp, warpgroup +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.utils import LayoutEnum +from . import copy_utils +from . import sm90_utils as quack_sm90_utils +from .cute_dsl_utils import ArgumentsBase, ParamsBase -from .cute_dsl_utils import ParamsBase, ArgumentsBase +# return PipelineStateWAdvance instead of PipelineState +from .pipeline import PipelineTmaCpAsync, make_pipeline_state from .tile_scheduler import ( - TileSchedulerOptions, - TileSchedulerArguments, + PersistenceMode, TileScheduler, - VarlenMTileSchedulerArguments, + TileSchedulerArguments, + TileSchedulerOptions, VarlenMTileScheduler, - PersistenceMode, + VarlenMTileSchedulerArguments, ) from .varlen_utils import VarlenArguments, VarlenManager -# return PipelineStateWAdvance instead of PipelineState -from .pipeline import make_pipeline_state, PipelineTmaCpAsync -from . import copy_utils -from . import sm90_utils as quack_sm90_utils """ A high-performance batched dense GEMM (C = A * B) example for the NVIDIA Hopper architecture @@ -178,9 +176,7 @@ def __init__( f"If tile_m == {tile_M}, CTA tile shape N must be divisible by 32 and <= {tile_N_max}" ) else: - if not ( - (tile_N % 16 == 0 and tile_N <= 256) or (tile_N % 32 == 0 and tile_N <= 512) - ): + if not ((tile_N % 16 == 0 and tile_N <= 256) or (tile_N % 32 == 0 and tile_N <= 512)): raise ValueError( "CTA tile shape N must be divisible by 16 and <= 256, or divisible by 32 and <= 512" ) @@ -200,9 +196,7 @@ def __init__( else: atom_layout_m, atom_layout_n = 1, 2 else: - atom_layout_m = ( - self.cta_tile_shape_mnk[0] // 64 if self.cta_tile_shape_mnk[0] < 256 else 2 - ) + atom_layout_m = self.cta_tile_shape_mnk[0] // 64 if self.cta_tile_shape_mnk[0] < 256 else 2 atom_layout_n = 1 assert atom_layout_m in [1, 2, 3] and atom_layout_n in [1, 2] else: @@ -238,9 +232,7 @@ def __init__( self.num_regs_load, self.num_regs_mma = 32, 160 else: heavy_register_pressure = regs_per_thread >= 208 - self.num_regs_load, self.num_regs_mma = ( - (40, 232) if not heavy_register_pressure else (24, 240) - ) + self.num_regs_load, self.num_regs_mma = (40, 232) if not heavy_register_pressure else (24, 240) else: if self.mma_warp_groups == 3: self.num_regs_load, self.num_regs_mma = 56, 152 @@ -400,14 +392,11 @@ def __call__( # Assume all strides are divisible by 128 bits except the last stride def new_stride(t: cute.Tensor): return tuple( - cute.assume(s, divby=128 // t.element_type.width) if not cute.is_static(s) else s - for s in t.stride + cute.assume(s, divby=128 // t.element_type.width) if not cute.is_static(s) else s for s in t.stride ) mA, mD = [ - cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) - if t is not None - else None + cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) if t is not None else None for t in (mA, mD) ] @@ -456,9 +445,7 @@ def new_stride(t: cute.Tensor): TileSchedulerCls = self.get_scheduler_class(varlen_m=varlen_args.mCuSeqlensM is not None) tile_sched_args = self.get_scheduler_arguments(mA, mB, mD, scheduler_args, varlen_args) tile_sched_params = TileSchedulerCls.to_underlying_arguments(tile_sched_args) - grid = TileSchedulerCls.get_grid_shape( - tile_sched_params, scheduler_args.max_active_clusters - ) + grid = TileSchedulerCls.get_grid_shape(tile_sched_params, scheduler_args.max_active_clusters) epi_smem_size = cute.cosize(self.epi_smem_layout_staged) if mD is not None else 0 epi_c_smem_size = cute.cosize(self.epi_c_smem_layout_staged) if mC is not None else 0 @@ -470,15 +457,11 @@ class SharedStorage: sched_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.sched_stage * 2] sched_data: cute.struct.MemRange[Int32, self.sched_stage * 4] sD: cute.struct.Align[ - cute.struct.MemRange[ - self.d_dtype if self.d_dtype is not None else Int32, epi_smem_size - ], + cute.struct.MemRange[self.d_dtype if self.d_dtype is not None else Int32, epi_smem_size], self.buffer_align_bytes, ] sC: cute.struct.Align[ - cute.struct.MemRange[ - self.c_dtype if self.c_dtype is not None else Int32, epi_c_smem_size - ], + cute.struct.MemRange[self.c_dtype if self.c_dtype is not None else Int32, epi_c_smem_size], self.buffer_align_bytes, ] epi: self.epi_get_smem_struct(epilogue_params) @@ -633,28 +616,21 @@ def kernel( self.num_epi_tensormaps, # Only used if not varlen_m len_m_static=Int32( - mA_mkl.shape[0] - if varlen_k or varlen_params.mAIdx is None - else varlen_params.mAIdx.shape[0] + mA_mkl.shape[0] if varlen_k or varlen_params.mAIdx is None else varlen_params.mAIdx.shape[0] ), len_k_static=Int32(mA_mkl.shape[1]), pingpong=self.pingpong, warp_idx=warp_idx, ) - TileSchedulerCls = partial( - TileSchedulerCls.create, tile_sched_params, sched_data, sched_pipeline - ) + TileSchedulerCls = partial(TileSchedulerCls.create, tile_sched_params, sched_data, sched_pipeline) # Cluster wait for barrier init pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk[:-1]) if warp_idx >= self.ab_load_warp_id: cute.arch.setmaxregister_decrease(self.num_regs_load) - if ( - warp_idx >= self.ab_load_warp_id - and warp_idx < self.ab_load_warp_id + self.num_ab_load_warps - ): + if warp_idx >= self.ab_load_warp_id and warp_idx < self.ab_load_warp_id + self.num_ab_load_warps: is_tma_warp = self.num_ab_load_warps == 1 or warp_idx == self.ab_load_warp_id # initialize tensormap for A & B varlen_manager.init_tensormap_AB(tma_atom_a, tma_atom_b, is_tma_warp) @@ -663,12 +639,8 @@ def kernel( # Get mcast mask cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) block_in_cluster_coord_mnk = cluster_layout_mnk.get_flat_coord(cta_rank_in_cluster) - a_mcast_mask = cute.make_layout_image_mask( - cluster_layout_mnk, block_in_cluster_coord_mnk, mode=1 - ) - b_mcast_mask = cute.make_layout_image_mask( - cluster_layout_mnk, block_in_cluster_coord_mnk, mode=0 - ) + a_mcast_mask = cute.make_layout_image_mask(cluster_layout_mnk, block_in_cluster_coord_mnk, mode=1) + b_mcast_mask = cute.make_layout_image_mask(cluster_layout_mnk, block_in_cluster_coord_mnk, mode=0) a_mcast_mask = a_mcast_mask if self.is_a_mcast else 0 b_mcast_mask = b_mcast_mask if self.is_b_mcast else 0 @@ -678,18 +650,14 @@ def kernel( is_scheduler_warp = is_scheduler_warp and cute.arch.block_idx_in_cluster() == 0 tile_scheduler = TileSchedulerCls() work_tile = tile_scheduler.initial_work_tile_info() - ab_producer_state = make_pipeline_state( - pipeline.PipelineUserType.Producer, self.ab_stage - ) + ab_producer_state = make_pipeline_state(pipeline.PipelineUserType.Producer, self.ab_stage) if const_expr(varlen_k): # wait tensormap initialization complete before update varlen_manager.fence_tensormap_init() while work_tile.is_valid_tile: tile_coord_mnkl = work_tile.tile_idx batch_idx = tile_coord_mnkl[3] - varlen_manager.update_tensormap_AB( - batch_idx, self.a_layout, self.b_layout, is_tma_warp - ) + varlen_manager.update_tensormap_AB(batch_idx, self.a_layout, self.b_layout, is_tma_warp) # Local_tile partition global tensors if const_expr(not self.gather_A): mA_mk = varlen_manager.offset_batch_A(mA_mkl, batch_idx) @@ -702,9 +670,7 @@ def kernel( else: mAIdx_mk = varlen_manager.offset_batch_AIdx(batch_idx) if const_expr(varlen_m): - gAIdx = cute.local_tile( - mAIdx_mk, (self.cta_tile_shape_mnk[0],), (tile_coord_mnkl[0],) - ) + gAIdx = cute.local_tile(mAIdx_mk, (self.cta_tile_shape_mnk[0],), (tile_coord_mnkl[0],)) # (M, K) mA_mk = mA_mkl else: @@ -712,9 +678,7 @@ def kernel( # (tile_K, RestK) gAIdx = cute.flat_divide(mAIdx_mk, (self.cta_tile_shape_mnk[2],)) # (tile_M, K) - mA_mk = cute.local_tile( - mA_mkl, (self.cta_tile_shape_mnk[0],), (tile_coord_mnkl[0], None) - ) + mA_mk = cute.local_tile(mA_mkl, (self.cta_tile_shape_mnk[0],), (tile_coord_mnkl[0], None)) # (bN, bK, RestK) gB_nk = cute.local_tile( varlen_manager.offset_batch_B(mB_nkl, batch_idx), @@ -731,9 +695,7 @@ def kernel( copy_A, _, _ = copy_utils.tma_get_copy_fn( tma_atom_a, cta_coord=block_in_cluster_coord_mnk[1], - cta_layout=cute.make_layout( - cute.slice_(cluster_layout_mnk, (0, None, 0)).shape - ), + cta_layout=cute.make_layout(cute.slice_(cluster_layout_mnk, (0, None, 0)).shape), src_tensor=gA_mk, dst_tensor=sA, mcast_mask=a_mcast_mask, @@ -743,9 +705,7 @@ def kernel( tiled_copy_A = self._make_gmem_tiled_copy_A( mA_mkl.element_type, self.a_layout, self.num_ab_load_warps * 32 ) - tidx = ( - cute.arch.thread_idx()[0] - cute.arch.WARP_SIZE * self.ab_load_warp_id - ) + tidx = cute.arch.thread_idx()[0] - cute.arch.WARP_SIZE * self.ab_load_warp_id thr_copy_A = tiled_copy_A.get_slice(tidx) copy_A, prefetch_A = None, None if const_expr(varlen_m): @@ -770,9 +730,7 @@ def kernel( copy_B, _, _ = copy_utils.tma_get_copy_fn( tma_atom_b, cta_coord=block_in_cluster_coord_mnk[0], - cta_layout=cute.make_layout( - cute.slice_(cluster_layout_mnk, (None, 0, 0)).shape - ), + cta_layout=cute.make_layout(cute.slice_(cluster_layout_mnk, (None, 0, 0)).shape), src_tensor=gB_nk, dst_tensor=sB, mcast_mask=b_mcast_mask, @@ -780,9 +738,7 @@ def kernel( ) k_tile_cnt = cute.ceil_div(len_k, self.cta_tile_shape_mnk[2]) if const_expr(not self.gather_A): - ab_producer_state = self.load_AB( - ab_pipeline, ab_producer_state, copy_A, copy_B, k_tile_cnt - ) + ab_producer_state = self.load_AB(ab_pipeline, ab_producer_state, copy_A, copy_B, k_tile_cnt) else: ab_producer_state = self.load_AB_gather_A( ab_pipeline, @@ -808,12 +764,9 @@ def kernel( if warp_idx < self.ab_load_warp_id: cute.arch.setmaxregister_increase(self.num_regs_mma) is_tma_warp = Boolean( - (not self.pingpong and warp_idx == 0) - or (self.pingpong and (warp_idx == 0 or warp_idx == 4)) - ) - varlen_manager.init_tensormap_epi( - tma_atom_d, self.epi_get_tma_atoms(epilogue_params), is_tma_warp + (not self.pingpong and warp_idx == 0) or (self.pingpong and (warp_idx == 0 or warp_idx == 4)) ) + varlen_manager.init_tensormap_epi(tma_atom_d, self.epi_get_tma_atoms(epilogue_params), is_tma_warp) tma_desc_d_ptr = varlen_manager.get_tma_desc_d_ptr() tma_desc_epi_ptrs = varlen_manager.get_tma_desc_epi_ptrs() # Partition global tensor for TiledMMA_A/B/C @@ -825,14 +778,10 @@ def kernel( self.mma_warp_groups if const_expr(not self.pingpong) else 1, stride=self.num_threads_per_warp_group, ) - thr_mma = tiled_mma.get_slice( - warp_group_thread_layout(warp_group_idx if not self.pingpong else 0) - ) + thr_mma = tiled_mma.get_slice(warp_group_thread_layout(warp_group_idx if not self.pingpong else 0)) # Make fragments - acc, tCrA, tCrB = quack_sm90_utils.partition_fragment_ABC( - thr_mma, self.cta_tile_shape_mnk, sA, sB - ) + acc, tCrA, tCrB = quack_sm90_utils.partition_fragment_ABC(thr_mma, self.cta_tile_shape_mnk, sA, sB) acc_slow = None if const_expr(self.fp8_slow_accum): acc_slow = cute.make_rmem_tensor(acc.shape, self.acc_dtype) @@ -849,12 +798,8 @@ def kernel( ab_read_state = make_pipeline_state(pipeline.PipelineUserType.Consumer, self.ab_stage) epi_store_pipeline = self.make_epi_store_pipeline() - epi_read_state = make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.epi_c_stage - ) - epi_producer_state = make_pipeline_state( - pipeline.PipelineUserType.Producer, self.epi_c_stage - ) + epi_read_state = make_pipeline_state(pipeline.PipelineUserType.Consumer, self.epi_c_stage) + epi_producer_state = make_pipeline_state(pipeline.PipelineUserType.Producer, self.epi_c_stage) tile_scheduler = TileSchedulerCls() work_tile = tile_scheduler.initial_work_tile_info() if const_expr(self.pingpong): @@ -880,14 +825,10 @@ def kernel( epi_shapes, epi_orders = self.epi_get_tensormap_update_shapes_orders( epilogue_params, varlen_params.cu_seqlens_m, batch_idx ) - varlen_manager.update_tensormap_epi( - batch_idx, self.d_layout, epi_shapes, epi_orders, is_tma_warp - ) + varlen_manager.update_tensormap_epi(batch_idx, self.d_layout, epi_shapes, epi_orders, is_tma_warp) len_k = varlen_manager.len_k(batch_idx) k_tile_cnt = cute.ceil_div(len_k, self.cta_tile_shape_mnk[2]) - ab_read_state = self.mma( - ab_pipeline, ab_read_state, mma_fn, acc, acc_slow, k_tile_cnt, warp_group_idx - ) + ab_read_state = self.mma(ab_pipeline, ab_read_state, mma_fn, acc, acc_slow, k_tile_cnt, warp_group_idx) if const_expr(varlen_k): if k_tile_cnt == 0: acc.fill(0.0) @@ -1207,9 +1148,7 @@ def epilogue( ) -> Tuple[cutlass.pipeline.PipelineState, cutlass.pipeline.PipelineState]: has_C = const_expr(tRS_rC is not None) has_D = const_expr(copy_D is not None) - epi_tile_shape = cute.zipped_divide( - cute.make_layout(self.cta_tile_shape_mnk[:2]), epi_tile - ).shape[1] + epi_tile_shape = cute.zipped_divide(cute.make_layout(self.cta_tile_shape_mnk[:2]), epi_tile).shape[1] # We iterate over epi tiles in the N dimension first before the M dimension epi_tile_layout = cute.make_ordered_layout(epi_tile_shape, order=(1, 0)) epi_tile_num = cute.size(epi_tile_shape) @@ -1314,11 +1253,7 @@ def get_scheduler_arguments( num_problems = ( mD.shape[2] if mD is not None - else ( - mB.shape[2] - if varlen_args.mCuSeqlensK is None - else varlen_args.mCuSeqlensK.shape[0] - 1 - ) + else (mB.shape[2] if varlen_args.mCuSeqlensK is None else varlen_args.mCuSeqlensK.shape[0] - 1) ) problem_shape_ntile_mnl = ( cute.ceil_div(mA.shape[0], self.cta_tile_shape_mnk[0]), @@ -1411,14 +1346,10 @@ def epi_end( ) -> None: pass - def epi_to_underlying_arguments( - self, args: EpilogueArguments, *, loc=None, ip=None - ) -> EpilogueParams: + def epi_to_underlying_arguments(self, args: EpilogueArguments, *, loc=None, ip=None) -> EpilogueParams: return self.EpilogueParams() - def epi_get_tma_atoms( - self, params: EpilogueParams, *, loc=None, ip=None - ) -> list[cute.CopyAtom]: + def epi_get_tma_atoms(self, params: EpilogueParams, *, loc=None, ip=None) -> list[cute.CopyAtom]: """Subclasses can override this""" return [] @@ -1488,9 +1419,7 @@ def epilog_smem_store_and_partition( tiled_copy_C_atom = self.epilog_smem_copy_atom(tiled_mma) # Doesn't work with tile_N % 8 == 0 but tile_n % 16 != since this always # get st.matrix with num_matrices=4 - copy_atom_r2s = sm90_utils.sm90_get_smem_store_op( - d_layout, elem_ty_d=dtype, elem_ty_acc=self.acc_dtype - ) + copy_atom_r2s = sm90_utils.sm90_get_smem_store_op(d_layout, elem_ty_d=dtype, elem_ty_acc=self.acc_dtype) tiled_copy_r2s = cute.make_tiled_copy_S(copy_atom_r2s, tiled_copy_C_atom) # (R2S, R2S_M, R2S_N, PIPE_D) thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) @@ -1531,12 +1460,8 @@ def epilog_gmem_copy_and_partition( # (bM, bN) gD = cute.local_tile(mD_mn, tile_shape_mn, tile_coord_mnkl[:2]) tDgD_for_tma_partition = cute.zipped_divide(gD, epi_tile) - is_s2g = isinstance( - atom.op, (cpasync.CopyBulkTensorTileS2GOp, cpasync.CopyReduceBulkTensorTileS2GOp) - ) - src_tensor, dst_tensor = ( - (sD, tDgD_for_tma_partition) if is_s2g else (tDgD_for_tma_partition, sD) - ) + is_s2g = isinstance(atom.op, (cpasync.CopyBulkTensorTileS2GOp, cpasync.CopyReduceBulkTensorTileS2GOp)) + src_tensor, dst_tensor = (sD, tDgD_for_tma_partition) if is_s2g else (tDgD_for_tma_partition, sD) return copy_utils.tma_get_copy_fn( atom, cta_coord=0, @@ -1558,9 +1483,7 @@ def make_ab_pipeline( # Each warp will contribute to the arrive count with the number of mcast size mcast_size = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 consumer_arrive_cnt = mcast_size * tiled_mma.size // cute.arch.WARP_SIZE - ab_pipeline_consumer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, consumer_arrive_cnt - ) + ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_arrive_cnt) pipeline_cls = pipeline.PipelineTmaAsync if not self.gather_A else PipelineTmaCpAsync return pipeline_cls.create( barrier_storage=ab_pipeline_mbar_ptr, @@ -1572,16 +1495,12 @@ def make_ab_pipeline( defer_sync=True, ) - def make_epi_pipeline( - self, c_smem_layout: cute.Layout | cute.ComposedLayout, epi_pipeline_mbar_ptr: cute.Pointer - ): + def make_epi_pipeline(self, c_smem_layout: cute.Layout | cute.ComposedLayout, epi_pipeline_mbar_ptr: cute.Pointer): # Threads/warps participating in this pipeline epi_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) # Each warp will contribute 1 to the arrive count consumer_arrive_cnt = self.num_epi_warps - epi_pipeline_consumer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, consumer_arrive_cnt - ) + epi_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_arrive_cnt) tma_copy_c_bytes = cute.size_in_bytes(self.c_dtype, c_smem_layout) return pipeline.PipelineTmaAsync.create( barrier_storage=epi_pipeline_mbar_ptr, @@ -1596,9 +1515,7 @@ def make_epi_store_pipeline(self): # Threads/warps participating in tma store pipeline num_epi_threads = self.num_epi_warps * cute.arch.WARP_SIZE epi_store_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_epi_threads) - return pipeline.PipelineTmaStore.create( - num_stages=self.epi_stage, producer_group=epi_store_producer_group - ) + return pipeline.PipelineTmaStore.create(num_stages=self.epi_stage, producer_group=epi_store_producer_group) def make_sched_pipeline( self, cluster_layout_mnk: cute.Layout, sched_pipeline_mbar_ptr: cute.Pointer, varlen_k: bool @@ -1610,12 +1527,9 @@ def make_sched_pipeline( # If pingpong and varlen_k, then all 8 mma warps will participate in the scheduler barrier # at each round. If pingpong and not varlen_k, then only 4 mma warp will participate. consumer_arrive_cnt = ( - (self.mma_warp_groups if not (self.pingpong and not varlen_k) else 1) * 4 - + self.num_ab_load_warps + (self.mma_warp_groups if not (self.pingpong and not varlen_k) else 1) * 4 + self.num_ab_load_warps ) * cluster_size - sched_pipeline_consumer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, consumer_arrive_cnt - ) + sched_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_arrive_cnt) return pipeline.PipelineAsync.create( barrier_storage=sched_pipeline_mbar_ptr, num_stages=self.sched_stage, @@ -1669,9 +1583,7 @@ def _compute_stages( a_shape = cute.slice_(cta_tile_shape_mnk, (None, 0, None)) b_shape = cute.slice_(cta_tile_shape_mnk, (0, None, None)) - ab_bytes_per_stage = ( - cute.size(a_shape) * a_dtype.width // 8 + cute.size(b_shape) * b_dtype.width // 8 - ) + ab_bytes_per_stage = cute.size(a_shape) * a_dtype.width // 8 + cute.size(b_shape) * b_dtype.width // 8 mbar_helpers_bytes = 1024 remaining_bytes = smem_capacity // occupancy - mbar_helpers_bytes - epi_bytes @@ -1740,9 +1652,7 @@ def _make_smem_layouts( c_dtype: Optional[Type[cutlass.Numeric]], c_layout: Optional[LayoutEnum], epi_c_stage: int, - ) -> Tuple[ - cute.ComposedLayout, cute.ComposedLayout, cute.ComposedLayout, Optional[cute.ComposedLayout] - ]: + ) -> Tuple[cute.ComposedLayout, cute.ComposedLayout, cute.ComposedLayout, Optional[cute.ComposedLayout]]: """Create shared memory layouts for A, B, and C tensors. :param cta_tile_shape_mnk: CTA tile shape (M,N,K) @@ -1799,16 +1709,12 @@ def _make_smem_layouts( epi_smem_layout_staged = None if d_dtype is not None: - epi_smem_layout_staged = quack_sm90_utils.make_smem_layout_epi( - d_dtype, d_layout, epi_tile, epi_stage - ) + epi_smem_layout_staged = quack_sm90_utils.make_smem_layout_epi(d_dtype, d_layout, epi_tile, epi_stage) epi_c_smem_layout_staged = None if c_dtype is not None: assert c_layout is not None - epi_c_smem_layout_staged = quack_sm90_utils.make_smem_layout_epi( - c_dtype, c_layout, epi_tile, epi_c_stage - ) + epi_c_smem_layout_staged = quack_sm90_utils.make_smem_layout_epi(c_dtype, c_layout, epi_tile, epi_c_stage) return ( a_smem_layout_staged, @@ -1846,9 +1752,7 @@ def _make_tma_epi_atoms_and_tensors( if op_type == "store" else cpasync.CopyReduceBulkTensorTileS2GOp(cute.ReductionOp.ADD) ) - tma_atom_d, tma_tensor_d = cpasync.make_tiled_tma_atom( - op, tensor_d, epi_smem_layout, d_cta_v_layout - ) + tma_atom_d, tma_tensor_d = cpasync.make_tiled_tma_atom(op, tensor_d, epi_smem_layout, d_cta_v_layout) return tma_atom_d, tma_tensor_d @staticmethod @@ -1872,11 +1776,7 @@ def _make_tma_atoms_and_tensors( :return: TMA atom and tensor :rtype: Tuple[cute.CopyAtom, cute.Tensor] """ - op = ( - cpasync.CopyBulkTensorTileG2SOp() - if mcast_dim == 1 - else cpasync.CopyBulkTensorTileG2SMulticastOp() - ) + op = cpasync.CopyBulkTensorTileG2SOp() if mcast_dim == 1 else cpasync.CopyBulkTensorTileG2SMulticastOp() tma_atom, tma_tensor = cpasync.make_tiled_tma_atom( op, tensor, @@ -1898,16 +1798,12 @@ def _make_gmem_tiled_copy_A(self, dtype, major_mode, num_threads, copy_bits=128) if shape_dim_1 > loads_per_cache_line: shape_dim_1 = math.gcd(shape_dim_1, loads_per_cache_line) # thread layout for copy - thread_layout = cute.make_layout( - (num_threads // shape_dim_1, shape_dim_1), stride=(shape_dim_1, 1) - ) + thread_layout = cute.make_layout((num_threads // shape_dim_1, shape_dim_1), stride=(shape_dim_1, 1)) if major_mode != LayoutEnum.ROW_MAJOR: shape_dim_0 = cute.size(self.cta_tile_shape_mnk[0]) // copy_elems if shape_dim_0 > loads_per_cache_line: shape_dim_0 = math.gcd(shape_dim_0, loads_per_cache_line) - thread_layout = cute.make_layout( - (shape_dim_0, num_threads // shape_dim_0), stride=(1, shape_dim_0) - ) + thread_layout = cute.make_layout((shape_dim_0, num_threads // shape_dim_0), stride=(1, shape_dim_0)) # Value layout for copy value_layout = ( cute.make_layout((1, copy_elems)) diff --git a/src/xorl/ops/quack/gemm_symmetric.py b/src/xorl/ops/quack/gemm_symmetric.py index fcb5afd2..421dbdc1 100644 --- a/src/xorl/ops/quack/gemm_symmetric.py +++ b/src/xorl/ops/quack/gemm_symmetric.py @@ -1,21 +1,23 @@ -from typing import Tuple, Optional, Callable from functools import partial +from typing import Callable, Optional, Tuple + +import cutlass +import cutlass.cute as cute +import cutlass.torch as cutlass_torch +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.hopper_helpers as sm90_utils_og +from cutlass import Boolean, Float32, Int32, const_expr +from cutlass.cute.runtime import make_ptr from torch import Tensor + +from . import copy_utils +from .cute_dsl_utils import get_device_capacity, get_max_active_clusters from .gemm_act import GemmActMixin, act_fn_map, gemm_act from .gemm_sm90 import GemmSm90 from .gemm_sm100 import GemmSm100 -from .tile_scheduler import TriangularTileScheduler from .gemm_wrapper_utils import GemmWrapperBase -from .cute_dsl_utils import get_device_capacity, get_max_active_clusters +from .tile_scheduler import TriangularTileScheduler from .varlen_utils import VarlenManager -from . import copy_utils -import cutlass -import cutlass.cute as cute -import cutlass.torch as cutlass_torch -from cutlass.cute.runtime import make_ptr -from cutlass import Int32, Float32, Boolean, const_expr -import cutlass.utils.hopper_helpers as sm90_utils_og -import cutlass.utils.blackwell_helpers as sm100_utils class GemmSymmetricMixin(GemmActMixin, GemmSm90): @@ -62,9 +64,7 @@ def epilogue( if self.arch == 100 else sm90_utils_og.sm90_get_smem_store_op ) - copy_atom_postact_r2s = get_smem_store_op( - self.postact_layout, self.postact_dtype, self.acc_dtype - ) + copy_atom_postact_r2s = get_smem_store_op(self.postact_layout, self.postact_dtype, self.acc_dtype) # tiled_copy_C_atom = self.epilog_smem_copy_atom(tiled_mma) # tiled_copy_postact_r2s = cute.make_tiled_copy_S(copy_atom_postact_r2s, tiled_copy_C_atom) tiled_copy_postact_r2s = cute.make_tiled_copy_S(copy_atom_postact_r2s, tiled_copy_r2s) @@ -82,9 +82,7 @@ def epilogue( ) # We iterate over epi tiles in the N dimension first before the M dimension - epi_tile_shape = cute.zipped_divide( - cute.make_layout(self.cta_tile_shape_mnk[:2]), epi_tile - ).shape[1] + epi_tile_shape = cute.zipped_divide(cute.make_layout(self.cta_tile_shape_mnk[:2]), epi_tile).shape[1] epi_tile_layout = cute.make_layout(epi_tile_shape, stride=(epi_tile_shape[1], 1)) epi_tile_num = cute.size(epi_tile_shape) num_prev_subtiles = tile_scheduler.num_tiles_executed * epi_tile_num @@ -198,7 +196,7 @@ def gemm_symmetric( alpha: float | Tensor = 1.0, beta: float | Tensor = 1.0, ) -> None: - # Tranpose D so the "activation" is a write to the mirrored tile + # Transpose D so the "activation" is a write to the mirrored tile PostAct = D.mT L, M, K, N, tensor_infos = GemmWrapperBase.validate_and_prepare_tensors( diff --git a/src/xorl/ops/quack/gemm_wrapper_utils.py b/src/xorl/ops/quack/gemm_wrapper_utils.py index b3ad9411..6d62214f 100644 --- a/src/xorl/ops/quack/gemm_wrapper_utils.py +++ b/src/xorl/ops/quack/gemm_wrapper_utils.py @@ -1,17 +1,16 @@ # Copyright (c) 2025, Tri Dao. -from typing import Optional, Tuple, Dict, Any from dataclasses import dataclass - -import torch -from torch import Tensor +from typing import Any, Dict, Optional, Tuple import cutlass.cute as cute +import torch from cutlass import Int32 from cutlass.cute.runtime import from_dlpack, make_ptr +from torch import Tensor from .cute_dsl_utils import torch2cute_dtype_map -from .varlen_utils import VarlenArguments from .tile_scheduler import TileSchedulerOptions +from .varlen_utils import VarlenArguments @dataclass @@ -30,9 +29,7 @@ def validate_tensor(tensor: Tensor, name: str, ndim: int) -> None: @staticmethod def validate_shape(tensor: Tensor, expected_shape: Tuple[int, ...], name: str) -> None: - assert tensor.shape == expected_shape, ( - f"{name} must have shape {expected_shape}, got {tensor.shape}" - ) + assert tensor.shape == expected_shape, f"{name} must have shape {expected_shape}, got {tensor.shape}" @staticmethod def get_major_order(tensor: Tensor, dims: Tuple[str, str, str]) -> str: @@ -52,9 +49,7 @@ def create_cute_tensor( # Tensor is already permuted to (dims[0], dims[1], dims[2]) or (dim[0], dim[1]) # If major is dims[1], leading_dim is 1; if major is dims[0], leading_dim is 0 leading_dim = 1 if major == dims[1] else 0 - return from_dlpack(tensor.detach(), assumed_align=assumed_align).mark_layout_dynamic( - leading_dim=leading_dim - ) + return from_dlpack(tensor.detach(), assumed_align=assumed_align).mark_layout_dynamic(leading_dim=leading_dim) @staticmethod def validate_and_prepare_tensors( @@ -75,9 +70,7 @@ def validate_and_prepare_tensors( # Validate A_idx if provided (for gather_A case) gather_A = A_idx is not None if gather_A: - assert cu_seqlens_m is not None or cu_seqlens_k is not None, ( - "gather_A requires either varlen_m or varlen_k" - ) + assert cu_seqlens_m is not None or cu_seqlens_k is not None, "gather_A requires either varlen_m or varlen_k" assert A_idx.dtype == torch.int32, f"A_idx must be int32, got {A_idx.dtype}" assert A_idx.dim() == 1, f"A_idx must be 1D, got {A_idx.dim()}D" @@ -96,9 +89,7 @@ def validate_and_prepare_tensors( L, N, K_B = B.shape assert K == K_B, f"K dimension mismatch: A has {K}, B has {K_B}" - assert cu_seqlens_m.shape == (L + 1,), ( - f"cu_seqlens_m must have shape ({L + 1},), got {cu_seqlens_m.shape}" - ) + assert cu_seqlens_m.shape == (L + 1,), f"cu_seqlens_m must have shape ({L + 1},), got {cu_seqlens_m.shape}" M = total_M dc_shape = (total_M, N) dc_ndim = 2 @@ -117,9 +108,7 @@ def validate_and_prepare_tensors( N, K_B = B.shape assert total_K == K_B, f"K dimension mismatch: expected {total_K}, B has {K_B}" L = cu_seqlens_k.shape[0] - 1 - assert cu_seqlens_k.shape == (L + 1,), ( - f"cu_seqlens_k must have shape ({L + 1},), got {cu_seqlens_k.shape}" - ) + assert cu_seqlens_k.shape == (L + 1,), f"cu_seqlens_k must have shape ({L + 1},), got {cu_seqlens_k.shape}" K = total_K dc_shape = (L, M, N) dc_ndim = 3 @@ -137,12 +126,8 @@ def validate_and_prepare_tensors( # Validate D and C shapes uniformly for tensor, name in [(D, "D"), (C, "C")]: if tensor is not None: - assert tensor.dim() == dc_ndim, ( - f"{name} must be {dc_ndim}D for this mode, got {tensor.dim()}D" - ) - assert tensor.shape == dc_shape, ( - f"{name} shape {tensor.shape} doesn't match expected {dc_shape}" - ) + assert tensor.dim() == dc_ndim, f"{name} must be {dc_ndim}D for this mode, got {tensor.dim()}D" + assert tensor.shape == dc_shape, f"{name} shape {tensor.shape} doesn't match expected {dc_shape}" tensors = { "A": GemmTensorInfo(A), @@ -154,20 +139,14 @@ def validate_and_prepare_tensors( if additional_tensors: for name, tensor in additional_tensors.items(): if tensor is not None: - assert tensor.dim() == dc_ndim, ( - f"{name} must be {dc_ndim}D for this mode, got {tensor.dim()}D" - ) - assert tensor.shape == dc_shape, ( - f"{name} shape {tensor.shape} doesn't match expected {dc_shape}" - ) + assert tensor.dim() == dc_ndim, f"{name} must be {dc_ndim}D for this mode, got {tensor.dim()}D" + assert tensor.shape == dc_shape, f"{name} shape {tensor.shape} doesn't match expected {dc_shape}" tensors[name] = GemmTensorInfo(tensor) return L, M, K, N, tensors @staticmethod - def permute_tensors( - tensors: Dict[str, GemmTensorInfo], varlen_m: bool = False, varlen_k: bool = False - ) -> None: + def permute_tensors(tensors: Dict[str, GemmTensorInfo], varlen_m: bool = False, varlen_k: bool = False) -> None: # Determine which tensors need permutation if varlen_m: # Only B needs permutation (3D tensor) @@ -200,14 +179,10 @@ def determine_major_orders( tensors[name].major = GemmWrapperBase.get_major_order(tensors[name].tensor, dims) @staticmethod - def create_cute_tensors( - tensors: Dict[str, GemmTensorInfo], major_configs: Dict[str, Tuple[str, str, str]] - ) -> None: + def create_cute_tensors(tensors: Dict[str, GemmTensorInfo], major_configs: Dict[str, Tuple[str, str, str]]) -> None: for name, info in tensors.items(): if info.tensor is not None and name in major_configs: - info.cute_tensor = GemmWrapperBase.create_cute_tensor( - info.tensor, info.major, major_configs[name] - ) + info.cute_tensor = GemmWrapperBase.create_cute_tensor(info.tensor, info.major, major_configs[name]) @staticmethod def create_scheduler_args( @@ -223,9 +198,7 @@ def create_scheduler_args( ) if tile_count_semaphore is not None else None, - batch_idx_permute=( - from_dlpack(batch_idx_permute, assumed_align=4).mark_layout_dynamic(leading_dim=0) - ) + batch_idx_permute=(from_dlpack(batch_idx_permute, assumed_align=4).mark_layout_dynamic(leading_dim=0)) if batch_idx_permute is not None else None, max_swizzle_size=Int32(max_swizzle_size), @@ -285,9 +258,7 @@ def create_varlen_args( ), mTensormaps=tensormaps_cute, mAIdx=( - from_dlpack(A_idx, assumed_align=4).mark_layout_dynamic(leading_dim=0) - if A_idx is not None - else None + from_dlpack(A_idx, assumed_align=4).mark_layout_dynamic(leading_dim=0) if A_idx is not None else None ), ) diff --git a/src/xorl/ops/quack/layout_utils.py b/src/xorl/ops/quack/layout_utils.py index 099e0daf..9b833661 100644 --- a/src/xorl/ops/quack/layout_utils.py +++ b/src/xorl/ops/quack/layout_utils.py @@ -3,7 +3,6 @@ import cutlass import cutlass.cute as cute - from cutlass import Int32, const_expr @@ -198,9 +197,7 @@ def convert_layout_acc_frgA(acc_layout: cute.Layout) -> cute.Layout: # For Sm90, FP16/BF16, convert acc_layout from ((2, 2, N / 8), MMA_M, MMA_N) to ((2, 2, 2), MMA_M, (N / 16, MMA_N)) # TODO: Sm90 FP8 if const_expr(cute.rank(acc_layout.shape[0]) == 3): # Sm90 - l = cute.logical_divide( - acc_layout, ((None, None, 2), None, None) - ) # ((2, 2, (2, N / 16)), MMA_M, MMA_N) + l = cute.logical_divide(acc_layout, ((None, None, 2), None, None)) # ((2, 2, (2, N / 16)), MMA_M, MMA_N) rA_mma_view = cute.make_layout( ( (l.shape[0][0], l.shape[0][1], l.shape[0][2][0]), @@ -235,9 +232,7 @@ def reshape_acc_to_frgA(acc: cute.Tensor) -> cute.Tensor: return cute.make_tensor(acc.iterator, convert_layout_acc_frgA(acc.layout)) -def convert_layout_zero_stride( - input: cute.Tensor | cute.Layout, ref_layout: cute.Layout -) -> cute.Layout: +def convert_layout_zero_stride(input: cute.Tensor | cute.Layout, ref_layout: cute.Layout) -> cute.Layout: layout = input.layout if const_expr(isinstance(input, cute.Tensor)) else input # Group the modes with non-zero stride in the ref_layout together, # and the modes with zero stride together @@ -267,11 +262,7 @@ def mma_partition_C_vec( assert cute.rank(sVec) == 2 assert sVec.stride[0] == 1 stage = sVec.shape[1] - shape = ( - (sVec.shape[0], expand_shape, stage) - if const_expr(is_colvec) - else (expand_shape, sVec.shape[0], stage) - ) + shape = (sVec.shape[0], expand_shape, stage) if const_expr(is_colvec) else (expand_shape, sVec.shape[0], stage) stride = (1, 0, sVec.stride[1]) if const_expr(is_colvec) else (0, 1, sVec.stride[1]) sVec_mma = cute.make_tensor(sVec.iterator, cute.make_layout(shape, stride=stride)) tC_sVec = make_acc_tensor_mn_view(thr_mma.partition_C(sVec_mma)) @@ -284,11 +275,7 @@ def mma_partition_A_vec( assert cute.rank(sVec) == 2 assert sVec.stride[0] == 1 stage = sVec.shape[1] - shape = ( - (sVec.shape[0], expand_shape, stage) - if const_expr(is_colvec) - else (expand_shape, sVec.shape[0], stage) - ) + shape = (sVec.shape[0], expand_shape, stage) if const_expr(is_colvec) else (expand_shape, sVec.shape[0], stage) stride = (1, 0, sVec.stride[1]) if const_expr(is_colvec) else (0, 1, sVec.stride[1]) sVec_mma = cute.make_tensor(sVec.iterator, cute.make_layout(shape, stride=stride)) tC_sVec = make_acc_tensor_mn_view(thr_mma.partition_A(sVec_mma)) diff --git a/src/xorl/ops/quack/linear.py b/src/xorl/ops/quack/linear.py index 13d417c6..93fae57b 100644 --- a/src/xorl/ops/quack/linear.py +++ b/src/xorl/ops/quack/linear.py @@ -5,11 +5,9 @@ import torch.nn as nn import torch.nn.functional as F from torch import Tensor -from torch.amp import custom_fwd, custom_bwd +from torch.amp import custom_bwd, custom_fwd - -from .gemm_interface import gemm, gemm_add_inplace, gemm_act, gemm_dact -from .gemm_interface import gemm_gated, gemm_dgated +from .gemm_interface import gemm, gemm_act, gemm_add_inplace, gemm_dact, gemm_dgated, gemm_gated def linear_fwd_convert_type(*tensors): @@ -90,16 +88,10 @@ def backward(cls, ctx, dout, *args): x, weight, weight_og = ctx.saved_tensors # weight_og is None if not ctx.fuse_grad_accum batch_shape = dout.shape[:-1] dout = dout.reshape(-1, dout.shape[-1]) - dbias = ( - dout.sum(0, dtype=ctx.bias_dtype) - if ctx.bias_dtype is not None and ctx.needs_input_grad[2] - else None - ) + dbias = dout.sum(0, dtype=ctx.bias_dtype) if ctx.bias_dtype is not None and ctx.needs_input_grad[2] else None dx = linear_bwd_compute_input_grad(ctx, dout, weight, cls.matmul_bwd_dx) dx = dx.reshape(*batch_shape, dx.shape[-1]) if dx is not None else None - dweight = linear_bwd_compute_weight_grad( - ctx, dout, x, weight_og, cls.matmul_bwd_dw, cls.matmul_bwd_dw_inplace - ) + dweight = linear_bwd_compute_weight_grad(ctx, dout, x, weight_og, cls.matmul_bwd_dw, cls.matmul_bwd_dw_inplace) # return extra Nones for other classes that inherit from LinearFunc return dx, dweight, dbias, *([None] * 10) @@ -123,9 +115,7 @@ class LinearActFunc(LinearFunc): # Use classmethod instead of staticmethod to allow inheritance @classmethod @custom_fwd(device_type="cuda") - def forward( - cls, ctx, x, weight, activation, bias=None, store_preact=True, fuse_grad_accum=False - ): + def forward(cls, ctx, x, weight, activation, bias=None, store_preact=True, fuse_grad_accum=False): """ x: (..., in_features) weight: (out_features, in_features) @@ -139,9 +129,7 @@ def forward( x, weight = linear_fwd_convert_type(x, weight) batch_shape = x.shape[:-1] x = x.reshape(-1, x.shape[-1]) - out, postact = cls.matmul_fwd_fn( - x, weight.T, bias=bias, activation=activation, store_preact=store_preact - ) + out, postact = cls.matmul_fwd_fn(x, weight.T, bias=bias, activation=activation, store_preact=store_preact) linear_fwd_postprocess(ctx, x, weight, weight_og, needs_x_w_grad=ctx.needs_input_grad[:2]) if out is not None: out = out.reshape(*batch_shape, out.shape[-1]) @@ -159,9 +147,7 @@ class LinearActUntunedFunc(LinearActFunc): matmul_bwd_dw_inplace = partial(gemm_add_inplace, dynamic_scheduler=True) -def linear_act_func( - x, weight, activation, bias=None, store_preact=True, fuse_grad_accum=False, tuned=True -): +def linear_act_func(x, weight, activation, bias=None, store_preact=True, fuse_grad_accum=False, tuned=True): fn_cls = LinearActFunc if tuned else LinearActUntunedFunc return fn_cls.apply(x, weight, activation, bias, store_preact, fuse_grad_accum) @@ -187,9 +173,7 @@ def forward(cls, ctx, preact, weight, x, activation, fuse_grad_accum=False): x = x.reshape(-1, x.shape[-1]) out = cls.matmul_fwd_fn(x, weight.T) # Store preact instead of x, we will recompute x in the backward pass - linear_fwd_postprocess( - ctx, preact, weight, weight_og, needs_x_w_grad=ctx.needs_input_grad[:2] - ) + linear_fwd_postprocess(ctx, preact, weight, weight_og, needs_x_w_grad=ctx.needs_input_grad[:2]) ctx.activation = activation return out.reshape(*batch_shape, out.shape[-1]) @@ -210,9 +194,7 @@ def backward(cls, ctx, dout): else: dpreact, x = None, None dpreact = dpreact.reshape(*batch_shape, dpreact.shape[-1]) if dpreact is not None else None - dweight = linear_bwd_compute_weight_grad( - ctx, dout, x, weight_og, cls.matmul_bwd_dw, cls.matmul_bwd_dw_inplace - ) + dweight = linear_bwd_compute_weight_grad(ctx, dout, x, weight_og, cls.matmul_bwd_dw, cls.matmul_bwd_dw_inplace) return dpreact, dweight, *([None] * 3) @@ -241,9 +223,7 @@ class LinearGatedUntunedFunc(LinearActFunc): matmul_bwd_dw_inplace = partial(gemm_add_inplace, dynamic_scheduler=True, tuned=False) -def linear_gated_func( - x, weight, activation, bias=None, store_preact=True, fuse_grad_accum=False, tuned=True -): +def linear_gated_func(x, weight, activation, bias=None, store_preact=True, fuse_grad_accum=False, tuned=True): fn_cls = LinearGatedFunc if tuned else LinearGatedUntunedFunc return fn_cls.apply(x, weight, activation, bias, store_preact, fuse_grad_accum) diff --git a/src/xorl/ops/quack/linear_cross_entropy.py b/src/xorl/ops/quack/linear_cross_entropy.py index 98ad3c47..0d8cf84d 100644 --- a/src/xorl/ops/quack/linear_cross_entropy.py +++ b/src/xorl/ops/quack/linear_cross_entropy.py @@ -1,11 +1,11 @@ # Copyright (c) 2025, Tri Dao -from typing import Optional, Literal +from typing import Literal, Optional import torch import torch.nn as nn import torch.nn.functional as F from torch import Tensor -from torch.amp import custom_fwd, custom_bwd +from torch.amp import custom_bwd, custom_fwd from .cross_entropy import cross_entropy, cross_entropy_fwd_out from .gemm_interface import gemm, gemm_add, gemm_add_inplace @@ -22,9 +22,7 @@ def linear_cross_entropy_func( inplace_backward: bool = False, ) -> Tensor: y = F.linear(x, weight, bias) # (..., V) - return cross_entropy( - y, target, ignore_index=ignore_index, reduction=reduction, inplace_backward=inplace_backward - ) + return cross_entropy(y, target, ignore_index=ignore_index, reduction=reduction, inplace_backward=inplace_backward) def linear_cross_entropy_func_ref( @@ -173,9 +171,7 @@ def backward(ctx, dloss): # dw = dloss * dw + dloss * (last_dlogits_chunk.T @ last_x_chunk) # We use alpha=dloss, beta=dloss if ctx.weight_dtype == dw.dtype: - gemm_add_inplace( - last_dlogits_chunk.T, last_x_chunk, dw, alpha=dloss, beta=dloss, tuned=tuned - ) + gemm_add_inplace(last_dlogits_chunk.T, last_x_chunk, dw, alpha=dloss, beta=dloss, tuned=tuned) else: dw = gemm_add( last_dlogits_chunk.T, @@ -215,9 +211,7 @@ def chunked_linear_cross_entropy( """ if reduction not in ["mean", "sum"]: raise ValueError(f"Invalid reduction: {reduction}") - loss = ChunkedLinearCrossEntropyFunction.apply( - x, weight, target, ignore_index, reduction, chunk_size, tuned - ) + loss = ChunkedLinearCrossEntropyFunction.apply(x, weight, target, ignore_index, reduction, chunk_size, tuned) return loss diff --git a/src/xorl/ops/quack/mlp.py b/src/xorl/ops/quack/mlp.py index 4929bfde..32fcc3ae 100644 --- a/src/xorl/ops/quack/mlp.py +++ b/src/xorl/ops/quack/mlp.py @@ -4,7 +4,7 @@ import torch.nn.functional as F from torch import Tensor -from .linear import linear_act_func, act_linear_func +from .linear import act_linear_func, linear_act_func def mlp_func(x, weight1, weight2, activation: str, fuse_grad_accum=False, tuned=True): diff --git a/src/xorl/ops/quack/pipeline.py b/src/xorl/ops/quack/pipeline.py index 5fc82b1d..276e2f20 100644 --- a/src/xorl/ops/quack/pipeline.py +++ b/src/xorl/ops/quack/pipeline.py @@ -1,15 +1,22 @@ # Copyright (c) 2025, Tri Dao. -from typing import Optional from dataclasses import dataclass +from typing import Optional import cutlass.cute as cute from cutlass import Boolean, Int32, const_expr -from cutlass.cutlass_dsl import if_generate, and_, dsl_user_op -from cutlass.pipeline import MbarrierArray, CooperativeGroup, PipelineOp -from cutlass.pipeline import PipelineTmaAsync, PipelineState, PipelineUserType -from cutlass.pipeline import PipelineTmaUmma -from cutlass.pipeline import Agent, agent_sync +from cutlass.cutlass_dsl import and_, dsl_user_op, if_generate +from cutlass.pipeline import ( + Agent, + CooperativeGroup, + MbarrierArray, + PipelineOp, + PipelineState, + PipelineTmaAsync, + PipelineTmaUmma, + PipelineUserType, + agent_sync, +) class PipelineStateWAdvance(PipelineState): @@ -24,9 +31,7 @@ def advance_iters(self, num_iterations: Int32, *, loc=None, ip=None): # This can be overridden by derived classes def __new_from_mlir_values__(self, values): - return PipelineStateWAdvance( - self.stages, Int32(values[0]), Int32(values[1]), Int32(values[2]) - ) + return PipelineStateWAdvance(self.stages, Int32(values[0]), Int32(values[1]), Int32(values[2])) def make_pipeline_state(type: PipelineUserType, stages: int): @@ -98,9 +103,7 @@ def producer_cpasync_commit(self, state: PipelineState, *, loc=None, ip=None): """ We need the mbarrier to track the completion of cp.async """ - cute.arch.cp_async_mbarrier_arrive_noinc( - self.producer_get_barrier(state, loc=loc, ip=ip), loc=loc, ip=ip - ) + cute.arch.cp_async_mbarrier_arrive_noinc(self.producer_get_barrier(state, loc=loc, ip=ip), loc=loc, ip=ip) class MbarrierArrayWDropCount(MbarrierArray): @@ -143,9 +146,7 @@ def __extract_mlir_values__(self): return [self.barrier_storage, self.drop_count] def __new_from_mlir_values__(self, values): - return MbarrierArrayWDropCount( - values[0], self.num_stages, (self.op_type, self.cg), self.tx_count, values[1] - ) + return MbarrierArrayWDropCount(values[0], self.num_stages, (self.op_type, self.cg), self.tx_count, values[1]) @dataclass(frozen=True) @@ -192,9 +193,7 @@ def create( :rtype: PipelineTmaUmma """ if not isinstance(barrier_storage, cute.Pointer): - raise TypeError( - f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" - ) + raise TypeError(f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}") producer_type = PipelineOp.TmaLoad consumer_type = PipelineOp.TCGen05Mma @@ -225,9 +224,7 @@ def create( # All threadblocks are leaders if not using clusters is_leader_cta = True else: - producer_mask = PipelineTmaUmma._compute_mcast_arrival_mask( - cta_layout_vmnk, mcast_mode_mn, loc=loc, ip=ip - ) + producer_mask = PipelineTmaUmma._compute_mcast_arrival_mask(cta_layout_vmnk, mcast_mode_mn, loc=loc, ip=ip) is_leader_cta = PipelineTmaUmma._compute_is_leader_cta(cta_layout_vmnk, loc=loc, ip=ip) cta_group = ( @@ -289,6 +286,4 @@ def producer_cpasync_commit(self, state: PipelineState, *, loc=None, ip=None): """ We need the mbarrier to track the completion of cp.async """ - cute.arch.cp_async_mbarrier_arrive_noinc( - self.producer_get_barrier(state, loc=loc, ip=ip), loc=loc, ip=ip - ) + cute.arch.cp_async_mbarrier_arrive_noinc(self.producer_get_barrier(state, loc=loc, ip=ip), loc=loc, ip=ip) diff --git a/src/xorl/ops/quack/reduce.py b/src/xorl/ops/quack/reduce.py index 11f7f574..bee0c455 100644 --- a/src/xorl/ops/quack/reduce.py +++ b/src/xorl/ops/quack/reduce.py @@ -6,7 +6,7 @@ import cutlass import cutlass.cute as cute -from cutlass import Int32, Int64, Float32, Boolean, const_expr +from cutlass import Boolean, Float32, Int32, Int64, const_expr from . import utils @@ -113,13 +113,9 @@ def row_reduce( hook_fn() if const_expr(reduction_buffer is not None): warps_per_row, cluster_n = reduction_buffer.shape[1] - assert cluster_n == 1 or mbar_ptr is not None, ( - "mbar_ptr must be provided for cluster reduction" - ) + assert cluster_n == 1 or mbar_ptr is not None, "mbar_ptr must be provided for cluster reduction" if const_expr(warps_per_row > 1 or cluster_n > 1): - val = block_or_cluster_reduce( - val, warp_op, reduction_buffer, mbar_ptr, phase=phase, init_val=init_val - ) + val = block_or_cluster_reduce(val, warp_op, reduction_buffer, mbar_ptr, phase=phase, init_val=init_val) return val @@ -151,13 +147,9 @@ def online_softmax_reduce( hook_fn() if const_expr(reduction_buffer is not None): rows_per_block, (warps_per_row, cluster_n) = reduction_buffer.shape - assert cluster_n == 1 or mbar_ptr is not None, ( - "mbar_ptr must be provided for cluster reduction" - ) + assert cluster_n == 1 or mbar_ptr is not None, "mbar_ptr must be provided for cluster reduction" if const_expr(warps_per_row > 1 or cluster_n > 1): - assert reduction_buffer.element_type == Int64, ( - "reduction_buffer must be of type cute.Int64" - ) + assert reduction_buffer.element_type == Int64, "reduction_buffer must be of type cute.Int64" lane_idx, warp_idx = cute.arch.lane_idx(), cute.arch.warp_idx() row_idx, col_idx = warp_idx // warps_per_row, warp_idx % warps_per_row if const_expr(mbar_ptr is None): @@ -167,9 +159,7 @@ def online_softmax_reduce( max_x_single_warp = -Float32.inf sum_exp_x = 0.0 if lane_idx < warps_per_row: - max_x_single_warp, sum_exp_x = utils.i64_to_f32x2( - reduction_buffer[row_idx, lane_idx] - ) + max_x_single_warp, sum_exp_x = utils.i64_to_f32x2(reduction_buffer[row_idx, lane_idx]) max_x_final = cute.arch.warp_reduction(max_x_single_warp, cute.arch.fmax) sum_exp_x *= cute.math.exp(max_x_single_warp - max_x_final, fastmath=True) sum_exp_x = cute.arch.warp_reduction(sum_exp_x, operator.add) @@ -188,9 +178,7 @@ def online_softmax_reduce( if lane_idx < cluster_n: utils.store_shared_remote( utils.f32x2_to_i64(max_x, sum_exp_x), - utils.elem_pointer( - reduction_buffer, (row_idx, (col_idx, cta_rank_in_cluster)) - ), + utils.elem_pointer(reduction_buffer, (row_idx, (col_idx, cta_rank_in_cluster))), mbar_ptr, peer_cta_rank_in_cluster=lane_idx, ) @@ -223,9 +211,7 @@ def online_softmax_reduce( @cute.jit -def sum_swap_shuffle( - X: cute.Tensor, elem_per_lane: int = 1, subwarp_size: int = 1, warp_size: int = 32 -) -> cute.Tensor: +def sum_swap_shuffle(X: cute.Tensor, elem_per_lane: int = 1, subwarp_size: int = 1, warp_size: int = 32) -> cute.Tensor: """ For warp reduction, we use Swap Shuffle The normal way to reduction among threads: @@ -237,25 +223,15 @@ def sum_swap_shuffle( After reduction, each half of threads should deal with a (N/2)x(N/2) sub-matrix independently in the following step. We can recursively do this until the problem size is 1. """ - assert ( - subwarp_size >= 1 - and subwarp_size <= 32 - and subwarp_size == 1 << int(math.log2(subwarp_size)) - ) - assert ( - warp_size <= 32 - and warp_size % subwarp_size == 0 - and warp_size == 1 << int(math.log2(warp_size)) - ) + assert subwarp_size >= 1 and subwarp_size <= 32 and subwarp_size == 1 << int(math.log2(subwarp_size)) + assert warp_size <= 32 and warp_size % subwarp_size == 0 and warp_size == 1 << int(math.log2(warp_size)) lane_idx = cute.arch.lane_idx() // subwarp_size X = cute.logical_divide(X, cute.make_layout(elem_per_lane)) # (elem_per_lane, M) numvec = cute.size(X, mode=[1]) assert numvec <= 32 // subwarp_size # If X has more values than warp_size // subwarp_size, we first do a normal warp reduction # to sum up values held by lanes further than size(X) away - for i in cutlass.range( - int(math.log2(numvec)), int(math.log2(warp_size // subwarp_size)), unroll_full=True - ): + for i in cutlass.range(int(math.log2(numvec)), int(math.log2(warp_size // subwarp_size)), unroll_full=True): for v in cutlass.range(cute.size(X), unroll_full=True): shfl_val = cute.arch.shuffle_sync_bfly(X[v], offset=(1 << i) * subwarp_size) X[v] = X[v] + shfl_val diff --git a/src/xorl/ops/quack/reduction_base.py b/src/xorl/ops/quack/reduction_base.py index bd0418e8..1b6c21ae 100644 --- a/src/xorl/ops/quack/reduction_base.py +++ b/src/xorl/ops/quack/reduction_base.py @@ -1,10 +1,10 @@ # Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao. -from typing import Type, Tuple, Optional +from typing import Optional, Tuple, Type import cutlass import cutlass.cute as cute -from cutlass import Int32, Int64, Float32, const_expr +from cutlass import Float32, Int32, Int64, const_expr from . import copy_utils @@ -38,9 +38,7 @@ def _get_tiled_copy(self, vecsize: int = 1): def _get_reduction_buffer_layout(self, tv_layout: cute.Layout, cluster_n: int): num_warps = cute.size(tv_layout, mode=[0]) // cute.arch.WARP_SIZE warps_per_row = ( - num_warps - if cute.rank(tv_layout.shape[0]) == 1 - else max(tv_layout.shape[0][0] // cute.arch.WARP_SIZE, 1) + num_warps if cute.rank(tv_layout.shape[0]) == 1 else max(tv_layout.shape[0][0] // cute.arch.WARP_SIZE, 1) ) return cute.make_ordered_layout( (num_warps // warps_per_row, (warps_per_row, cluster_n), self.stage), @@ -56,9 +54,7 @@ def _allocate_reduction_buffer_and_mbar( byte_alignment=8, ) if const_expr(self.cluster_n > 1): - mbar_ptr = smem.allocate_array( - Int64, num_elems=self.stage if not is_persistent else self.stage * 2 - ) + mbar_ptr = smem.allocate_array(Int64, num_elems=self.stage if not is_persistent else self.stage * 2) else: mbar_ptr = None return reduction_buffer, mbar_ptr @@ -75,9 +71,7 @@ def _initialize_cluster( if tidx < self.stage: # Initialize full barrier cute.arch.mbarrier_init(mbar_ptr + tidx, 1) if const_expr(is_persistent): # Initialize empty barrier - cute.arch.mbarrier_init( - mbar_ptr + self.stage + tidx, num_warps * self.cluster_n - ) + cute.arch.mbarrier_init(mbar_ptr + self.stage + tidx, num_warps * self.cluster_n) cute.arch.mbarrier_init_fence() # Cluster arrive after barrier init cute.arch.cluster_arrive_relaxed() diff --git a/src/xorl/ops/quack/rmsnorm.py b/src/xorl/ops/quack/rmsnorm.py index 5b029e13..43f239f9 100644 --- a/src/xorl/ops/quack/rmsnorm.py +++ b/src/xorl/ops/quack/rmsnorm.py @@ -1,25 +1,21 @@ # Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao. import math -from typing import Optional, Tuple, Type from functools import partial +from typing import Optional, Tuple, Type import cuda.bindings.driver as cuda - import cutlass import cutlass.cute as cute -from cutlass import Float32, Int32, const_expr - import torch +from cutlass import Float32, Int32, const_expr from torch import Tensor -from . import utils -from . import copy_utils -from . import layout_utils +from . import copy_utils, layout_utils, utils from .compile_utils import make_fake_tensor as fake_tensor +from .cute_dsl_utils import torch2cute_dtype_map from .reduce import row_reduce from .reduction_base import ReductionBase -from .cute_dsl_utils import torch2cute_dtype_map class RMSNorm(ReductionBase): @@ -73,16 +69,12 @@ def __call__( tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=vecsize) num_threads = tiled_copy.size mW, mB = [ - layout_utils.expand(mT, dim=0, size=tiler_mn[0]) if const_expr(mT is not None) else None - for mT in (mW, mB) + layout_utils.expand(mT, dim=0, size=tiler_mn[0]) if const_expr(mT is not None) else None for mT in (mW, mB) ] mRstd, mMean = [ - layout_utils.expand(mT, dim=1, size=self.N) if const_expr(mT is not None) else None - for mT in (mRstd, mMean) + layout_utils.expand(mT, dim=1, size=self.N) if const_expr(mT is not None) else None for mT in (mRstd, mMean) ] - self.kernel( - mX, mW, mB, mRes, mO, mResO, mRstd, mMean, eps, tiler_mn, tiled_copy, threads_per_row - ).launch( + self.kernel(mX, mW, mB, mRes, mO, mResO, mRstd, mMean, eps, tiler_mn, tiled_copy, threads_per_row).launch( grid=[cute.ceil_div(mX.shape[0], tiler_mn[0]), self.cluster_n, 1], block=[num_threads, 1, 1], cluster=[1, self.cluster_n, 1] if const_expr(self.cluster_n > 1) else None, @@ -111,9 +103,7 @@ def kernel( tv_layout = tiled_copy.layout_tv_tiled smem = cutlass.utils.SmemAllocator() - sX = smem.allocate_tensor( - mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16 - ) + sX = smem.allocate_tensor(mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16) if const_expr(mRes is not None): sRes = smem.allocate_tensor( mRes.element_type, @@ -130,8 +120,7 @@ def kernel( for mT in (mX, mRes, mO, mResO, mRstd, mMean, idX) ] gW, gB = [ - cute.local_tile(mT, tiler_mn, (0, cluster_y)) if const_expr(mT is not None) else None - for mT in (mW, mB) + cute.local_tile(mT, tiler_mn, (0, cluster_y)) if const_expr(mT is not None) else None for mT in (mW, mB) ] thr_copy_X = tiled_copy.get_slice(tidx) @@ -161,11 +150,7 @@ def kernel( self._initialize_cluster(tidx, mbar_ptr, num_warps) is_even_N = const_expr(shape[1] == tiler_mn[1] * self.cluster_n) - tXpX = ( - copy_utils.predicate_k(thr_copy_X.partition_S(cX), limit=shape[1]) - if not is_even_N - else None - ) + tXpX = copy_utils.predicate_k(thr_copy_X.partition_S(cX), limit=shape[1]) if not is_even_N else None # Each copy will use the same predicate copy = partial(copy_utils.copy, pred=tXpX) @@ -251,11 +236,7 @@ def kernel( rstd = cute.math.rsqrt(sum_sq_x / shape[1] + eps, fastmath=True) if const_expr(mRstd is not None): # Only the thread corresponding to column 0 writes out the rstd to gmem - if ( - tXcX[0][1] == 0 - and row < shape[0] - and (self.cluster_n == 1 or cute.arch.block_idx_in_cluster() == 0) - ): + if tXcX[0][1] == 0 and row < shape[0] and (self.cluster_n == 1 or cute.arch.block_idx_in_cluster() == 0): tXrRstd[0] = rstd if const_expr(self.delay_w_load): if const_expr(mW is not None): @@ -323,8 +304,7 @@ def _rmsnorm_fwd( _, N = x.shape dtype, out_dtype, weight_dtype, bias_dtype, res_dtype, res_out_dtype = [ - torch2cute_dtype_map[t.dtype] if t is not None else None - for t in [x, out, weight, bias, residual, residual_out] + torch2cute_dtype_map[t.dtype] if t is not None else None for t in [x, out, weight, bias, residual, residual_out] ] compile_key = ( dtype, @@ -343,8 +323,7 @@ def _rmsnorm_fwd( all_dtypes = [dtype, out_dtype, res_dtype, weight_dtype, bias_dtype, res_out_dtype] div = math.gcd(N, *(128 // dt.width for dt in all_dtypes if dt is not None)) x_cute, out_cute, res_cute, res_out_cute = [ - fake_tensor(dt, (batch_sym, N), div) - for dt in [dtype, out_dtype, res_dtype, res_out_dtype] + fake_tensor(dt, (batch_sym, N), div) for dt in [dtype, out_dtype, res_dtype, res_out_dtype] ] weight_cute, bias_cute = [fake_tensor(dt, (N,), div) for dt in [weight_dtype, bias_dtype]] rstd_cute = fake_tensor(Float32, (batch_sym,)) if rstd is not None else None @@ -363,9 +342,7 @@ def _rmsnorm_fwd( cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), options="--enable-tvm-ffi", ) - _rmsnorm_fwd.compile_cache[compile_key]( - x, weight, bias, residual, out, residual_out, rstd, mean, eps - ) + _rmsnorm_fwd.compile_cache[compile_key](x, weight, bias, residual, out, residual_out, rstd, mean, eps) _rmsnorm_fwd.compile_cache = {} @@ -390,9 +367,7 @@ def rmsnorm_fwd( if residual is not None: residual_dtype = residual.dtype if residual is not None or (residual_dtype is not None and residual_dtype != x.dtype): - residual_out = torch.empty_like( - x, dtype=residual_dtype if residual_dtype is not None else x.dtype - ) + residual_out = torch.empty_like(x, dtype=residual_dtype if residual_dtype is not None else x.dtype) else: residual_out = None _rmsnorm_fwd(x, weight, out, bias, rstd, None, residual, residual_out, eps, False) @@ -486,13 +461,9 @@ def __call__( vecsize = math.gcd(self.N, 128 // largest_dtype_width) tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=vecsize) num_threads = tiled_copy.size - mW = ( - layout_utils.expand(mW, dim=0, size=tiler_mn[0]) if const_expr(mW is not None) else None - ) + mW = layout_utils.expand(mW, dim=0, size=tiler_mn[0]) if const_expr(mW is not None) else None num_blocks = sm_count - self.kernel( - mX, mW, mdO, mdResO, mRstd, mdX, mdW, mdB, mdRes, tiler_mn, tiled_copy, threads_per_row - ).launch( + self.kernel(mX, mW, mdO, mdResO, mRstd, mdX, mdW, mdB, mdRes, tiler_mn, tiled_copy, threads_per_row).launch( grid=[num_blocks, self.cluster_n, 1], block=[num_threads, 1, 1], cluster=[1, self.cluster_n, 1] if self.cluster_n > 1 else None, @@ -531,9 +502,7 @@ def kernel( smem_layout = cute.make_ordered_layout((tiler_mn[0], tiler_mn[1], 2), order=(1, 0, 2)) sX = smem.allocate_tensor(mX.element_type, smem_layout, byte_alignment=16) sdO = smem.allocate_tensor(mdO.element_type, smem_layout, byte_alignment=16) - reduction_buffer, mbar_ptr = self._allocate_reduction_buffer_and_mbar( - smem, tv_layout, is_persistent=True - ) + reduction_buffer, mbar_ptr = self._allocate_reduction_buffer_and_mbar(smem, tv_layout, is_persistent=True) if const_expr(mbar_ptr is not None): mbar_full_ptr, mbar_empty_ptr = mbar_ptr, mbar_ptr + 2 else: @@ -547,9 +516,7 @@ def kernel( ] gW = cute.local_tile(mW, tiler_mn, (0, cluster_y)) if mW is not None else None gdW, gdB = [ - cute.local_tile(mT, (1, tiler_mn[1]), (bidx_start, cluster_y)) - if const_expr(mT is not None) - else None + cute.local_tile(mT, (1, tiler_mn[1]), (bidx_start, cluster_y)) if const_expr(mT is not None) else None for mT in (mdW, mdB) ] @@ -564,9 +531,7 @@ def kernel( tXgdRes = thr_copy_X.partition_D(gdRes) tXcX = thr_copy_X.partition_S(cX)[(0, None), None, None, None] - tXrX, tXrdO, tXrdX = [ - cute.make_fragment_like(thr[None, None, None, 0]) for thr in (tXgX, tXgdO, tXgdX) - ] + tXrX, tXrdO, tXrdX = [cute.make_fragment_like(thr[None, None, None, 0]) for thr in (tXgX, tXgdO, tXgdX)] tXrdResO = None if const_expr(mdResO is not None): tXrdResO = cute.make_fragment_like(tXgdResO[None, None, None, 0]) @@ -575,11 +540,7 @@ def kernel( tXrdRes = cute.make_fragment_like(tXgdRes[None, None, None, 0]) # This doesn't change across iterations - tXpX = ( - None - if is_even_N - else copy_utils.predicate_k(thr_copy_X.partition_S(cX[None, None, 0]), limit=shape[1]) - ) + tXpX = None if is_even_N else copy_utils.predicate_k(thr_copy_X.partition_S(cX[None, None, 0]), limit=shape[1]) # Each copy will use the same number of elements as X copy = partial(copy_utils.copy, pred=tXpX) @@ -645,12 +606,8 @@ def kernel( ) else: if const_expr(tiler_mn[0] > 1): - utils.fill_oob( - tXsX[None, None, None, stage ^ 1], None, fill_value=mX.element_type.zero - ) - utils.fill_oob( - tXsdO[None, None, None, stage ^ 1], None, fill_value=mdO.element_type.zero - ) + utils.fill_oob(tXsX[None, None, None, stage ^ 1], None, fill_value=mX.element_type.zero) + utils.fill_oob(tXsdO[None, None, None, stage ^ 1], None, fill_value=mdO.element_type.zero) cute.arch.cp_async_commit_group() rstd = cutlass.Float.zero if row < M or tiler_mn[0] == 1: @@ -692,9 +649,7 @@ def kernel( cute.arch.sync_warp() lane_idx = cute.arch.lane_idx() if lane_idx < self.cluster_n: - cute.arch.mbarrier_arrive( - mbar_empty_ptr + stage, peer_cta_rank_in_cluster=lane_idx - ) + cute.arch.mbarrier_arrive(mbar_empty_ptr + stage, peer_cta_rank_in_cluster=lane_idx) if const_expr(self.reload_wdy == "smem"): cute.autovec_copy(tXsdO[None, None, None, stage], tXrdO) @@ -739,9 +694,7 @@ def kernel( if row == 0: for i in cutlass.range_constexpr(1, const_expr(tiler_mn[0])): tXrdW_other = cute.make_fragment_like(tXrdW) - tXsdW_other = cute.make_tensor( - tXsdW.iterator + i * sdW.stride[0], tXsdW.layout - ) + tXsdW_other = cute.make_tensor(tXsdW.iterator + i * sdW.stride[0], tXsdW.layout) cute.autovec_copy(tXsdW_other, tXrdW_other) tXrdW.store(tXrdW.load() + tXrdW_other.load()) copy(tXrdW, tXgdW) @@ -760,9 +713,7 @@ def kernel( if row == 0: for i in cutlass.range_constexpr(1, const_expr(tiler_mn[0])): tXrdB_other = cute.make_fragment_like(tXrdB) - tXsdB_other = cute.make_tensor( - tXsdB.iterator + i * sdB.stride[0], tXsdB.layout - ) + tXsdB_other = cute.make_tensor(tXsdB.iterator + i * sdB.stride[0], tXsdB.layout) cute.autovec_copy(tXsdB_other, tXrdB_other) tXrdB.store(tXrdB.load() + tXrdB_other.load()) copy(tXrdB, tXgdB) @@ -784,17 +735,13 @@ def kernel( def _get_sm_count(N: int, device: torch.device) -> int: # This should be tuned on how many CTAs can be launched on each SM - sm_count_multiple = ( - 16 if N <= 256 else (8 if N <= 1024 else (4 if N <= 2048 else (2 if N <= 4096 else 1))) - ) + sm_count_multiple = 16 if N <= 256 else (8 if N <= 1024 else (4 if N <= 2048 else (2 if N <= 4096 else 1))) sm_count = torch.cuda.get_device_properties(device).multi_processor_count # By right, if we're using cluster, this should be cluster_count not sm_count. # But for cluster >= 4, due to quantization we would need to query active max cluster. # Instead we just do sm_count * 2, which is reasonably larger than active_cluster_count to # avoid wave quantization. - sm_count = ( - sm_count * sm_count_multiple if N <= 8192 else sm_count // 2 if N <= 16384 else sm_count * 2 - ) + sm_count = sm_count * sm_count_multiple if N <= 8192 else sm_count // 2 if N <= 16384 else sm_count * 2 return sm_count @@ -841,9 +788,7 @@ def _rmsnorm_bwd( if dresidual_out is not None: assert dresidual_out.shape == x.shape assert dresidual_out.is_cuda - assert dresidual_out.dtype in supported_types, ( - "Residual must be float16, bfloat16, or float32" - ) + assert dresidual_out.dtype in supported_types, "Residual must be float16, bfloat16, or float32" if dresidual is not None: assert dresidual.shape == x.shape assert dresidual.is_cuda @@ -873,17 +818,12 @@ def _rmsnorm_bwd( all_dtypes = [dtype, dout_dtype, dx_dtype, dres_dtype, dres_out_dtype] div = math.gcd(N, *(128 // dt.width for dt in all_dtypes if dt is not None)) x_cute, dout_cute, dx_cute, dres_out_cute, dres_cute = [ - fake_tensor(dt, (batch_sym, N), div) - for dt in [dtype, dout_dtype, dx_dtype, dres_out_dtype, dres_dtype] + fake_tensor(dt, (batch_sym, N), div) for dt in [dtype, dout_dtype, dx_dtype, dres_out_dtype, dres_dtype] ] weight_cute = fake_tensor(weight_dtype, (N,), div) rstd_cute = fake_tensor(Float32, (batch_sym,)) - dw_partial_cute = ( - fake_tensor(Float32, (batch_partial_sym, N), div) if dw_partial is not None else None - ) - db_partial_cute = ( - fake_tensor(Float32, (batch_partial_sym, N), div) if db_partial is not None else None - ) + dw_partial_cute = fake_tensor(Float32, (batch_partial_sym, N), div) if dw_partial is not None else None + db_partial_cute = fake_tensor(Float32, (batch_partial_sym, N), div) if db_partial is not None else None _rmsnorm_bwd.compile_cache[compile_key] = cute.compile( RMSNormBackward(dtype, N), x_cute, @@ -931,9 +871,7 @@ def rmsnorm_bwd( dw_partial = None db_partial = torch.empty(sm_count, N, device=device, dtype=torch.float32) if has_bias else None - _rmsnorm_bwd( - x, weight, dout, rstd, dx, dw_partial, db_partial, dresidual_out, dresidual, sm_count - ) + _rmsnorm_bwd(x, weight, dout, rstd, dx, dw_partial, db_partial, dresidual_out, dresidual, sm_count) # we have summed the partial gradients in fp32, now we convert back to the weight dtype dw = dw_partial.sum(dim=0).to(weight.dtype) if weight is not None else None @@ -1050,9 +988,7 @@ class QuackRMSNorm(torch.nn.RMSNorm): eps (float): A small constant for numerical stability """ - def __init__( - self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True, device=None, dtype=None - ): + def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True, device=None, dtype=None): super().__init__(dim, eps, elementwise_affine, device=device, dtype=dtype) def forward(self, x: Tensor) -> Tensor: diff --git a/src/xorl/ops/quack/sm100_utils.py b/src/xorl/ops/quack/sm100_utils.py index 2c12a38b..2dafb00b 100644 --- a/src/xorl/ops/quack/sm100_utils.py +++ b/src/xorl/ops/quack/sm100_utils.py @@ -33,9 +33,7 @@ def make_smem_layout_cpasync_a( """ is_k_major = tiled_mma.op.a_major_mode == OperandMajorMode.K - a_smem_shape = tiled_mma.partition_shape_A( - cute.dice(mma_tiler_mnk, (1, None, 1), loc=loc, ip=ip) - ) + a_smem_shape = tiled_mma.partition_shape_A(cute.dice(mma_tiler_mnk, (1, None, 1), loc=loc, ip=ip)) a_smem_shape_mn_k = ( cute.size(a_smem_shape[0][0], loc=loc, ip=ip) * a_smem_shape[1], cute.size(a_smem_shape[0][1], loc=loc, ip=ip) * a_smem_shape[2], diff --git a/src/xorl/ops/quack/sm90_utils.py b/src/xorl/ops/quack/sm90_utils.py index 36bf3251..a8746bb3 100644 --- a/src/xorl/ops/quack/sm90_utils.py +++ b/src/xorl/ops/quack/sm90_utils.py @@ -1,13 +1,13 @@ # Copyright (c) 2025, Tri Dao. -from typing import Type, Union, Optional +from typing import Optional, Type, Union import cutlass import cutlass.cute as cute import cutlass.utils.hopper_helpers as sm90_utils_og +from cutlass import Boolean, Float32, Int32, const_expr from cutlass.cute.nvgpu import warpgroup from cutlass.cutlass_dsl import Numeric, dsl_user_op -from cutlass import Float32, Int32, Boolean, const_expr from cutlass.utils import LayoutEnum @@ -98,9 +98,7 @@ def gemm_zero_init( swap_AB: bool = False, ) -> cute.Tensor: if const_expr(swap_AB): - return gemm_zero_init( - tiled_mma, shape[::-1], tCrB, tCrA, B_idx, A_idx, wg_wait, swap_AB=False - ) + return gemm_zero_init(tiled_mma, shape[::-1], tCrB, tCrA, B_idx, A_idx, wg_wait, swap_AB=False) else: acc = cute.make_rmem_tensor(tiled_mma.partition_shape_C(shape), Float32) rA = tCrA if const_expr(A_idx is None) else tCrA[None, None, None, A_idx] @@ -146,9 +144,7 @@ def partition_fragment_ABC( assert sB is not None tCrB = thr_mma.make_fragment_B(thr_mma.partition_B(sB)) else: - acc = cute.make_rmem_tensor( - thr_mma.partition_shape_C((shape_mnk[1], shape_mnk[0])), Float32 - ) + acc = cute.make_rmem_tensor(thr_mma.partition_shape_C((shape_mnk[1], shape_mnk[0])), Float32) if const_expr(not is_rs): assert sB is not None tCrB = thr_mma.make_fragment_A(thr_mma.partition_A(sB)) diff --git a/src/xorl/ops/quack/softmax.py b/src/xorl/ops/quack/softmax.py index c01157d7..18bcf849 100644 --- a/src/xorl/ops/quack/softmax.py +++ b/src/xorl/ops/quack/softmax.py @@ -1,23 +1,20 @@ # Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao. import math -from typing import Type from functools import partial - -import torch +from typing import Type import cuda.bindings.driver as cuda - import cutlass import cutlass.cute as cute -from cutlass import Int64, Float32, const_expr +import torch +from cutlass import Float32, Int64, const_expr -from . import utils -from . import copy_utils +from . import copy_utils, utils from .compile_utils import make_fake_tensor as fake_tensor -from .reduce import row_reduce, online_softmax_reduce -from .reduction_base import ReductionBase from .cute_dsl_utils import torch2cute_dtype_map +from .reduce import online_softmax_reduce, row_reduce +from .reduction_base import ReductionBase class Softmax(ReductionBase): @@ -60,9 +57,7 @@ def __call__( assert mX.element_type == self.dtype self._set_cluster_n() largest_dtype_width = const_expr(max(t.element_type.width for t in [mX, mO])) - tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy( - vecsize=128 // largest_dtype_width - ) + tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=128 // largest_dtype_width) num_threads = tiled_copy.size self.kernel(mX, mO, tiler_mn, tiled_copy, threads_per_row).launch( grid=[cute.ceil_div(mX.shape[0], tiler_mn[0]), self.cluster_n, 1], @@ -92,9 +87,7 @@ def kernel( gX, gO, cX = [cute.local_tile(mT, tiler_mn, (bidx, cluster_y)) for mT in (mX, mO, idX)] smem = cutlass.utils.SmemAllocator() - sX = smem.allocate_tensor( - mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16 - ) + sX = smem.allocate_tensor(mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16) reduction_buffer, mbar_ptr = self._allocate_reduction_buffer_and_mbar(smem, tv_layout) thr_copy_X = tiled_copy.get_slice(tidx) @@ -106,11 +99,7 @@ def kernel( tXrX, tXrO = [cute.make_fragment_like(thr) for thr in (tXgX, tXgO)] is_even_N = const_expr(shape[1] == tiler_mn[1] * self.cluster_n) - tXpX = ( - None - if is_even_N - else copy_utils.predicate_k(thr_copy_X.partition_S(cX), limit=shape[1]) - ) + tXpX = None if is_even_N else copy_utils.predicate_k(thr_copy_X.partition_S(cX), limit=shape[1]) # Each copy will use the same predicate copy = partial(copy_utils.copy, pred=tXpX) @@ -239,9 +228,7 @@ def __call__( assert mdY.element_type == self.dtype self._set_cluster_n() largest_dtype_width = const_expr(max(t.element_type.width for t in [mdY, mY, mdX])) - tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy( - vecsize=128 // largest_dtype_width - ) + tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=128 // largest_dtype_width) num_threads = tiled_copy.size self.kernel(mdY, mY, mdX, tiler_mn, tiled_copy, threads_per_row).launch( grid=[cute.ceil_div(mdY.shape[0], tiler_mn[0]), self.cluster_n, 1], @@ -268,17 +255,13 @@ def kernel( shape = mdY.shape idX = cute.make_identity_tensor(shape) # slice for CTAs - gdY, gY, gdX, cX = [ - cute.local_tile(mT, tiler_mn, (bidx, cluster_y)) for mT in (mdY, mY, mdX, idX) - ] + gdY, gY, gdX, cX = [cute.local_tile(mT, tiler_mn, (bidx, cluster_y)) for mT in (mdY, mY, mdX, idX)] smem = cutlass.utils.SmemAllocator() sdY = smem.allocate_tensor( mdY.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16 ) - sY = smem.allocate_tensor( - mY.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16 - ) + sY = smem.allocate_tensor(mY.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16) reduction_buffer, mbar_ptr = self._allocate_reduction_buffer_and_mbar(smem, tv_layout) thr_copy = tiled_copy.get_slice(tidx) @@ -292,9 +275,7 @@ def kernel( tdYrdY, tYrY, tdXrdX = [cute.make_fragment_like(thr) for thr in (tdYgdY, tYgY, tdXgdX)] is_even_N = const_expr(shape[1] == tiler_mn[1] * self.cluster_n) - tXpX = ( - None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) - ) + tXpX = None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) # Each copy will use the same predicate copy = partial(copy_utils.copy, pred=tXpX) @@ -353,9 +334,7 @@ def _softmax_backward(dy: torch.Tensor, y: torch.Tensor, dx: torch.Tensor) -> No if compile_key not in _softmax_backward.compile_cache: batch_sym = cute.sym_int() div = math.gcd(128 // dtype.width, N) - dy_cute, y_cute, dx_cute = [ - fake_tensor(dt, (batch_sym, N), div) for dt in [dtype, y_dtype, dx_dtype] - ] + dy_cute, y_cute, dx_cute = [fake_tensor(dt, (batch_sym, N), div) for dt in [dtype, y_dtype, dx_dtype]] softmax_backward_op = SoftmaxBackward(dtype, N) _softmax_backward.compile_cache[compile_key] = cute.compile( softmax_backward_op, diff --git a/src/xorl/ops/quack/sort/__init__.py b/src/xorl/ops/quack/sort/__init__.py index efdc8303..ce7ade9d 100644 --- a/src/xorl/ops/quack/sort/__init__.py +++ b/src/xorl/ops/quack/sort/__init__.py @@ -1,4 +1,5 @@ from .bitonic_sort import bitonic_sort, bitonic_topk from .sorting_networks import optimal_sort + __all__ = ["bitonic_sort", "bitonic_topk", "optimal_sort"] diff --git a/src/xorl/ops/quack/sort/bitonic_sort.py b/src/xorl/ops/quack/sort/bitonic_sort.py index 77fa8207..7374cda4 100644 --- a/src/xorl/ops/quack/sort/bitonic_sort.py +++ b/src/xorl/ops/quack/sort/bitonic_sort.py @@ -5,11 +5,11 @@ import cutlass import cutlass.cute as cute -from cutlass import Int32, Float32, const_expr +from cutlass import Float32, Int32, const_expr from .. import utils -from .utils import compare_and_swap from .sorting_networks import optimal_sort +from .utils import compare_and_swap @cute.jit diff --git a/src/xorl/ops/quack/sort/generate_sorting_networks.py b/src/xorl/ops/quack/sort/generate_sorting_networks.py old mode 100644 new mode 100755 index 25d10151..d83afc70 --- a/src/xorl/ops/quack/sort/generate_sorting_networks.py +++ b/src/xorl/ops/quack/sort/generate_sorting_networks.py @@ -9,7 +9,8 @@ import argparse import os import re -from typing import List, Tuple, Dict +from typing import Dict, List, Tuple + # Network strings from bertdobbelaere.github.io/sorting_networks.html # Copy-paste network strings here, then run initialize_networks() to parse them @@ -178,9 +179,7 @@ def add_network_from_string(size: int, network_str: str, description: str = ""): return False -def generate_networks_dict( - networks_data: Dict[int, Tuple[int, int, List[List[Tuple[int, int]]]]] -) -> str: +def generate_networks_dict(networks_data: Dict[int, Tuple[int, int, List[List[Tuple[int, int]]]]]) -> str: """Generate the global networks dictionary.""" lines = ["networks = {"] @@ -296,9 +295,7 @@ def main(): default=64, help="Maximum sorting network size to generate (default: 32)", ) - parser.add_argument( - "--stats", "-s", action="store_true", help="Print statistics about the optimal networks" - ) + parser.add_argument("--stats", "-s", action="store_true", help="Print statistics about the optimal networks") args = parser.parse_args() diff --git a/src/xorl/ops/quack/sort/utils.py b/src/xorl/ops/quack/sort/utils.py index 0237e88a..d542cb4d 100644 --- a/src/xorl/ops/quack/sort/utils.py +++ b/src/xorl/ops/quack/sort/utils.py @@ -5,9 +5,7 @@ @cute.jit -def compare_and_swap( - arr: cute.Tensor, i: int, j: int, ascending: bool = True, use_selection: bool = False -) -> None: +def compare_and_swap(arr: cute.Tensor, i: int, j: int, ascending: bool = True, use_selection: bool = False) -> None: """Compare and swap elements at indices i and j in ascending or descending order.""" if const_expr(use_selection): a, b = arr[i], arr[j] diff --git a/src/xorl/ops/quack/tensormap_manager.py b/src/xorl/ops/quack/tensormap_manager.py index 9241f8cd..fe9e22d0 100644 --- a/src/xorl/ops/quack/tensormap_manager.py +++ b/src/xorl/ops/quack/tensormap_manager.py @@ -1,13 +1,13 @@ # Copyright (c) 2025, Tri Dao. -from typing import Tuple from dataclasses import dataclass +from typing import Tuple import cutlass import cutlass.cute as cute -from cutlass.cutlass_dsl import Boolean, const_expr, Int32 -from cutlass.utils import TensorMapUpdateMode, TensorMapManager from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import Boolean, Int32, const_expr +from cutlass.utils import TensorMapManager, TensorMapUpdateMode @dataclass(frozen=True) @@ -41,9 +41,7 @@ def update_tensormap( # updates before touching tensormap in global memory if is_manager_warp: if const_expr(self.tensormap_update_mode == TensorMapUpdateMode.SMEM): - for copy_atom, tensor, smem_ptr in zip( - tma_copy_atom, tensor_gmem, tensormap_smem_ptr - ): + for copy_atom, tensor, smem_ptr in zip(tma_copy_atom, tensor_gmem, tensormap_smem_ptr): cute.nvgpu.cpasync.update_tma_descriptor(copy_atom, tensor, smem_ptr) # wait until it's safe to update tensormap in global memory with cute.arch.elect_one(): @@ -55,9 +53,7 @@ def update_tensormap( for gmem_ptr, smem_ptr in zip(tensormap_gmem_ptr, tensormap_smem_ptr): cute.nvgpu.cpasync.cp_fence_tma_desc_release(gmem_ptr, smem_ptr) else: - for copy_atom, tensor, gmem_ptr in zip( - tma_copy_atom, tensor_gmem, tensormap_gmem_ptr - ): + for copy_atom, tensor, gmem_ptr in zip(tma_copy_atom, tensor_gmem, tensormap_gmem_ptr): cute.nvgpu.cpasync.update_tma_descriptor(copy_atom, tensor, gmem_ptr) cute.arch.sync_warp() cute.nvgpu.cpasync.fence_tma_desc_release() diff --git a/src/xorl/ops/quack/tile_scheduler.py b/src/xorl/ops/quack/tile_scheduler.py index 53496722..374c4f5e 100644 --- a/src/xorl/ops/quack/tile_scheduler.py +++ b/src/xorl/ops/quack/tile_scheduler.py @@ -1,17 +1,17 @@ # Copyright (c) 2025, Tri Dao. -from typing import Tuple, Optional from dataclasses import dataclass from enum import IntEnum +from typing import Optional, Tuple import cutlass import cutlass.cute as cute -from cutlass import Int32, Float32, Boolean, const_expr +from cutlass import Boolean, Float32, Int32, const_expr from . import utils +from .cute_dsl_utils import ArgumentsBase, ParamsBase from .fast_math import FastDivmod from .pipeline import PipelineStateWAdvance -from .cute_dsl_utils import ArgumentsBase, ParamsBase class RasterOrderOption(IntEnum): @@ -36,17 +36,11 @@ class PersistenceMode(IntEnum): def get_raster_order_from_option( raster_order_option: RasterOrderOption, problem_shape_ncluster_mn: cute.Shape, group_size: Int32 ) -> RasterOrder: - raster_order = ( - RasterOrder.AlongM - if raster_order_option == RasterOrderOption.AlongM - else RasterOrder.AlongN - ) + raster_order = RasterOrder.AlongM if raster_order_option == RasterOrderOption.AlongM else RasterOrder.AlongN if raster_order_option == RasterOrderOption.Heuristic: problem_blocks_m = cute.round_up(problem_shape_ncluster_mn[0], group_size) problem_blocks_n = cute.round_up(problem_shape_ncluster_mn[1], group_size) - raster_order = ( - RasterOrder.AlongM if problem_blocks_n > problem_blocks_m else RasterOrder.AlongN - ) + raster_order = RasterOrder.AlongM if problem_blocks_n > problem_blocks_m else RasterOrder.AlongN return raster_order @@ -93,22 +87,14 @@ def create(args: TileSchedulerArguments, *, loc=None, ip=None) -> "TileScheduler cluster_shape_mn = const_expr(cute.select(args.cluster_shape_mnk, mode=[0, 1])) problem_shape_ntile_mn = cute.select(args.problem_shape_ntile_mnl, mode=[0, 1]) problem_shape_ncluster_mn = cute.ceil_div(problem_shape_ntile_mn, cluster_shape_mn) - problem_shape_ncluster_mnl = problem_shape_ncluster_mn + ( - args.problem_shape_ntile_mnl[2], - ) + problem_shape_ncluster_mnl = problem_shape_ncluster_mn + (args.problem_shape_ntile_mnl[2],) num_clusters_per_problem = cute.size(problem_shape_ncluster_mn) - raster_order = get_raster_order_from_option( - args.raster_order, problem_shape_ncluster_mn, args.group_size - ) + raster_order = get_raster_order_from_option(args.raster_order, problem_shape_ncluster_mn, args.group_size) ncluster_fast = ( - problem_shape_ncluster_mn[0] - if raster_order == RasterOrder.AlongM - else problem_shape_ncluster_mn[1] + problem_shape_ncluster_mn[0] if raster_order == RasterOrder.AlongM else problem_shape_ncluster_mn[1] ) ncluster_slow = ( - problem_shape_ncluster_mn[1] - if raster_order == RasterOrder.AlongM - else problem_shape_ncluster_mn[0] + problem_shape_ncluster_mn[1] if raster_order == RasterOrder.AlongM else problem_shape_ncluster_mn[0] ) group_size = min(args.group_size, ncluster_fast) group_size_tail = ncluster_fast % group_size @@ -125,9 +111,7 @@ def create(args: TileSchedulerArguments, *, loc=None, ip=None) -> "TileScheduler # Don't divide by 0 FastDivmod(group_size_tail if group_size_tail > 0 else 1), FastDivmod(num_clusters_in_group), - args.tile_count_semaphore - if const_expr(args.persistence_mode == PersistenceMode.DYNAMIC) - else None, + args.tile_count_semaphore if const_expr(args.persistence_mode == PersistenceMode.DYNAMIC) else None, args.batch_idx_permute, cluster_shape_mn, args.persistence_mode, @@ -208,8 +192,7 @@ def create( ) stages = 0 if const_expr( - params.persistence_mode - in [PersistenceMode.STATIC, PersistenceMode.DYNAMIC, PersistenceMode.CLC] + params.persistence_mode in [PersistenceMode.STATIC, PersistenceMode.DYNAMIC, PersistenceMode.CLC] ): assert sched_smem is not None assert scheduler_pipeline is not None @@ -246,9 +229,9 @@ def get_grid_shape( params.problem_shape_ncluster_mnl[2], ) else: - num_ctas_in_problem = cute.size( - params.problem_shape_ncluster_mnl, loc=loc, ip=ip - ) * cute.size(params.cluster_shape_mn) + num_ctas_in_problem = cute.size(params.problem_shape_ncluster_mnl, loc=loc, ip=ip) * cute.size( + params.cluster_shape_mn + ) num_ctas_per_cluster = cute.size(params.cluster_shape_mn, loc=loc, ip=ip) # Total ctas that can run in one wave num_ctas_per_wave = max_active_clusters * num_ctas_per_cluster @@ -257,9 +240,7 @@ def get_grid_shape( return (*params.cluster_shape_mn, num_persistent_clusters) @cute.jit - def _swizzle_cta( - self, cluster_id_in_problem: Int32, *, loc=None, ip=None - ) -> Tuple[Int32, Int32]: + def _swizzle_cta(self, cluster_id_in_problem: Int32, *, loc=None, ip=None) -> Tuple[Int32, Int32]: # CTA Swizzle to promote L2 data reuse params = self.params group_id, id_in_group = divmod(cluster_id_in_problem, params.num_clusters_in_group_fdd) @@ -323,14 +304,8 @@ def _delinearize_work_idx( if const_expr(bidz is not None): bidz_ = bidz cid_m, cid_n = self._swizzle_cta(cluster_id_in_problem, loc=loc, ip=ip) - pid_m, pid_n = self._cluster_id_to_cta_id( - cid_m, cid_n, block_zero_only=block_zero_only, loc=loc, ip=ip - ) - batch_idx = ( - bidz_ - if const_expr(params.batch_idx_permute is None) - else params.batch_idx_permute[bidz_] - ) + pid_m, pid_n = self._cluster_id_to_cta_id(cid_m, cid_n, block_zero_only=block_zero_only, loc=loc, ip=ip) + batch_idx = bidz_ if const_expr(params.batch_idx_permute is None) else params.batch_idx_permute[bidz_] tile_coord_mnkl = (pid_m, pid_n, None, batch_idx) return cutlass.utils.WorkTileInfo(tile_coord_mnkl, is_valid) @@ -344,9 +319,7 @@ def get_current_work(self, *, loc=None, ip=None) -> cutlass.utils.WorkTileInfo: # return self._delinearize_work_idx(loc=loc, ip=ip) else: self._scheduler_pipeline.consumer_wait(self._pipeline_state) - pid_m, pid_n, batch_idx, is_valid_i32 = [ - self._sched_smem[i, self._pipeline_state.index] for i in range(4) - ] + pid_m, pid_n, batch_idx, is_valid_i32 = [self._sched_smem[i, self._pipeline_state.index] for i in range(4)] # Need this fence since the STAS from the producer is using the async proxy. # Without this, we get race condition / deadlock. if const_expr(cute.size(params.cluster_shape_mn) > 1): @@ -408,17 +381,13 @@ def _fetch_next_work_idx(self, *, loc=None, ip=None) -> Int32 | Tuple[Int32, Int bidy // params.cluster_shape_mn[1], bidz, ) - cluster_idx, batch_idx = type(self)._cluster_idx_to_work_idx_batch( - params, cluster_idx, loc=loc, ip=ip - ) + cluster_idx, batch_idx = type(self)._cluster_idx_to_work_idx_batch(params, cluster_idx, loc=loc, ip=ip) return cluster_idx, batch_idx, Boolean(valid) else: return Int32(0) @cute.jit - def write_work_tile_to_smem( - self, work_tile_info: cutlass.utils.WorkTileInfo, *, loc=None, ip=None - ): + def write_work_tile_to_smem(self, work_tile_info: cutlass.utils.WorkTileInfo, *, loc=None, ip=None): params = self.params if const_expr(self._sched_smem is not None): # producer phase is always consumer_phase ^ 1 @@ -572,16 +541,12 @@ class Params(ParamsBase): @staticmethod @cute.jit - def create( - args: TileSchedulerArguments, *, loc=None, ip=None - ) -> "TriangularTileScheduler.Params": + def create(args: TileSchedulerArguments, *, loc=None, ip=None) -> "TriangularTileScheduler.Params": assert args.cluster_shape_mnk[2] == 1 cluster_shape_mn = const_expr(cute.select(args.cluster_shape_mnk, mode=[0, 1])) problem_shape_ntile_mn = cute.select(args.problem_shape_ntile_mnl, mode=[0, 1]) problem_shape_ncluster_mn = cute.ceil_div(problem_shape_ntile_mn, cluster_shape_mn) - problem_shape_ncluster_mnl = problem_shape_ncluster_mn + ( - args.problem_shape_ntile_mnl[2], - ) + problem_shape_ncluster_mnl = problem_shape_ncluster_mn + (args.problem_shape_ntile_mnl[2],) cluster_m = problem_shape_ncluster_mn[0] # Assume that each cluster is responsible for a square tile num_clusters_per_problem = cluster_m * (cluster_m + 1) // 2 @@ -600,9 +565,7 @@ def create( FastDivmod(group_size_tail if group_size_tail > 0 else 1), FastDivmod(group_size * group_size), FastDivmod((group_size_tail if group_size_tail > 0 else 1) * group_size), - args.tile_count_semaphore - if const_expr(args.persistence_mode == PersistenceMode.DYNAMIC) - else None, + args.tile_count_semaphore if const_expr(args.persistence_mode == PersistenceMode.DYNAMIC) else None, cluster_shape_mn, args.persistence_mode, ) @@ -627,8 +590,7 @@ def create( ) stages = 0 if const_expr( - params.persistence_mode - in [PersistenceMode.STATIC, PersistenceMode.DYNAMIC, PersistenceMode.CLC] + params.persistence_mode in [PersistenceMode.STATIC, PersistenceMode.DYNAMIC, PersistenceMode.CLC] ): assert sched_smem is not None assert scheduler_pipeline is not None @@ -674,32 +636,19 @@ def get_grid_shape( return (*params.cluster_shape_mn, num_persistent_clusters) @cute.jit - def _swizzle_cta( - self, cluster_id_in_problem: Int32, *, loc=None, ip=None - ) -> Tuple[Int32, Int32]: + def _swizzle_cta(self, cluster_id_in_problem: Int32, *, loc=None, ip=None) -> Tuple[Int32, Int32]: # CTA Swizzle to promote L2 data reuse params = self.params group_size = params.group_size_fdd.divisor - group_id = ( - utils.ceil( - (utils.sqrt(2 * cluster_id_in_problem + 2.25) - 0.5) * params.group_size_inv_f32 - ) - - 1 - ) + group_id = utils.ceil((utils.sqrt(2 * cluster_id_in_problem + 2.25) - 0.5) * params.group_size_inv_f32) - 1 cid_m_start = group_id * group_size id_in_group = cluster_id_in_problem - (cid_m_start * (cid_m_start + 1)) // 2 - group_size_actual = ( - group_size - if group_id < params.num_groups_regular - else params.group_size_tail_fdd.divisor - ) + group_size_actual = group_size if group_id < params.num_groups_regular else params.group_size_tail_fdd.divisor group_col, group_remainder = Int32(0), Int32(0) if group_id < params.num_groups_regular: group_col, group_remainder = divmod(id_in_group, params.group_size_mul_group_size_fdd) else: # tail part - group_col, group_remainder = divmod( - id_in_group, params.group_size_tail_mul_group_size_fdd - ) + group_col, group_remainder = divmod(id_in_group, params.group_size_tail_mul_group_size_fdd) cid_m_in_group, cid_n_in_group = Int32(0), Int32(0) if id_in_group >= group_size_actual * group_size * group_id: # triangular tail cid_m_in_group, cid_n_in_group = triangular_idx_to_coord(group_remainder) @@ -728,11 +677,7 @@ def _delinearize_work_idx( if const_expr(params.persistence_mode == PersistenceMode.NONE): is_valid = self.num_tiles_executed == 0 else: - is_valid = ( - work_idx - < params.num_clusters_per_problem_fdd.divisor - * params.problem_shape_ncluster_mnl[2] - ) + is_valid = work_idx < params.num_clusters_per_problem_fdd.divisor * params.problem_shape_ncluster_mnl[2] pid_m, pid_n, batch_idx = Int32(0), Int32(0), Int32(0) if is_valid: if const_expr(params.persistence_mode in [PersistenceMode.NONE, PersistenceMode.CLC]): @@ -743,9 +688,7 @@ def _delinearize_work_idx( if const_expr(bidz is not None): bidz_ = bidz cid_m, cid_n = self._swizzle_cta(cluster_id_in_problem, loc=loc, ip=ip) - pid_m, pid_n = self._cluster_id_to_cta_id( - cid_m, cid_n, block_zero_only=block_zero_only, loc=loc, ip=ip - ) + pid_m, pid_n = self._cluster_id_to_cta_id(cid_m, cid_n, block_zero_only=block_zero_only, loc=loc, ip=ip) batch_idx = bidz_ tile_coord_mnkl = (pid_m, pid_n, None, batch_idx) # tidx, _, _ = cute.arch.thread_idx() @@ -786,9 +729,7 @@ class Params(ParamsBase): @staticmethod @cute.jit - def create( - args: TileSchedulerArguments, *, loc=None, ip=None - ) -> "VarlenMTileScheduler.Params": + def create(args: TileSchedulerArguments, *, loc=None, ip=None) -> "VarlenMTileScheduler.Params": assert args.cluster_shape_mnk[2] == 1 cluster_shape_mn = const_expr(cute.select(args.cluster_shape_mnk, mode=[0, 1])) # problem_shape_ntile_mnl[0] will be None for VarlenM @@ -797,20 +738,14 @@ def create( None, cute.ceil_div(problem_shape_ntile_mn[1], cluster_shape_mn[1]), ) - problem_shape_ncluster_mnl = problem_shape_ncluster_mn + ( - args.problem_shape_ntile_mnl[2], - ) + problem_shape_ncluster_mnl = problem_shape_ncluster_mn + (args.problem_shape_ntile_mnl[2],) raster_order = const_expr( RasterOrder.AlongM if args.raster_order == RasterOrderOption.AlongM else RasterOrder.AlongN # For Heuristic we also use AlongN ) - ncluster_fast = problem_shape_ncluster_mn[ - 0 if raster_order == RasterOrder.AlongM else 1 - ] - ncluster_slow = problem_shape_ncluster_mn[ - 1 if raster_order == RasterOrder.AlongM else 0 - ] + ncluster_fast = problem_shape_ncluster_mn[0 if raster_order == RasterOrder.AlongM else 1] + ncluster_slow = problem_shape_ncluster_mn[1 if raster_order == RasterOrder.AlongM else 0] if const_expr(ncluster_fast is not None): group_size = min(args.group_size, ncluster_fast) group_size_tail = ncluster_fast % group_size @@ -829,14 +764,10 @@ def create( group_size, FastDivmod(group_size) if ncluster_fast is not None else None, # Don't divide by 0 - FastDivmod(group_size_tail if group_size_tail > 0 else 1) - if group_size_tail is not None - else None, + FastDivmod(group_size_tail if group_size_tail > 0 else 1) if group_size_tail is not None else None, FastDivmod(num_clusters_in_group) if num_clusters_in_group is not None else None, args.tile_shape_mn, - args.tile_count_semaphore - if const_expr(args.persistence_mode == PersistenceMode.DYNAMIC) - else None, + args.tile_count_semaphore if const_expr(args.persistence_mode == PersistenceMode.DYNAMIC) else None, cluster_shape_mn, args.persistence_mode, ) @@ -895,8 +826,7 @@ def create( ) stages = 0 if const_expr( - params.persistence_mode - in [PersistenceMode.STATIC, PersistenceMode.DYNAMIC, PersistenceMode.CLC] + params.persistence_mode in [PersistenceMode.STATIC, PersistenceMode.DYNAMIC, PersistenceMode.CLC] ): assert sched_smem is not None assert scheduler_pipeline is not None @@ -959,16 +889,12 @@ def _swizzle_cta( cid_slow, cid_fast_in_group = divmod(id_in_group, params.group_size_tail_fdd) else: assert params.raster_order == RasterOrder.AlongM - group_size_actual = cutlass.min( - params.group_size, num_clusters_m - group_id * params.group_size - ) + group_size_actual = cutlass.min(params.group_size, num_clusters_m - group_id * params.group_size) cid_slow = id_in_group // group_size_actual cid_fast_in_group = id_in_group - cid_slow * group_size_actual if group_id % 2 == 1: # serpentine order ncluster_slow = ( - params.problem_shape_ncluster_mnl[1] - if params.raster_order == RasterOrder.AlongM - else num_clusters_m + params.problem_shape_ncluster_mnl[1] if params.raster_order == RasterOrder.AlongM else num_clusters_m ) cid_slow = ncluster_slow - 1 - cid_slow cid_fast = group_id * params.group_size + cid_fast_in_group @@ -978,9 +904,7 @@ def _swizzle_cta( return cid_m, cid_n @cute.jit - def _get_num_m_blocks( - self, lane: Int32, bidb_start: Int32, block_size: cutlass.Constexpr[int] - ) -> Int32: + def _get_num_m_blocks(self, lane: Int32, bidb_start: Int32, block_size: cutlass.Constexpr[int]) -> Int32: num_batch = self.params.problem_shape_ncluster_mnl[2] batch_idx = lane + bidb_start cur_cu_seqlen = Int32(0) @@ -989,9 +913,7 @@ def _get_num_m_blocks( next_cu_seqlen = cute.arch.shuffle_sync_down(cur_cu_seqlen, offset=1) seqlen = next_cu_seqlen - cur_cu_seqlen return ( - cute.ceil_div(seqlen, block_size) - if batch_idx < num_batch and lane < cute.arch.WARP_SIZE - 1 - else Int32(0) + cute.ceil_div(seqlen, block_size) if batch_idx < num_batch and lane < cute.arch.WARP_SIZE - 1 else Int32(0) ) @cute.jit @@ -1020,15 +942,11 @@ def _delinearize_work_idx( is_valid = is_valid_ if is_valid: while problems_end_tile <= next_tile_idx: - num_clusters_m = self._get_num_m_blocks( - lane_idx, bidb_start=batch_idx, block_size=block_size - ) + num_clusters_m = self._get_num_m_blocks(lane_idx, bidb_start=batch_idx, block_size=block_size) num_clusters = num_clusters_m * params.problem_shape_ncluster_mnl[1] num_clusters_cumulative = utils.warp_prefix_sum(num_clusters, lane_idx) # Total number of blocks for the next 31 problems, same for all lanes - clusters_in_problems = cute.arch.shuffle_sync( - num_clusters_cumulative, cute.arch.WARP_SIZE - 1 - ) + clusters_in_problems = cute.arch.shuffle_sync(num_clusters_cumulative, cute.arch.WARP_SIZE - 1) problems_end_tile += clusters_in_problems if problems_end_tile <= next_tile_idx: batch_idx += cute.arch.WARP_SIZE - 1 @@ -1049,9 +967,7 @@ def _delinearize_work_idx( # The next problem to process is the first one that does not have ending tile position # that is greater than or equal to tile index. batch_idx_in_problems = cute.arch.popc( - cute.arch.vote_ballot_sync( - problems_start_tile + num_clusters_cumulative <= next_tile_idx - ) + cute.arch.vote_ballot_sync(problems_start_tile + num_clusters_cumulative <= next_tile_idx) ) batch_idx += batch_idx_in_problems num_clusters_prev_lane = ( @@ -1064,9 +980,7 @@ def _delinearize_work_idx( cluster_id_in_problem = next_tile_idx - num_work_idx_before_cur_batch # if cute.arch.thread_idx()[0] == 128: cute.printf("SingleTileVarlenScheduler: tile_idx=%d, batch_idx=%d, cid_n=%d, cid_m=%d, is_valid = %d", self._tile_idx, batch_idx, cid_n, cid_m, is_valid) cid_m, cid_n = self._swizzle_cta(cluster_id_in_problem, num_clusters_m, loc=loc, ip=ip) - pid_m, pid_n = self._cluster_id_to_cta_id( - cid_m, cid_n, block_zero_only=block_zero_only, loc=loc, ip=ip - ) + pid_m, pid_n = self._cluster_id_to_cta_id(cid_m, cid_n, block_zero_only=block_zero_only, loc=loc, ip=ip) tile_coord_mnkl = (pid_m, pid_n, None, batch_idx) self._current_batch_idx = batch_idx self._num_work_idx_before_cur_batch = num_work_idx_before_cur_batch diff --git a/src/xorl/ops/quack/topk.py b/src/xorl/ops/quack/topk.py index a87eea14..40c19d69 100644 --- a/src/xorl/ops/quack/topk.py +++ b/src/xorl/ops/quack/topk.py @@ -2,22 +2,19 @@ import math from functools import partial -from typing import Type, Optional - -import torch +from typing import Optional, Type import cuda.bindings.driver as cuda - import cutlass import cutlass.cute as cute -from cutlass import Int32, Float32, const_expr +import torch +from cutlass import Float32, Int32, const_expr -from . import utils -from . import copy_utils +from . import copy_utils, utils from .compile_utils import make_fake_tensor as fake_tensor -from .reduction_base import ReductionBase -from .reduce import row_reduce from .cute_dsl_utils import torch2cute_dtype_map +from .reduce import row_reduce +from .reduction_base import ReductionBase from .sort.bitonic_sort import bitonic_topk @@ -48,9 +45,7 @@ def _get_tiled_copy(self): cols_per_block = num_threads // threads_per_row num_blocks_N = cute.ceil_div(min(N, 16384) // vecsize, threads_per_row) tiler_mn = (cols_per_block, vecsize * num_blocks_N * threads_per_row) - tiled_copy = copy_utils.tiled_copy_2d( - self.dtype, threads_per_row, num_threads, num_copy_elems=vecsize - ) + tiled_copy = copy_utils.tiled_copy_2d(self.dtype, threads_per_row, num_threads, num_copy_elems=vecsize) return tiled_copy, tiler_mn, threads_per_row @cute.jit @@ -98,9 +93,7 @@ def kernel( tXrX = cute.make_fragment_like(tXgX) is_even_N = const_expr(shape[1] == tiler_mn[1]) - tXpX = ( - None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) - ) + tXpX = None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) copy = partial(copy_utils.copy, pred=tXpX) if tXcX[0][0] < shape[0]: @@ -175,9 +168,7 @@ def kernel( # Get max from thread 0 (topk_vals[0] is the max since sorted descending) max_val = cute.arch.shuffle_sync(topk_vals[0], offset=0, mask_and_clamp=mask_and_clamp) log2_e = math.log2(math.e) - exp_x = cute.math.exp2( - topk_vals_split.load() * log2_e - (max_val * log2_e), fastmath=True - ) + exp_x = cute.math.exp2(topk_vals_split.load() * log2_e - (max_val * log2_e), fastmath=True) denom = cute.arch.warp_reduction_sum( exp_x.reduce(cute.ReductionOp.ADD, init_val=0.0, reduction_profile=0), threads_in_group=threads_per_row, @@ -215,9 +206,7 @@ def kernel( @torch.library.custom_op("quack::_topk_fwd", mutates_args={"values", "indices"}) -def _topk_fwd( - x: torch.Tensor, k: int, softmax: bool, values: torch.Tensor, indices: torch.Tensor -) -> None: +def _topk_fwd(x: torch.Tensor, k: int, softmax: bool, values: torch.Tensor, indices: torch.Tensor) -> None: """Top-k forward pass. Args: x: Input tensor of shape (M, N) @@ -295,9 +284,7 @@ def _get_tiled_copy(self, N: int, vecsize: Optional[int] = None): cols_per_block = num_threads // threads_per_row num_blocks_N = cute.ceil_div(N // vecsize, threads_per_row) tiler_mn = (cols_per_block, vecsize * num_blocks_N * threads_per_row) - tiled_copy = copy_utils.tiled_copy_2d( - self.dtype, threads_per_row, num_threads, num_copy_elems=vecsize - ) + tiled_copy = copy_utils.tiled_copy_2d(self.dtype, threads_per_row, num_threads, num_copy_elems=vecsize) return tiled_copy, tiler_mn, threads_per_row @cute.jit @@ -315,9 +302,7 @@ def __call__( assert mIndices.element_type == Int32 self._set_cluster_n() largest_dtype_width = const_expr( - max( - *(t.element_type.width for t in [mdValues, mValues, mIndices, mdX] if t is not None) - ) + max(*(t.element_type.width for t in [mdValues, mValues, mIndices, mdX] if t is not None)) ) vecsize = math.gcd(self.N, 128 // largest_dtype_width) tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(self.N, vecsize=vecsize) @@ -388,9 +373,7 @@ def kernel( is_even_N = const_expr(shape[1] == tiler_mn[1]) tXpV = copy_utils.predicate_k(thr_copy.partition_S(cTopK), limit=mdValues.shape[1]) - tXpX = ( - None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) - ) + tXpX = None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) copy_k = partial(copy_utils.copy, pred=tXpV) copy_dx = partial(copy_utils.copy, pred=tXpX) @@ -429,9 +412,7 @@ def kernel( for n in cutlass.range(tXrdV.shape[2], unroll_full=True): if tXpV[rest_v, 0, n]: for v in cutlass.range(tXrdV.shape[0][0], unroll_full=True): - sdX[row - tile_row_start, tXrI[(v, rest_v), 0, n]] = grad_cvt[ - (v, rest_v), 0, n - ] + sdX[row - tile_row_start, tXrI[(v, rest_v), 0, n]] = grad_cvt[(v, rest_v), 0, n] cute.arch.barrier() # Read from smem to rmem, then write to gmem diff --git a/src/xorl/ops/quack/utils.py b/src/xorl/ops/quack/utils.py index 892d91b4..1dab5c94 100644 --- a/src/xorl/ops/quack/utils.py +++ b/src/xorl/ops/quack/utils.py @@ -5,10 +5,9 @@ import cutlass import cutlass.cute as cute - from cutlass import Float32, Int32, const_expr -from cutlass.cutlass_dsl import T, dsl_user_op from cutlass._mlir.dialects import llvm, nvvm, vector +from cutlass.cutlass_dsl import T, dsl_user_op @dsl_user_op @@ -26,9 +25,7 @@ def load_scalar_or_pointer(x: Float32 | cute.Pointer) -> Float32: @dsl_user_op -def set_block_rank( - smem_ptr: cute.Pointer, peer_cta_rank_in_cluster: Int32, *, loc=None, ip=None -) -> Int32: +def set_block_rank(smem_ptr: cute.Pointer, peer_cta_rank_in_cluster: Int32, *, loc=None, ip=None) -> Int32: """Map the given smem pointer to the address at another CTA rank in the cluster.""" smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value() return Int32( @@ -54,12 +51,8 @@ def store_shared_remote( loc=None, ip=None, ) -> None: - remote_smem_ptr_i32 = set_block_rank( - smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip - ).ir_value() - remote_mbar_ptr_i32 = set_block_rank( - mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip - ).ir_value() + remote_smem_ptr_i32 = set_block_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value() + remote_mbar_ptr_i32 = set_block_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value() if const_expr(isinstance(val, float)): val = Float32(val) assert isinstance(val, (Float32, Int32, cutlass.Int64)), "val must be Float32, Int32, or Int64" @@ -89,12 +82,8 @@ def store_shared_remote_x4( loc=None, ip=None, ) -> None: - remote_smem_ptr_i32 = set_block_rank( - smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip - ).ir_value() - remote_mbar_ptr_i32 = set_block_rank( - mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip - ).ir_value() + remote_smem_ptr_i32 = set_block_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value() + remote_mbar_ptr_i32 = set_block_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value() assert isinstance(val0, (Float32, Int32)), "val must be Float32, or Int32" dtype = Float32 if isinstance(val0, Float32) else Int32 suffix = {Float32: "f32", Int32: "s32"}[dtype] @@ -189,13 +178,9 @@ def fill_oob(tXsX: cute.Tensor, tXpX: Optional[cute.Tensor], fill_value: cute.Nu @dsl_user_op def f32x2_to_i64(a: Float32, b: Float32, *, loc=None, ip=None) -> cutlass.Int64: - vec_f32x2 = vector.from_elements( - T.vector(2, T.f32()), (a.ir_value(), b.ir_value()), loc=loc, ip=ip - ) + vec_f32x2 = vector.from_elements(T.vector(2, T.f32()), (a.ir_value(), b.ir_value()), loc=loc, ip=ip) vec_i64x1 = vector.bitcast(T.vector(1, T.i64()), vec_f32x2) - res = cutlass.Int64( - vector.extract(vec_i64x1, dynamic_position=[], static_position=[0], loc=loc, ip=ip) - ) + res = cutlass.Int64(vector.extract(vec_i64x1, dynamic_position=[], static_position=[0], loc=loc, ip=ip)) return res @@ -203,12 +188,8 @@ def f32x2_to_i64(a: Float32, b: Float32, *, loc=None, ip=None) -> cutlass.Int64: def i64_to_f32x2(c: cutlass.Int64, *, loc=None, ip=None) -> Tuple[Float32, Float32]: vec_i64x1 = vector.from_elements(T.vector(1, T.i64()), (c.ir_value(),), loc=loc, ip=ip) vec_f32x2 = vector.bitcast(T.vector(2, T.f32()), vec_i64x1) - res0 = Float32( - vector.extract(vec_f32x2, dynamic_position=[], static_position=[0], loc=loc, ip=ip) - ) - res1 = Float32( - vector.extract(vec_f32x2, dynamic_position=[], static_position=[1], loc=loc, ip=ip) - ) + res0 = Float32(vector.extract(vec_f32x2, dynamic_position=[], static_position=[0], loc=loc, ip=ip)) + res1 = Float32(vector.extract(vec_f32x2, dynamic_position=[], static_position=[1], loc=loc, ip=ip)) return res0, res1 @@ -232,14 +213,10 @@ def atomic_inc_i32(a: int | Int32, gmem_ptr: cute.Pointer, *, loc=None, ip=None) # * NVVM call based on nvvm version if CUDA_VERSION.major == 12 and CUDA_VERSION.minor == 9: # Old API: requires explicit result type as first positional argument - return nvvm.atomicrmw( - res=T.i32(), op=nvvm.AtomicOpKind.INC, ptr=gmem_ptr.llvm_ptr, a=Int32(a).ir_value() - ) + return nvvm.atomicrmw(res=T.i32(), op=nvvm.AtomicOpKind.INC, ptr=gmem_ptr.llvm_ptr, a=Int32(a).ir_value()) else: # New API: infers result type automatically - return nvvm.atomicrmw( - op=nvvm.AtomicOpKind.INC, ptr=gmem_ptr.llvm_ptr, a=Int32(a).ir_value() - ) + return nvvm.atomicrmw(op=nvvm.AtomicOpKind.INC, ptr=gmem_ptr.llvm_ptr, a=Int32(a).ir_value()) @dsl_user_op @@ -249,14 +226,10 @@ def atomic_add_i32(a: int | Int32, gmem_ptr: cute.Pointer, *, loc=None, ip=None) # * NVVM call based on nvvm version if CUDA_VERSION.major == 12 and CUDA_VERSION.minor == 9: # Old API: requires explicit result type as first positional argument - return nvvm.atomicrmw( - res=T.i32(), op=nvvm.AtomicOpKind.ADD, ptr=gmem_ptr.llvm_ptr, a=Int32(a).ir_value() - ) + return nvvm.atomicrmw(res=T.i32(), op=nvvm.AtomicOpKind.ADD, ptr=gmem_ptr.llvm_ptr, a=Int32(a).ir_value()) else: # New API: infers result type automatically - return nvvm.atomicrmw( - op=nvvm.AtomicOpKind.ADD, ptr=gmem_ptr.llvm_ptr, a=Int32(a).ir_value() - ) + return nvvm.atomicrmw(op=nvvm.AtomicOpKind.ADD, ptr=gmem_ptr.llvm_ptr, a=Int32(a).ir_value()) @dsl_user_op diff --git a/src/xorl/ops/quack/varlen_utils.py b/src/xorl/ops/quack/varlen_utils.py index b265cfbc..3acab74e 100644 --- a/src/xorl/ops/quack/varlen_utils.py +++ b/src/xorl/ops/quack/varlen_utils.py @@ -1,11 +1,11 @@ # Copyright (c) 2025, Tri Dao. -from typing import Optional from dataclasses import dataclass +from typing import Optional import cutlass import cutlass.cute as cute -from cutlass import Int32, Boolean, const_expr +from cutlass import Boolean, Int32, const_expr from cutlass.utils import LayoutEnum from .cute_dsl_utils import ArgumentsBase, ParamsBase @@ -131,9 +131,7 @@ def create( params.tensormaps[tensormap_workspace_idx, 0, None].iterator ) tensormap_b_ptr = tensormap_manager.get_tensormap_ptr( - params.tensormaps[ - tensormap_workspace_idx, 1 if not gather_A else 0, None - ].iterator + params.tensormaps[tensormap_workspace_idx, 1 if not gather_A else 0, None].iterator ) return VarlenManager( params, @@ -202,12 +200,8 @@ def init_tensormap_AB( ) -> None: if const_expr(self.varlen_k): if const_expr(not self.gather_A): - self.tensormap_manager.init_tensormap_from_atom( - tma_atom_a, self._tensormap_a_ptr, is_manager_warp - ) - self.tensormap_manager.init_tensormap_from_atom( - tma_atom_b, self._tensormap_b_ptr, is_manager_warp - ) + self.tensormap_manager.init_tensormap_from_atom(tma_atom_a, self._tensormap_a_ptr, is_manager_warp) + self.tensormap_manager.init_tensormap_from_atom(tma_atom_b, self._tensormap_b_ptr, is_manager_warp) def init_tensormap_epi( self, @@ -217,13 +211,9 @@ def init_tensormap_epi( ) -> None: if const_expr(self.varlen_m): if const_expr(self._tensormap_d_ptr is not None): - self.tensormap_manager.init_tensormap_from_atom( - tma_atom_d, self._tensormap_d_ptr, is_manager_warp - ) + self.tensormap_manager.init_tensormap_from_atom(tma_atom_d, self._tensormap_d_ptr, is_manager_warp) for tma_atom, tensormap_epi_ptr in zip(tma_atoms_epi, self._tensormap_epi_ptrs): - self.tensormap_manager.init_tensormap_from_atom( - tma_atom, tensormap_epi_ptr, is_manager_warp - ) + self.tensormap_manager.init_tensormap_from_atom(tma_atom, tensormap_epi_ptr, is_manager_warp) def fence_tensormap_init(self) -> None: self.tensormap_manager.fence_tensormap_initialization() @@ -270,9 +260,7 @@ def update_tensormap_epi( self._is_group_changed = Boolean(batch_idx != self._last_batch_idx) self._last_batch_idx = batch_idx # Cute-DSL doesn't like this under if statement - order_d = ( - (0 if const_expr(d_layout.is_m_major_c()) else 1) if d_layout is not None else None - ) + order_d = (0 if const_expr(d_layout.is_m_major_c()) else 1) if d_layout is not None else None if self._is_group_changed: # construct tensor A/B based on real address, shape and stride information cu_seqlens_m = self.params.cu_seqlens_m @@ -314,25 +302,19 @@ def fence_tensormap_update_epi(self, is_manager_warp: bool | Boolean = True) -> def get_tma_desc_a_ptr(self) -> Optional[cute.Pointer]: tma_desc_a_ptr = None if const_expr(self.varlen_k and self._tensormap_a_ptr is not None): - tma_desc_a_ptr = self.tensormap_manager.get_tensormap_ptr( - self._tensormap_a_ptr, cute.AddressSpace.generic - ) + tma_desc_a_ptr = self.tensormap_manager.get_tensormap_ptr(self._tensormap_a_ptr, cute.AddressSpace.generic) return tma_desc_a_ptr def get_tma_desc_b_ptr(self) -> Optional[cute.Pointer]: tma_desc_b_ptr = None if const_expr(self.varlen_k): - tma_desc_b_ptr = self.tensormap_manager.get_tensormap_ptr( - self._tensormap_b_ptr, cute.AddressSpace.generic - ) + tma_desc_b_ptr = self.tensormap_manager.get_tensormap_ptr(self._tensormap_b_ptr, cute.AddressSpace.generic) return tma_desc_b_ptr def get_tma_desc_d_ptr(self) -> Optional[cute.Pointer]: tma_desc_d_ptr = None if const_expr(self.varlen_m and self._tensormap_d_ptr is not None): - tma_desc_d_ptr = self.tensormap_manager.get_tensormap_ptr( - self._tensormap_d_ptr, cute.AddressSpace.generic - ) + tma_desc_d_ptr = self.tensormap_manager.get_tensormap_ptr(self._tensormap_d_ptr, cute.AddressSpace.generic) return tma_desc_d_ptr def get_tma_desc_epi_ptrs(self) -> list[Optional[cute.Pointer]]: diff --git a/src/xorl/ops/quantize/__init__.py b/src/xorl/ops/quantize/__init__.py index 6e43bfaf..c26089c0 100644 --- a/src/xorl/ops/quantize/__init__.py +++ b/src/xorl/ops/quantize/__init__.py @@ -1,15 +1,16 @@ +from .block_fp8_gkn_quantize import block_fp8_dequantize_gkn, block_fp8_quantize_gkn +from .block_fp8_quantize import block_fp8_dequantize, block_fp8_gemm, block_fp8_quantize from .fp4_codec import FP4_E2M1_MAX, FP8_E4M3_MAX -from .nf4_codec import NF4_TABLE, NF4_MIN_STEP, get_nf4_lut -from .int4_quantize import int4_quantize, int4_dequantize -from .mxfp4_quantize import mxfp4_quantize, mxfp4_dequantize -from .nf4_quantize import nf4_quantize, nf4_dequantize -from .nvfp4_quantize import nvfp4_quantize, nvfp4_dequantize -from .int4_gkn_quantize import int4_quantize_gkn, int4_dequantize_gkn -from .mxfp4_gkn_quantize import mxfp4_quantize_gkn, mxfp4_dequantize_gkn -from .nf4_gkn_quantize import nf4_quantize_gkn, nf4_dequantize_gkn -from .nvfp4_gkn_quantize import nvfp4_quantize_gkn, nvfp4_dequantize_gkn -from .block_fp8_quantize import block_fp8_quantize, block_fp8_dequantize, block_fp8_gemm -from .block_fp8_gkn_quantize import block_fp8_quantize_gkn, block_fp8_dequantize_gkn +from .int4_gkn_quantize import int4_dequantize_gkn, int4_quantize_gkn +from .int4_quantize import int4_dequantize, int4_quantize +from .mxfp4_gkn_quantize import mxfp4_dequantize_gkn, mxfp4_quantize_gkn +from .mxfp4_quantize import mxfp4_dequantize, mxfp4_quantize +from .nf4_codec import NF4_MIN_STEP, NF4_TABLE, get_nf4_lut +from .nf4_gkn_quantize import nf4_dequantize_gkn, nf4_quantize_gkn +from .nf4_quantize import nf4_dequantize, nf4_quantize +from .nvfp4_gkn_quantize import nvfp4_dequantize_gkn, nvfp4_quantize_gkn +from .nvfp4_quantize import nvfp4_dequantize, nvfp4_quantize + __all__ = [ "FP4_E2M1_MAX", diff --git a/src/xorl/ops/quantize/block_fp8_gkn_quantize.py b/src/xorl/ops/quantize/block_fp8_gkn_quantize.py index b4da21b4..5131c260 100644 --- a/src/xorl/ops/quantize/block_fp8_gkn_quantize.py +++ b/src/xorl/ops/quantize/block_fp8_gkn_quantize.py @@ -23,28 +23,24 @@ # scale tensor shapes consistent. # --------------------------------------------------------------------------- -_quant_configs = [ - triton.Config({}, num_warps=nw, num_stages=ns) - for nw in [4, 8] - for ns in [3, 4, 5] -] +_quant_configs = [triton.Config({}, num_warps=nw, num_stages=ns) for nw in [4, 8] for ns in [3, 4, 5]] -_dequant_configs = [ - triton.Config({}, num_warps=nw, num_stages=ns) - for nw in [4, 8] - for ns in [3, 4, 5] -] +_dequant_configs = [triton.Config({}, num_warps=nw, num_stages=ns) for nw in [4, 8] for ns in [3, 4, 5]] # --------------------------------------------------------------------------- # 2D weight quantization kernel (GKN) # --------------------------------------------------------------------------- + @triton.autotune(configs=_quant_configs, key=["M", "N"]) @triton.jit def _block_fp8_quantize_gkn_kernel( - x_ptr, y_ptr, s_ptr, - M, N, + x_ptr, + y_ptr, + s_ptr, + M, + N, BLOCK_SIZE: tl.constexpr, ): """2D block-based quantization kernel for weight matrices. @@ -71,7 +67,8 @@ def _block_fp8_quantize_gkn_kernel( def block_fp8_quantize_gkn( - x: torch.Tensor, block_size: int = 128, + x: torch.Tensor, + block_size: int = 128, ) -> Tuple[torch.Tensor, torch.Tensor]: """Quantize a 2D weight matrix to FP8 using 2D block-based quantization. @@ -92,14 +89,19 @@ def block_fp8_quantize_gkn( y = torch.empty_like(x, dtype=torch.float8_e4m3fn) s = torch.empty( (triton.cdiv(M, block_size), triton.cdiv(N, block_size)), - dtype=torch.float32, device=x.device, + dtype=torch.float32, + device=x.device, ) grid = lambda meta: ( triton.cdiv(M, block_size), triton.cdiv(N, block_size), ) _block_fp8_quantize_gkn_kernel[grid]( - x, y, s, M, N, + x, + y, + s, + M, + N, BLOCK_SIZE=block_size, ) return y, s @@ -109,11 +111,15 @@ def block_fp8_quantize_gkn( # 2D weight dequantization kernel (GKN) # --------------------------------------------------------------------------- + @triton.autotune(configs=_dequant_configs, key=["M", "N"]) @triton.jit def _block_fp8_dequantize_gkn_kernel( - x_ptr, s_ptr, y_ptr, - M, N, + x_ptr, + s_ptr, + y_ptr, + M, + N, BLOCK_SIZE: tl.constexpr, ): """2D block-based dequantization kernel for weight matrices. @@ -134,9 +140,7 @@ def _block_fp8_dequantize_gkn_kernel( tl.store(y_ptr + offs, y, mask=mask) -def block_fp8_dequantize_gkn( - x: torch.Tensor, s: torch.Tensor, block_size: int = 128 -) -> torch.Tensor: +def block_fp8_dequantize_gkn(x: torch.Tensor, s: torch.Tensor, block_size: int = 128) -> torch.Tensor: """Dequantize a 2D weight matrix using 2D block-based quantization. Args: diff --git a/src/xorl/ops/quantize/block_fp8_quantize.py b/src/xorl/ops/quantize/block_fp8_quantize.py index d026bc30..7639f7ec 100644 --- a/src/xorl/ops/quantize/block_fp8_quantize.py +++ b/src/xorl/ops/quantize/block_fp8_quantize.py @@ -18,6 +18,7 @@ # 1D quantization kernel + wrapper # --------------------------------------------------------------------------- + @triton.jit def _block_fp8_quantize_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr): """1D block-based quantization kernel for FP8 conversion. @@ -41,9 +42,7 @@ def _block_fp8_quantize_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr): tl.store(s_ptr + pid, s) -def block_fp8_quantize( - x: torch.Tensor, block_size: int = 128 -) -> Tuple[torch.Tensor, torch.Tensor]: +def block_fp8_quantize(x: torch.Tensor, block_size: int = 128) -> Tuple[torch.Tensor, torch.Tensor]: """Quantize a tensor to FP8 using 1D block-based quantization. Divides the input tensor into blocks along the last dimension and quantizes @@ -77,6 +76,7 @@ def block_fp8_quantize( # 1D dequantization kernel + wrapper # --------------------------------------------------------------------------- + @triton.jit def _block_fp8_dequantize_kernel(y_ptr, s_ptr, x_ptr, BLOCK_SIZE: tl.constexpr): """1D block-based dequantization kernel. @@ -92,9 +92,7 @@ def _block_fp8_dequantize_kernel(y_ptr, s_ptr, x_ptr, BLOCK_SIZE: tl.constexpr): tl.store(x_ptr + offs, x) -def block_fp8_dequantize( - y: torch.Tensor, s: torch.Tensor, block_size: int = 128 -) -> torch.Tensor: +def block_fp8_dequantize(y: torch.Tensor, s: torch.Tensor, block_size: int = 128) -> torch.Tensor: """Dequantize a tensor using 1D block-based quantization along the last dimension. Args: @@ -183,9 +181,7 @@ def _block_fp8_gemm_kernel( tl.store(c_ptrs, c, mask=mask) -def block_fp8_gemm( - a: torch.Tensor, a_s: torch.Tensor, b: torch.Tensor, b_s: torch.Tensor -) -> torch.Tensor: +def block_fp8_gemm(a: torch.Tensor, a_s: torch.Tensor, b: torch.Tensor, b_s: torch.Tensor) -> torch.Tensor: """Block-wise FP8 matrix multiplication with per-block scaling. Computes C = A @ B^T where A and B are block-quantized FP8 tensors. diff --git a/src/xorl/ops/quantize/fp4_codec.py b/src/xorl/ops/quantize/fp4_codec.py index 46fc5a92..af7a7c2e 100644 --- a/src/xorl/ops/quantize/fp4_codec.py +++ b/src/xorl/ops/quantize/fp4_codec.py @@ -1,4 +1,5 @@ """Shared FP4 E2M1 encode/decode helpers for MXFP4 and NVFP4 kernels.""" + import triton import triton.language as tl diff --git a/src/xorl/ops/quantize/int4_gkn_quantize.py b/src/xorl/ops/quantize/int4_gkn_quantize.py index 8ccf2a94..5e25ba01 100644 --- a/src/xorl/ops/quantize/int4_gkn_quantize.py +++ b/src/xorl/ops/quantize/int4_gkn_quantize.py @@ -11,15 +11,18 @@ from typing import Tuple import torch -from torch import Tensor import triton import triton.language as tl +from torch import Tensor @triton.jit def _int4_quantize_gkn_kernel( - X, Out, Scale, - K, N, + X, + Out, + Scale, + K, + N, GROUP_SIZE: tl.constexpr, TILE_N: tl.constexpr, ): @@ -74,8 +77,11 @@ def _int4_quantize_gkn_kernel( @triton.jit def _int4_dequantize_gkn_kernel( - Packed, Scale, Out, - K, N, + Packed, + Scale, + Out, + K, + N, GROUP_SIZE: tl.constexpr, TILE_N: tl.constexpr, ): @@ -105,7 +111,7 @@ def _int4_dequantize_gkn_kernel( packed = tl.load(pack_addrs, mask=pack_mask, other=0).to(tl.int32) # [GS//2, TN] # Unpack lo/hi nibbles - lo = (packed & 0xF).to(tl.float32) # even K values [GS//2, TN] + lo = (packed & 0xF).to(tl.float32) # even K values [GS//2, TN] hi = ((packed >> 4) & 0xF).to(tl.float32) # odd K values [GS//2, TN] # Dequantize: (nibble - 8) * scale → bf16 @@ -155,17 +161,21 @@ def int4_quantize_gkn(x: Tensor, group_size: int = -1) -> Tuple[Tensor, Tensor]: grid = (num_groups, (N + TILE_N - 1) // TILE_N) _int4_quantize_gkn_kernel[grid]( - x.contiguous(), packed, scales, - K, N, - gs, TILE_N, + x.contiguous(), + packed, + scales, + K, + N, + gs, + TILE_N, num_warps=4, ) return packed, scales -def int4_dequantize_gkn(packed: Tensor, scales: Tensor, K: int, N: int, - group_size: int = -1, - out_dtype: torch.dtype = torch.bfloat16) -> Tensor: +def int4_dequantize_gkn( + packed: Tensor, scales: Tensor, K: int, N: int, group_size: int = -1, out_dtype: torch.dtype = torch.bfloat16 +) -> Tensor: """Dequantize INT4 packed [K//2, N] back to [K, N] in G,K,N format. Args: @@ -191,9 +201,13 @@ def int4_dequantize_gkn(packed: Tensor, scales: Tensor, K: int, N: int, scales_f32 = scales.float() if scales.dtype != torch.float32 else scales _int4_dequantize_gkn_kernel[grid]( - packed.contiguous(), scales_f32.contiguous(), result, - K, N, - gs, TILE_N, + packed.contiguous(), + scales_f32.contiguous(), + result, + K, + N, + gs, + TILE_N, num_warps=4, ) diff --git a/src/xorl/ops/quantize/int4_quantize.py b/src/xorl/ops/quantize/int4_quantize.py index 4dbad3ca..eb62edc3 100644 --- a/src/xorl/ops/quantize/int4_quantize.py +++ b/src/xorl/ops/quantize/int4_quantize.py @@ -1,15 +1,18 @@ """INT4 quantization/dequantization via Triton kernels.""" + from typing import Tuple import torch -from torch import Tensor import triton import triton.language as tl +from torch import Tensor @triton.jit def _int4_quantize_kernel( - X, Out, Scale, + X, + Out, + Scale, total_elems, GROUP_SIZE: tl.constexpr, TILE_ELEMS: tl.constexpr, @@ -26,8 +29,7 @@ def _int4_quantize_kernel( amax = tl.max(tl.abs(x_2d), axis=1) scale = tl.maximum(amax, 1e-12) * (2.0 / 15.0) scale_offs = tl.arange(0, groups_per_tile) - tl.store(Scale + base_group + scale_offs, scale, - mask=(base_group + scale_offs) < (total_elems // GROUP_SIZE)) + tl.store(Scale + base_group + scale_offs, scale, mask=(base_group + scale_offs) < (total_elems // GROUP_SIZE)) inv_scale_2d = 1.0 / tl.expand_dims(scale, axis=1) q_2d = tl.maximum(0, tl.minimum(15, (x_2d * inv_scale_2d + 8.5).to(tl.int32))) q = tl.reshape(q_2d, (TILE_ELEMS,)) @@ -36,13 +38,14 @@ def _int4_quantize_kernel( packed = tl.sum(tl.reshape(shifted, (half_tile, 2)), axis=1).to(tl.uint8) out_base = pid * half_tile out_offs = tl.arange(0, half_tile) - tl.store(Out + out_base + out_offs, packed, - mask=(out_base + out_offs) < (total_elems // 2)) + tl.store(Out + out_base + out_offs, packed, mask=(out_base + out_offs) < (total_elems // 2)) @triton.jit def _int4_dequantize_kernel( - Packed, Scale, Out, + Packed, + Scale, + Out, total_elems, GROUP_SIZE: tl.constexpr, TILE_ELEMS: tl.constexpr, @@ -59,9 +62,9 @@ def _int4_dequantize_kernel( lo = packed & 0xF hi = (packed >> 4) & 0xF scale_offs = tl.arange(0, groups_per_tile) - scale = tl.load(Scale + base_group + scale_offs, - mask=(base_group + scale_offs) < (total_elems // GROUP_SIZE), - other=1.0).to(tl.float32) + scale = tl.load( + Scale + base_group + scale_offs, mask=(base_group + scale_offs) < (total_elems // GROUP_SIZE), other=1.0 + ).to(tl.float32) scale_2d = tl.expand_dims(scale, axis=1) lo_2d = tl.reshape(lo.to(tl.float32), (groups_per_tile, half_gs)) hi_2d = tl.reshape(hi.to(tl.float32), (groups_per_tile, half_gs)) @@ -107,8 +110,9 @@ def int4_quantize(x: Tensor, group_size: int = -1) -> Tuple[Tensor, Tensor]: return packed.reshape(M, K // 2), scales.reshape(M, num_groups) -def int4_dequantize(packed: Tensor, scales: Tensor, M: int, K: int, - group_size: int = -1, out_dtype: torch.dtype = torch.bfloat16) -> Tensor: +def int4_dequantize( + packed: Tensor, scales: Tensor, M: int, K: int, group_size: int = -1, out_dtype: torch.dtype = torch.bfloat16 +) -> Tensor: gs = K if group_size == -1 else group_size n = M * K packed_flat = packed.reshape(-1) diff --git a/src/xorl/ops/quantize/mxfp4_gkn_quantize.py b/src/xorl/ops/quantize/mxfp4_gkn_quantize.py index 4b1cd76c..2cc21229 100644 --- a/src/xorl/ops/quantize/mxfp4_gkn_quantize.py +++ b/src/xorl/ops/quantize/mxfp4_gkn_quantize.py @@ -11,17 +11,20 @@ from typing import Tuple import torch -from torch import Tensor import triton import triton.language as tl +from torch import Tensor -from .fp4_codec import _fp4_encode, _fp4_decode +from .fp4_codec import _fp4_decode, _fp4_encode @triton.jit def _mxfp4_quantize_gkn_kernel( - X, Out, Scale, - K, N, + X, + Out, + Scale, + K, + N, BLOCK_SIZE: tl.constexpr, TILE_N: tl.constexpr, ): @@ -81,8 +84,11 @@ def _mxfp4_quantize_gkn_kernel( @triton.jit def _mxfp4_dequantize_gkn_kernel( - Packed, Scale, Out, - K, N, + Packed, + Scale, + Out, + K, + N, BLOCK_SIZE: tl.constexpr, TILE_N: tl.constexpr, ): @@ -112,7 +118,7 @@ def _mxfp4_dequantize_gkn_kernel( packed = tl.load(pack_addrs, mask=pack_mask, other=0).to(tl.int32) # [BS//2, TN] # Unpack and decode FP4 (element-wise, works directly on 2D) - lo = packed & 0xF # even K [BS//2, TN] + lo = packed & 0xF # even K [BS//2, TN] hi = (packed >> 4) & 0xF # odd K [BS//2, TN] val_lo = _fp4_decode(lo) val_hi = _fp4_decode(hi) @@ -163,16 +169,19 @@ def mxfp4_quantize_gkn(x: Tensor, block_size: int = 32) -> Tuple[Tensor, Tensor] grid = (num_blocks, (N + TILE_N - 1) // TILE_N) _mxfp4_quantize_gkn_kernel[grid]( - x.contiguous(), packed, scales, - K, N, - block_size, TILE_N, + x.contiguous(), + packed, + scales, + K, + N, + block_size, + TILE_N, num_warps=4, ) return packed, scales.to(torch.float16) -def mxfp4_dequantize_gkn(packed: Tensor, scales: Tensor, K: int, N: int, - block_size: int = 32) -> Tensor: +def mxfp4_dequantize_gkn(packed: Tensor, scales: Tensor, K: int, N: int, block_size: int = 32) -> Tensor: """Dequantize MXFP4 packed [K//2, N] back to [K, N] in G,K,N format. Args: @@ -192,9 +201,13 @@ def mxfp4_dequantize_gkn(packed: Tensor, scales: Tensor, K: int, N: int, grid = (num_blocks, (N + TILE_N - 1) // TILE_N) _mxfp4_dequantize_gkn_kernel[grid]( - packed.contiguous(), scales.float().contiguous(), result, - K, N, - block_size, TILE_N, + packed.contiguous(), + scales.float().contiguous(), + result, + K, + N, + block_size, + TILE_N, num_warps=2, ) return result diff --git a/src/xorl/ops/quantize/mxfp4_quantize.py b/src/xorl/ops/quantize/mxfp4_quantize.py index 12587e00..c3bb8a6b 100644 --- a/src/xorl/ops/quantize/mxfp4_quantize.py +++ b/src/xorl/ops/quantize/mxfp4_quantize.py @@ -1,12 +1,13 @@ """MXFP4 quantization/dequantization via Triton kernels.""" + from typing import Tuple import torch -from torch import Tensor import triton import triton.language as tl +from torch import Tensor -from .fp4_codec import _fp4_encode, _fp4_decode +from .fp4_codec import _fp4_decode, _fp4_encode _QUANT_CONFIGS = { @@ -28,7 +29,9 @@ @triton.jit def _mxfp4_quantize_kernel( - X, Out, Scale, + X, + Out, + Scale, total_elems, BLOCK_SIZE: tl.constexpr, TILE_ELEMS: tl.constexpr, @@ -47,8 +50,7 @@ def _mxfp4_quantize_kernel( # Bitwise E8M0 rounding: multiply by sqrt(2) then zero mantissa bits scale = ((scale * 1.4142135623730951).to(tl.int32, bitcast=True) & 0x7F800000).to(tl.float32, bitcast=True) scale_offs = tl.arange(0, blocks_per_tile) - tl.store(Scale + base_block + scale_offs, scale, - mask=(base_block + scale_offs) < (total_elems // BLOCK_SIZE)) + tl.store(Scale + base_block + scale_offs, scale, mask=(base_block + scale_offs) < (total_elems // BLOCK_SIZE)) inv_scale_2d = 1.0 / tl.expand_dims(scale, axis=1) scaled_2d = x_2d * inv_scale_2d scaled = tl.reshape(scaled_2d, (TILE_ELEMS,)) @@ -58,13 +60,14 @@ def _mxfp4_quantize_kernel( packed = tl.sum(tl.reshape(shifted, (half_tile, 2)), axis=1).to(tl.uint8) out_base = pid * half_tile out_offs = tl.arange(0, half_tile) - tl.store(Out + out_base + out_offs, packed, - mask=(out_base + out_offs) < (total_elems // 2)) + tl.store(Out + out_base + out_offs, packed, mask=(out_base + out_offs) < (total_elems // 2)) @triton.jit def _mxfp4_dequantize_kernel( - Packed, Scale, Out, + Packed, + Scale, + Out, total_elems, BLOCK_SIZE: tl.constexpr, TILE_ELEMS: tl.constexpr, @@ -76,8 +79,7 @@ def _mxfp4_dequantize_kernel( base_block = pid * blocks_per_tile scale_offs = tl.arange(0, blocks_per_tile) scale_mask = (base_block + scale_offs) < (total_elems // BLOCK_SIZE) - scale_vec = tl.load(Scale + base_block + scale_offs, - mask=scale_mask, other=1.0).to(tl.float32) + scale_vec = tl.load(Scale + base_block + scale_offs, mask=scale_mask, other=1.0).to(tl.float32) pack_base = pid * half_tile pack_offs = tl.arange(0, half_tile) pack_mask = (pack_base + pack_offs) < (total_elems // 2) @@ -120,14 +122,19 @@ def mxfp4_quantize(x: Tensor, block_size: int = 32) -> Tuple[Tensor, Tensor]: return packed, scales -def mxfp4_dequantize(packed: Tensor, scales: Tensor, num_elements: int, - block_size: int = 32) -> Tensor: +def mxfp4_dequantize(packed: Tensor, scales: Tensor, num_elements: int, block_size: int = 32) -> Tensor: assert num_elements % block_size == 0 packed_flat = packed.reshape(-1) out_i32 = torch.empty(num_elements // 2, dtype=torch.int32, device=packed_flat.device) te, nw = _get_config(block_size, num_elements, is_dequant=True) grid = (num_elements + te - 1) // te _mxfp4_dequantize_kernel[(grid,)]( - packed_flat, scales, out_i32, num_elements, block_size, te, num_warps=nw, + packed_flat, + scales, + out_i32, + num_elements, + block_size, + te, + num_warps=nw, ) return out_i32.view(torch.bfloat16) diff --git a/src/xorl/ops/quantize/nf4_codec.py b/src/xorl/ops/quantize/nf4_codec.py index 45f3ed69..64f6d7e3 100644 --- a/src/xorl/ops/quantize/nf4_codec.py +++ b/src/xorl/ops/quantize/nf4_codec.py @@ -8,10 +8,12 @@ Encoding: comparison chain against midpoints (15 comparisons, branchless). Decoding: LUT gather (pass pointer, tl.load with code as offset, L1-cached). """ + import torch import triton import triton.language as tl + # NF4 quantization levels (sorted, normalized to [-1, 1]) NF4_TABLE = [ -1.0, @@ -79,7 +81,5 @@ def get_nf4_lut(device) -> torch.Tensor: """Get (or create) the NF4 decode LUT on the given device.""" key = str(device) if key not in _NF4_LUT_CACHE: - _NF4_LUT_CACHE[key] = torch.tensor( - NF4_TABLE, dtype=torch.float32, device=device - ) + _NF4_LUT_CACHE[key] = torch.tensor(NF4_TABLE, dtype=torch.float32, device=device) return _NF4_LUT_CACHE[key] diff --git a/src/xorl/ops/quantize/nf4_gkn_quantize.py b/src/xorl/ops/quantize/nf4_gkn_quantize.py index 516277be..a5d7da0d 100644 --- a/src/xorl/ops/quantize/nf4_gkn_quantize.py +++ b/src/xorl/ops/quantize/nf4_gkn_quantize.py @@ -11,17 +11,20 @@ from typing import Tuple import torch -from torch import Tensor import triton import triton.language as tl +from torch import Tensor -from .nf4_codec import _nf4_encode, _nf4_decode, get_nf4_lut +from .nf4_codec import _nf4_decode, _nf4_encode, get_nf4_lut @triton.jit def _nf4_quantize_gkn_kernel( - X, Out, Scale, - K, N, + X, + Out, + Scale, + K, + N, GROUP_SIZE: tl.constexpr, TILE_N: tl.constexpr, ): @@ -79,8 +82,12 @@ def _nf4_quantize_gkn_kernel( @triton.jit def _nf4_dequantize_gkn_kernel( - Packed, Scale, LUT, Out, - K, N, + Packed, + Scale, + LUT, + Out, + K, + N, GROUP_SIZE: tl.constexpr, TILE_N: tl.constexpr, ): @@ -160,16 +167,19 @@ def nf4_quantize_gkn(x: Tensor, group_size: int = 64) -> Tuple[Tensor, Tensor]: grid = (num_groups, (N + TILE_N - 1) // TILE_N) _nf4_quantize_gkn_kernel[grid]( - x.contiguous(), packed, scales, - K, N, - group_size, TILE_N, + x.contiguous(), + packed, + scales, + K, + N, + group_size, + TILE_N, num_warps=4, ) return packed, scales -def nf4_dequantize_gkn(packed: Tensor, scales: Tensor, - K: int, N: int, group_size: int = 64) -> Tensor: +def nf4_dequantize_gkn(packed: Tensor, scales: Tensor, K: int, N: int, group_size: int = 64) -> Tensor: """Dequantize NF4 packed [K//2, N] back to [K, N] in G,K,N format. Args: @@ -192,9 +202,14 @@ def nf4_dequantize_gkn(packed: Tensor, scales: Tensor, lut = get_nf4_lut(packed.device) _nf4_dequantize_gkn_kernel[grid]( - packed.contiguous(), scales_f32.contiguous(), lut, result, - K, N, - group_size, TILE_N, + packed.contiguous(), + scales_f32.contiguous(), + lut, + result, + K, + N, + group_size, + TILE_N, num_warps=4, ) return result diff --git a/src/xorl/ops/quantize/nf4_quantize.py b/src/xorl/ops/quantize/nf4_quantize.py index 6fe31800..2c4ad468 100644 --- a/src/xorl/ops/quantize/nf4_quantize.py +++ b/src/xorl/ops/quantize/nf4_quantize.py @@ -6,19 +6,22 @@ Storage: packed uint8 (2 nf4 codes per byte) + float32 scales per group. """ + from typing import Tuple import torch -from torch import Tensor import triton import triton.language as tl +from torch import Tensor -from .nf4_codec import _nf4_encode, _nf4_decode, get_nf4_lut +from .nf4_codec import _nf4_decode, _nf4_encode, get_nf4_lut @triton.jit def _nf4_quantize_kernel( - X, Out, Scale, + X, + Out, + Scale, total_elems, GROUP_SIZE: tl.constexpr, TILE_ELEMS: tl.constexpr, @@ -38,8 +41,7 @@ def _nf4_quantize_kernel( scale = tl.maximum(amax, 1e-12) # Store scales scale_offs = tl.arange(0, groups_per_tile) - tl.store(Scale + base_group + scale_offs, scale, - mask=(base_group + scale_offs) < (total_elems // GROUP_SIZE)) + tl.store(Scale + base_group + scale_offs, scale, mask=(base_group + scale_offs) < (total_elems // GROUP_SIZE)) # Normalize to [-1, 1] and encode inv_scale_2d = 1.0 / tl.expand_dims(scale, axis=1) norm_2d = x_2d * inv_scale_2d @@ -52,13 +54,15 @@ def _nf4_quantize_kernel( packed = tl.sum(tl.reshape(shifted, (half_tile, 2)), axis=1).to(tl.uint8) out_base = pid * half_tile out_offs = tl.arange(0, half_tile) - tl.store(Out + out_base + out_offs, packed, - mask=(out_base + out_offs) < (total_elems // 2)) + tl.store(Out + out_base + out_offs, packed, mask=(out_base + out_offs) < (total_elems // 2)) @triton.jit def _nf4_dequantize_kernel( - Packed, Scale, LUT, Out, + Packed, + Scale, + LUT, + Out, total_elems, GROUP_SIZE: tl.constexpr, TILE_ELEMS: tl.constexpr, @@ -74,9 +78,9 @@ def _nf4_dequantize_kernel( base_group = pid * groups_per_tile # Load scales scale_offs = tl.arange(0, groups_per_tile) - scale = tl.load(Scale + base_group + scale_offs, - mask=(base_group + scale_offs) < (total_elems // GROUP_SIZE), - other=1.0).to(tl.float32) + scale = tl.load( + Scale + base_group + scale_offs, mask=(base_group + scale_offs) < (total_elems // GROUP_SIZE), other=1.0 + ).to(tl.float32) # Load packed bytes pack_base = pid * half_tile pack_offs = tl.arange(0, half_tile) @@ -138,14 +142,18 @@ def nf4_quantize(x: Tensor, group_size: int = 64) -> Tuple[Tensor, Tensor]: te, nw = _get_config(group_size, n) grid = (n + te - 1) // te _nf4_quantize_kernel[(grid,)]( - x_flat, packed, scales, n, - group_size, te, num_warps=nw, + x_flat, + packed, + scales, + n, + group_size, + te, + num_warps=nw, ) return packed, scales -def nf4_dequantize(packed: Tensor, scales: Tensor, num_elements: int, - group_size: int = 64) -> Tensor: +def nf4_dequantize(packed: Tensor, scales: Tensor, num_elements: int, group_size: int = 64) -> Tensor: """Dequantize NF4 packed data back to bfloat16 (flat output). Args: @@ -165,7 +173,13 @@ def nf4_dequantize(packed: Tensor, scales: Tensor, num_elements: int, te, nw = _get_config(group_size, num_elements, is_dequant=True) grid = (num_elements + te - 1) // te _nf4_dequantize_kernel[(grid,)]( - packed_flat, scales_flat, lut, out_i32, num_elements, - group_size, te, num_warps=nw, + packed_flat, + scales_flat, + lut, + out_i32, + num_elements, + group_size, + te, + num_warps=nw, ) return out_i32.view(torch.bfloat16) diff --git a/src/xorl/ops/quantize/nvfp4_gkn_quantize.py b/src/xorl/ops/quantize/nvfp4_gkn_quantize.py index b3bbf81f..85026f74 100644 --- a/src/xorl/ops/quantize/nvfp4_gkn_quantize.py +++ b/src/xorl/ops/quantize/nvfp4_gkn_quantize.py @@ -12,17 +12,21 @@ from typing import Optional, Tuple import torch -from torch import Tensor import triton import triton.language as tl +from torch import Tensor -from .fp4_codec import FP4_E2M1_MAX, FP8_E4M3_MAX, _fp4_encode, _fp4_decode +from .fp4_codec import FP4_E2M1_MAX, FP8_E4M3_MAX, _fp4_decode, _fp4_encode @triton.jit def _nvfp4_quantize_gkn_kernel( - X, Out, Amax, GlobalAmax, - K, N, + X, + Out, + Amax, + GlobalAmax, + K, + N, BLOCK_SIZE: tl.constexpr, TILE_N: tl.constexpr, ): @@ -85,8 +89,12 @@ def _nvfp4_quantize_gkn_kernel( @triton.jit def _nvfp4_scale_convert_gkn_kernel( - Amax, ScaleOut, GlobalAmaxInOut, - n_scales, FP4_MAX_INV, FP8_MAX, + Amax, + ScaleOut, + GlobalAmaxInOut, + n_scales, + FP4_MAX_INV, + FP8_MAX, BLOCK: tl.constexpr, ): """Convert raw amax to fp8 block_scales (operates on flat amax array).""" @@ -102,8 +110,12 @@ def _nvfp4_scale_convert_gkn_kernel( @triton.jit def _nvfp4_dequantize_gkn_kernel( - Packed, Scale, GlobalScale, Out, - K, N, + Packed, + Scale, + GlobalScale, + Out, + K, + N, BLOCK_SIZE: tl.constexpr, TILE_N: tl.constexpr, ): @@ -161,8 +173,9 @@ def _next_pow2(n): return 1 << (n - 1).bit_length() -def nvfp4_quantize_gkn(x: Tensor, block_size: int = 16, - global_amax: Optional[Tensor] = None) -> Tuple[Tensor, Tensor, Tensor]: +def nvfp4_quantize_gkn( + x: Tensor, block_size: int = 16, global_amax: Optional[Tensor] = None +) -> Tuple[Tensor, Tensor, Tensor]: """Quantize a [K, N] weight tensor in G,K,N format to NVFP4. Groups are formed along the K (first) dimension. @@ -190,9 +203,14 @@ def nvfp4_quantize_gkn(x: Tensor, block_size: int = 16, grid = (num_blocks, (N + TILE_N - 1) // TILE_N) _nvfp4_quantize_gkn_kernel[grid]( - x.contiguous(), packed, amax, global_amax_buf, - K, N, - block_size, TILE_N, + x.contiguous(), + packed, + amax, + global_amax_buf, + K, + N, + block_size, + TILE_N, num_warps=4, ) @@ -205,15 +223,21 @@ def nvfp4_quantize_gkn(x: Tensor, block_size: int = 16, block_scales = torch.empty(n_scales, dtype=torch.float8_e4m3fn, device=x.device) SC_BLOCK = 1024 _nvfp4_scale_convert_gkn_kernel[((n_scales + SC_BLOCK - 1) // SC_BLOCK,)]( - amax.reshape(-1), block_scales, global_amax_buf, n_scales, - 1.0 / FP4_E2M1_MAX, FP8_E4M3_MAX, SC_BLOCK, + amax.reshape(-1), + block_scales, + global_amax_buf, + n_scales, + 1.0 / FP4_E2M1_MAX, + FP8_E4M3_MAX, + SC_BLOCK, num_warps=4, ) return packed, block_scales.reshape(num_blocks, N), global_amax_buf -def nvfp4_dequantize_gkn(packed: Tensor, block_scales: Tensor, global_scale: Tensor, - K: int, N: int, block_size: int = 16) -> Tensor: +def nvfp4_dequantize_gkn( + packed: Tensor, block_scales: Tensor, global_scale: Tensor, K: int, N: int, block_size: int = 16 +) -> Tensor: """Dequantize NVFP4 packed [K//2, N] back to [K, N] in G,K,N format. Args: @@ -236,9 +260,14 @@ def nvfp4_dequantize_gkn(packed: Tensor, block_scales: Tensor, global_scale: Ten gs = global_scale if global_scale.numel() == 1 else global_scale.reshape(1) _nvfp4_dequantize_gkn_kernel[grid]( - packed.contiguous(), block_scales.contiguous(), gs, result, - K, N, - block_size, TILE_N, + packed.contiguous(), + block_scales.contiguous(), + gs, + result, + K, + N, + block_size, + TILE_N, num_warps=1, ) return result diff --git a/src/xorl/ops/quantize/nvfp4_quantize.py b/src/xorl/ops/quantize/nvfp4_quantize.py index 8c97856b..c83eac59 100644 --- a/src/xorl/ops/quantize/nvfp4_quantize.py +++ b/src/xorl/ops/quantize/nvfp4_quantize.py @@ -1,12 +1,13 @@ """NVFP4 quantization/dequantization via Triton kernels.""" + from typing import Optional, Tuple import torch -from torch import Tensor import triton import triton.language as tl +from torch import Tensor -from .fp4_codec import FP4_E2M1_MAX, FP8_E4M3_MAX, _fp4_encode, _fp4_decode +from .fp4_codec import FP4_E2M1_MAX, FP8_E4M3_MAX, _fp4_decode, _fp4_encode _QUANT_CONFIGS = { @@ -22,7 +23,10 @@ @triton.jit def _nvfp4_quantize_kernel( - X, Out, Amax, GlobalAmax, + X, + Out, + Amax, + GlobalAmax, total_elems, BLOCK_SIZE: tl.constexpr, TILE_ELEMS: tl.constexpr, @@ -40,27 +44,28 @@ def _nvfp4_quantize_kernel( amax = tl.max(tl.abs(x_2d), axis=1) amax = tl.maximum(amax, 1e-12) # Compute codes before stores to overlap compute with memory - inv_amax_2d = (6.0 / tl.expand_dims(amax, axis=1)) + inv_amax_2d = 6.0 / tl.expand_dims(amax, axis=1) scaled_2d = x_2d * inv_amax_2d scaled_2d = tl.minimum(tl.maximum(scaled_2d, -6.0), 6.0) scaled = tl.reshape(scaled_2d, (TILE_ELEMS,)) codes = _fp4_encode(scaled) scale_offs = tl.arange(0, blocks_per_tile) - tl.store(Amax + base_block + scale_offs, amax, - mask=(base_block + scale_offs) < (total_elems // BLOCK_SIZE)) + tl.store(Amax + base_block + scale_offs, amax, mask=(base_block + scale_offs) < (total_elems // BLOCK_SIZE)) tl.atomic_max(GlobalAmax, tl.max(amax)) is_odd = (offs & 1).to(tl.int32) shifted = (codes & 0xF) << (is_odd * 4) packed = tl.sum(tl.reshape(shifted, (half_tile, 2)), axis=1).to(tl.uint8) out_base = pid * half_tile out_offs = tl.arange(0, half_tile) - tl.store(Out + out_base + out_offs, packed, - mask=(out_base + out_offs) < (total_elems // 2)) + tl.store(Out + out_base + out_offs, packed, mask=(out_base + out_offs) < (total_elems // 2)) @triton.jit def _nvfp4_dequantize_kernel( - Packed, Scale, GlobalScale, Out, + Packed, + Scale, + GlobalScale, + Out, total_elems, BLOCK_SIZE: tl.constexpr, TILE_ELEMS: tl.constexpr, @@ -73,8 +78,7 @@ def _nvfp4_dequantize_kernel( gs = tl.load(GlobalScale).to(tl.float32) scale_offs = tl.arange(0, blocks_per_tile) scale_mask = (base_block + scale_offs) < (total_elems // BLOCK_SIZE) - scale_vec = tl.load(Scale + base_block + scale_offs, - mask=scale_mask, other=1.0).to(tl.float32) * gs + scale_vec = tl.load(Scale + base_block + scale_offs, mask=scale_mask, other=1.0).to(tl.float32) * gs pack_base = pid * half_tile pack_offs = tl.arange(0, half_tile) pack_mask = (pack_base + pack_offs) < (total_elems // 2) @@ -95,9 +99,7 @@ def _nvfp4_dequantize_kernel( @triton.jit -def _nvfp4_scale_convert_kernel(Amax, ScaleOut, GlobalAmaxInOut, n_blocks, - FP4_MAX_INV, FP8_MAX, - BLOCK: tl.constexpr): +def _nvfp4_scale_convert_kernel(Amax, ScaleOut, GlobalAmaxInOut, n_blocks, FP4_MAX_INV, FP8_MAX, BLOCK: tl.constexpr): """Convert raw amax to fp8 block_scales and update GlobalAmaxInOut to global_scale.""" pid = tl.program_id(0) offs = pid * BLOCK + tl.arange(0, BLOCK) @@ -109,8 +111,9 @@ def _nvfp4_scale_convert_kernel(Amax, ScaleOut, GlobalAmaxInOut, n_blocks, tl.store(GlobalAmaxInOut, gamax * FP4_MAX_INV / FP8_MAX) -def nvfp4_quantize(x: Tensor, block_size: int = 16, - global_amax: Optional[Tensor] = None) -> Tuple[Tensor, Tensor, Tensor]: +def nvfp4_quantize( + x: Tensor, block_size: int = 16, global_amax: Optional[Tensor] = None +) -> Tuple[Tensor, Tensor, Tensor]: assert x.dim() == 2 M, K = x.shape n = M * K @@ -124,8 +127,14 @@ def nvfp4_quantize(x: Tensor, block_size: int = 16, te = min(te, n) grid = (n + te - 1) // te _nvfp4_quantize_kernel[(grid,)]( - x_flat, packed, amax, global_amax_buf, n, - block_size, te, num_warps=nw, + x_flat, + packed, + amax, + global_amax_buf, + n, + block_size, + te, + num_warps=nw, ) # Override global amax with EMA-tracked value if provided if global_amax is not None: @@ -133,14 +142,21 @@ def nvfp4_quantize(x: Tensor, block_size: int = 16, block_scales = torch.empty(n_scales, dtype=torch.float8_e4m3fn, device=x.device) SC_BLOCK = 1024 _nvfp4_scale_convert_kernel[((n_scales + SC_BLOCK - 1) // SC_BLOCK,)]( - amax, block_scales, global_amax_buf, n_scales, - 1.0 / FP4_E2M1_MAX, FP8_E4M3_MAX, SC_BLOCK, num_warps=4, + amax, + block_scales, + global_amax_buf, + n_scales, + 1.0 / FP4_E2M1_MAX, + FP8_E4M3_MAX, + SC_BLOCK, + num_warps=4, ) return packed, block_scales, global_amax_buf -def nvfp4_dequantize(packed: Tensor, block_scales: Tensor, global_scale: Tensor, - num_elements: int, block_size: int = 16) -> Tensor: +def nvfp4_dequantize( + packed: Tensor, block_scales: Tensor, global_scale: Tensor, num_elements: int, block_size: int = 16 +) -> Tensor: assert num_elements % block_size == 0 packed_flat = packed.reshape(-1) out_i32 = torch.empty(num_elements // 2, dtype=torch.int32, device=packed_flat.device) @@ -149,7 +165,13 @@ def nvfp4_dequantize(packed: Tensor, block_scales: Tensor, global_scale: Tensor, grid = (num_elements + te - 1) // te gs = global_scale if global_scale.numel() == 1 else global_scale.reshape(1) _nvfp4_dequantize_kernel[(grid,)]( - packed_flat, block_scales, gs, out_i32, num_elements, - block_size, te, num_warps=nw, + packed_flat, + block_scales, + gs, + out_i32, + num_elements, + block_size, + te, + num_warps=nw, ) return out_i32.view(torch.bfloat16) diff --git a/src/xorl/optim/muon.py b/src/xorl/optim/muon.py index f58c4f47..c318795d 100644 --- a/src/xorl/optim/muon.py +++ b/src/xorl/optim/muon.py @@ -32,6 +32,7 @@ from ..utils import logging + logger = logging.get_logger(__name__) @@ -176,7 +177,8 @@ def _muon_step(self, group: dict) -> None: if "momentum_buffer" not in state: buf_dtype = self._momentum_dtype or grad_local.dtype state["momentum_buffer"] = torch.zeros_like( - grad_local, dtype=buf_dtype, + grad_local, + dtype=buf_dtype, ) if not self._logged_dtypes: logger.info_rank0( diff --git a/src/xorl/optim/optimizer.py b/src/xorl/optim/optimizer.py index 86a7642d..857c1ad0 100644 --- a/src/xorl/optim/optimizer.py +++ b/src/xorl/optim/optimizer.py @@ -193,7 +193,9 @@ def state_dict( f"Key clash detected while merging state dict for optimizer '{name}': {', '.join(sorted(overlap))}" ) else: - logger.info_rank0(f"MultiOptimizer merged '{name}' state dict ({len(sd)} keys, total {len(merged) + len(sd)})") + logger.info_rank0( + f"MultiOptimizer merged '{name}' state dict ({len(sd)} keys, total {len(merged) + len(sd)})" + ) merged.update(sd) return merged @@ -304,16 +306,26 @@ def _get_optimizer_cls_and_kwargs( if optimizer_type == "adamw": foreach = not fused ctor_kwargs = dict( - lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, - fused=fused, foreach=foreach, **kwargs, + lr=lr, + betas=betas, + eps=eps, + weight_decay=weight_decay, + fused=fused, + foreach=foreach, + **kwargs, ) return AdamW, ctor_kwargs elif optimizer_type == "anyprecision_adamw": state_dtype = _ANYPRECISION_STATE_DTYPES[optimizer_dtype] ctor_kwargs = dict( - lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, - momentum_dtype=state_dtype, variance_dtype=state_dtype, - compensation_buffer_dtype=state_dtype, **kwargs, + lr=lr, + betas=betas, + eps=eps, + weight_decay=weight_decay, + momentum_dtype=state_dtype, + variance_dtype=state_dtype, + compensation_buffer_dtype=state_dtype, + **kwargs, ) return AnyPrecisionAdamW, ctor_kwargs elif optimizer_type == "sgd": @@ -338,8 +350,7 @@ def _get_optimizer_cls_and_kwargs( return Muon, ctor_kwargs else: raise ValueError( - f"Unsupported optimizer type: '{optimizer_type}'. " - f"Supported: adamw, anyprecision_adamw, sgd, muon." + f"Unsupported optimizer type: '{optimizer_type}'. Supported: adamw, anyprecision_adamw, sgd, muon." ) @@ -365,8 +376,14 @@ def _create_optimizer( - adamw/anyprecision_adamw: any extra kwargs forwarded to constructor """ cls, ctor_kwargs = _get_optimizer_cls_and_kwargs( - optimizer_type, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, - fused=fused, optimizer_dtype=optimizer_dtype, optimizer_kwargs=optimizer_kwargs, + optimizer_type, + lr=lr, + betas=betas, + eps=eps, + weight_decay=weight_decay, + fused=fused, + optimizer_dtype=optimizer_dtype, + optimizer_kwargs=optimizer_kwargs, ) return cls(param_groups, **ctor_kwargs) @@ -489,8 +506,17 @@ def build_optimizer( # EP-aware routing: for FSDP2+EP, split params into EP and non-EP groups and build two optimizers. if _should_build_ep_aware(model, param_groups): return build_ep_fsdp2_optimizer( - model, lr, betas, eps, weight_decay, fused, optimizer_type, optimizer_dtype, - param_groups, no_decay_modules, no_decay_params, + model, + lr, + betas, + eps, + weight_decay, + fused, + optimizer_type, + optimizer_dtype, + param_groups, + no_decay_modules, + no_decay_params, optimizer_kwargs=optimizer_kwargs, ) @@ -499,17 +525,20 @@ def build_optimizer( # Muon optimizer: split params into Muon (2D+ matrices) and AdamW (rest) if optimizer_type == "muon": muon_params, muon_names, adamw_params, adamw_names = _classify_muon_params(model) - logger.info_rank0( - f"Muon optimizer: {len(muon_params)} Muon params, {len(adamw_params)} AdamW params" - ) + logger.info_rank0(f"Muon optimizer: {len(muon_params)} Muon params, {len(adamw_params)} AdamW params") logger.info_rank0(f"Muon params: {muon_names}") logger.info_rank0(f"AdamW params: {adamw_names}") muon_lr = kwargs.get("muon_lr", 0.02) param_groups = _make_muon_param_groups( - model, muon_params, adamw_params, - muon_lr=muon_lr, adamw_lr=lr, weight_decay=weight_decay, - no_decay_modules=no_decay_modules, no_decay_params=no_decay_params, + model, + muon_params, + adamw_params, + muon_lr=muon_lr, + adamw_lr=lr, + weight_decay=weight_decay, + no_decay_modules=no_decay_modules, + no_decay_params=no_decay_params, ) elif param_groups is None: # Build param groups with weight decay splitting @@ -531,9 +560,14 @@ def build_optimizer( param_groups.append({"params": no_decay_parameters, "weight_decay": 0.0}) return _create_optimizer( - optimizer_type, param_groups, - lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, - fused=fused, optimizer_dtype=optimizer_dtype, + optimizer_type, + param_groups, + lr=lr, + betas=betas, + eps=eps, + weight_decay=weight_decay, + fused=fused, + optimizer_dtype=optimizer_dtype, optimizer_kwargs=optimizer_kwargs, ) @@ -598,14 +632,24 @@ def build_ep_fsdp2_optimizer( muon_lr = kwargs.get("muon_lr", 0.02) ep_groups = _make_muon_param_groups( - model, ep_muon, ep_adamw, - muon_lr=muon_lr, adamw_lr=lr, weight_decay=weight_decay, - no_decay_modules=no_decay_modules, no_decay_params=no_decay_params, + model, + ep_muon, + ep_adamw, + muon_lr=muon_lr, + adamw_lr=lr, + weight_decay=weight_decay, + no_decay_modules=no_decay_modules, + no_decay_params=no_decay_params, ) non_ep_groups = _make_muon_param_groups( - model, non_ep_muon, non_ep_adamw, - muon_lr=muon_lr, adamw_lr=lr, weight_decay=weight_decay, - no_decay_modules=no_decay_modules, no_decay_params=no_decay_params, + model, + non_ep_muon, + non_ep_adamw, + muon_lr=muon_lr, + adamw_lr=lr, + weight_decay=weight_decay, + no_decay_modules=no_decay_modules, + no_decay_params=no_decay_params, ) else: ep_groups = _make_param_groups_for_subset(model, ep_params, weight_decay, no_decay_modules, no_decay_params) @@ -615,9 +659,14 @@ def build_ep_fsdp2_optimizer( def _build(groups: Sequence[Dict[str, Any]]) -> Optimizer: return _create_optimizer( - optimizer_type, groups, - lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, - fused=fused, optimizer_dtype=optimizer_dtype, + optimizer_type, + groups, + lr=lr, + betas=betas, + eps=eps, + weight_decay=weight_decay, + fused=fused, + optimizer_dtype=optimizer_dtype, optimizer_kwargs=optimizer_kwargs, ) diff --git a/src/xorl/qlora/__init__.py b/src/xorl/qlora/__init__.py index ade69244..cf22245a 100644 --- a/src/xorl/qlora/__init__.py +++ b/src/xorl/qlora/__init__.py @@ -1,19 +1,21 @@ -from .modules.linear import QLoRALinear, prefetch_aqn_noise -from .modules.nvfp4_linear import NvFP4QLoRALinear +from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules + from .modules.block_fp8_linear import BlockFP8QLoRALinear -from .modules.nf4_linear import NF4QLoRALinear +from .modules.linear import QLoRALinear, prefetch_aqn_noise from .modules.moe_experts import QLoRAMoeExperts +from .modules.nf4_linear import NF4QLoRALinear +from .modules.nvfp4_linear import NvFP4QLoRALinear from .utils import ( + detect_prequantized_block_fp8, + detect_prequantized_nvfp4, inject_qlora_into_model, - save_qlora_checkpoint, - maybe_requant_qlora, - maybe_quantize_qlora, maybe_load_and_quantize_moe_qlora, - detect_prequantized_nvfp4, - detect_prequantized_block_fp8, maybe_load_prequantized_qlora, + maybe_quantize_qlora, + maybe_requant_qlora, + save_qlora_checkpoint, ) -from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules + __all__ = [ "QLoRALinear", diff --git a/src/xorl/qlora/modules/__init__.py b/src/xorl/qlora/modules/__init__.py index 110723f8..c5c88910 100644 --- a/src/xorl/qlora/modules/__init__.py +++ b/src/xorl/qlora/modules/__init__.py @@ -1,8 +1,9 @@ -from .linear import QLoRALinear, prefetch_aqn_noise -from .nvfp4_linear import NvFP4QLoRALinear from .block_fp8_linear import BlockFP8QLoRALinear -from .nf4_linear import NF4QLoRALinear +from .linear import QLoRALinear, prefetch_aqn_noise from .moe_experts import QLoRAMoeExperts +from .nf4_linear import NF4QLoRALinear +from .nvfp4_linear import NvFP4QLoRALinear + __all__ = [ "QLoRALinear", diff --git a/src/xorl/qlora/modules/block_fp8_linear.py b/src/xorl/qlora/modules/block_fp8_linear.py index cd3c4d08..77040254 100644 --- a/src/xorl/qlora/modules/block_fp8_linear.py +++ b/src/xorl/qlora/modules/block_fp8_linear.py @@ -7,7 +7,7 @@ import torch.nn as nn from torch import Tensor -from xorl.ops.quantize import block_fp8_quantize_gkn, block_fp8_dequantize_gkn +from xorl.ops.quantize import block_fp8_dequantize_gkn, block_fp8_quantize_gkn from xorl.qlora.modules.linear import QLoRALinear @@ -32,9 +32,16 @@ def __init__( aqn_alpha: float = 1.0, ): super().__init__( - in_features, out_features, r=r, lora_alpha=lora_alpha, - quant_format="block_fp8", quant_group_size=128, - bias=bias, device=device, enable_aqn=enable_aqn, aqn_alpha=aqn_alpha, + in_features, + out_features, + r=r, + lora_alpha=lora_alpha, + quant_format="block_fp8", + quant_group_size=128, + bias=bias, + device=device, + enable_aqn=enable_aqn, + aqn_alpha=aqn_alpha, ) # block_fp8: 1 fp8 byte per element -> in_features bytes -> in_features // 4 float32 elements pw_cols = in_features // 4 @@ -50,14 +57,25 @@ def __init__( self.reset_lora_parameters() @classmethod - def from_module(cls, module: nn.Module, r: int = 16, lora_alpha: int = 16, - enable_aqn: bool = False, aqn_alpha: float = 1.0, **kwargs) -> "BlockFP8QLoRALinear": + def from_module( + cls, + module: nn.Module, + r: int = 16, + lora_alpha: int = 16, + enable_aqn: bool = False, + aqn_alpha: float = 1.0, + **kwargs, + ) -> "BlockFP8QLoRALinear": """Create from a bf16 nn.Linear by quantizing its weight to block_fp8.""" qlora = cls( - in_features=module.in_features, out_features=module.out_features, - r=r, lora_alpha=lora_alpha, - bias=module.bias is not None, device=module.weight.device, - enable_aqn=enable_aqn, aqn_alpha=aqn_alpha, + in_features=module.in_features, + out_features=module.out_features, + r=r, + lora_alpha=lora_alpha, + bias=module.bias is not None, + device=module.weight.device, + enable_aqn=enable_aqn, + aqn_alpha=aqn_alpha, ) qlora._quantize_and_store(module.weight.detach()) return qlora @@ -75,9 +93,7 @@ def _dequantize_weight(self) -> Tensor: M, K = self.out_features, self.in_features uint8_data = self._read_packed_weight_uint8() fp8_w = uint8_data.view(torch.float8_e4m3fn).reshape(M, K) - scales = self._recover_tensor( - self.weight_block_scales, self._scale_dtypes["weight_block_scales"] - ) + scales = self._recover_tensor(self.weight_block_scales, self._scale_dtypes["weight_block_scales"]) return block_fp8_dequantize_gkn(fp8_w, scales, self.quant_group_size) @torch.compiler.disable @@ -85,9 +101,7 @@ def _compute_aqn_step(self) -> Tensor: """block_fp8: 0.125 * block_scale (FP8 E4M3 ULP at unit).""" M, K = self.out_features, self.in_features bs = self.quant_group_size - scales = self._recover_tensor( - self.weight_block_scales, self._scale_dtypes["weight_block_scales"] - ).float() + scales = self._recover_tensor(self.weight_block_scales, self._scale_dtypes["weight_block_scales"]).float() step = scales.repeat_interleave(bs, dim=0).repeat_interleave(bs, dim=1)[:M, :K] return (0.125 * step).contiguous() diff --git a/src/xorl/qlora/modules/linear.py b/src/xorl/qlora/modules/linear.py index 6630313a..489a246d 100644 --- a/src/xorl/qlora/modules/linear.py +++ b/src/xorl/qlora/modules/linear.py @@ -54,9 +54,7 @@ def __init__( aqn_alpha: float = 1.0, ): if type(self) is QLoRALinear: - raise TypeError( - "QLoRALinear is abstract. Use NvFP4QLoRALinear, BlockFP8QLoRALinear, or NF4QLoRALinear." - ) + raise TypeError("QLoRALinear is abstract. Use NvFP4QLoRALinear, BlockFP8QLoRALinear, or NF4QLoRALinear.") super().__init__() self.in_features = in_features self.out_features = out_features @@ -113,9 +111,10 @@ def from_module(cls, module: nn.Module, r: int = 16, lora_alpha: int = 16, **kwa based on ``quant_format`` kwarg (default: "nvfp4"). """ if cls is QLoRALinear: - from xorl.qlora.modules.nvfp4_linear import NvFP4QLoRALinear from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear from xorl.qlora.modules.nf4_linear import NF4QLoRALinear + from xorl.qlora.modules.nvfp4_linear import NvFP4QLoRALinear + fmt = kwargs.pop("quant_format", "nvfp4") kwargs.pop("quant_group_size", None) # subclass sets this if fmt == "block_fp8": @@ -146,9 +145,10 @@ def from_quantized( ) -> "QLoRALinear": """Create QLoRALinear from pre-quantized weights (skip quantization).""" if cls is QLoRALinear: - from xorl.qlora.modules.nvfp4_linear import NvFP4QLoRALinear from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear from xorl.qlora.modules.nf4_linear import NF4QLoRALinear + from xorl.qlora.modules.nvfp4_linear import NvFP4QLoRALinear + if quant_format == "block_fp8": cls = BlockFP8QLoRALinear elif quant_format == "nf4": @@ -215,11 +215,12 @@ def _write_packed_weight(self, uint8_data: Tensor) -> None: target_shape = self.packed_weight_f32.shape f32_data = f32_data.reshape(target_shape) self.packed_weight_f32 = nn.Parameter(f32_data, requires_grad=False) - elif hasattr(self.packed_weight_f32, '_local_tensor'): + elif hasattr(self.packed_weight_f32, "_local_tensor"): local = self.packed_weight_f32._local_tensor full_shape = self.packed_weight_f32.shape f32_data = f32_data.reshape(full_shape) from torch.distributed._tensor import Shard as _Shard + placements = self.packed_weight_f32.placements shard_dim = 0 for p in placements: @@ -250,7 +251,7 @@ def _write_packed_weight(self, uint8_data: Tensor) -> None: def _read_packed_weight_uint8(self) -> Tensor: """Read the packed weight as uint8 data (unpack float32 -> uint8).""" f32 = self.packed_weight_f32.data - if hasattr(f32, '_local_tensor'): + if hasattr(f32, "_local_tensor"): f32 = f32.full_tensor() elif f32.dtype != torch.float32: raise RuntimeError( @@ -294,6 +295,7 @@ def quantize_weight(self) -> None: return w = self.weight.data from torch.distributed._tensor import DTensor + if isinstance(w, DTensor): w = w.full_tensor() self._quantize_and_store(w) @@ -323,10 +325,12 @@ def _aqn_start_noise(self, device: torch.device, dtype: torch.dtype) -> None: M, K = self.out_features, self.in_features if self._aqn_stream is None: self._aqn_stream = torch.cuda.Stream(device=device) - if (self._aqn_noise_buf is None - or self._aqn_noise_buf.shape != (M, K) - or self._aqn_noise_buf.dtype != dtype - or self._aqn_noise_buf.device != device): + if ( + self._aqn_noise_buf is None + or self._aqn_noise_buf.shape != (M, K) + or self._aqn_noise_buf.dtype != dtype + or self._aqn_noise_buf.device != device + ): self._aqn_noise_buf = torch.empty(M, K, device=device, dtype=dtype) self._aqn_stream.wait_stream(torch.cuda.current_stream(device)) with torch.cuda.stream(self._aqn_stream): diff --git a/src/xorl/qlora/modules/moe_experts.py b/src/xorl/qlora/modules/moe_experts.py index 46a9ff61..979470b3 100644 --- a/src/xorl/qlora/modules/moe_experts.py +++ b/src/xorl/qlora/modules/moe_experts.py @@ -27,18 +27,19 @@ import torch import torch.nn as nn -import torch.nn.functional as F from torch import Tensor +from xorl.lora.modules.base import LoraModule +from xorl.ops.group_gemm.kernel.lora_utils import compute_lora_scaling from xorl.ops.quantize import ( - nvfp4_quantize_gkn, + block_fp8_dequantize_gkn, + block_fp8_quantize_gkn, + nf4_dequantize_gkn, + nf4_quantize_gkn, nvfp4_dequantize_gkn, + nvfp4_quantize_gkn, ) from xorl.ops.quantize.fp4_codec import FP4_E2M1_MAX, FP8_E4M3_MAX -from xorl.ops.quantize import block_fp8_quantize_gkn, block_fp8_dequantize_gkn -from xorl.ops.quantize import nf4_quantize_gkn, nf4_dequantize_gkn -from xorl.ops.group_gemm.kernel.lora_utils import compute_lora_scaling -from xorl.lora.modules.base import LoraModule class QLoRAMoeExperts(LoraModule, nn.Module): @@ -112,16 +113,31 @@ def __init__( # automatically replicated by the EP plan. shared_exp = 1 if hybrid_shared else num_experts self._create_lora_params( - "gate_proj", shared_exp, num_experts, r, - hidden_size, intermediate_size, device, + "gate_proj", + shared_exp, + num_experts, + r, + hidden_size, + intermediate_size, + device, ) self._create_lora_params( - "up_proj", shared_exp, num_experts, r, - hidden_size, intermediate_size, device, + "up_proj", + shared_exp, + num_experts, + r, + hidden_size, + intermediate_size, + device, ) self._create_lora_params( - "down_proj", num_experts, (1 if hybrid_shared else num_experts), r, - intermediate_size, hidden_size, device, + "down_proj", + num_experts, + (1 if hybrid_shared else num_experts), + r, + intermediate_size, + hidden_size, + device, ) self._scale_dtypes = {"gate": {}, "up": {}, "down": {}} @@ -142,20 +158,30 @@ def __init__( # then call self.reset_lora_parameters() def _create_lora_params( - self, name: str, A_experts: int, B_experts: int, r: int, - in_features: int, out_features: int, device: Optional[torch.device] = None, + self, + name: str, + A_experts: int, + B_experts: int, + r: int, + in_features: int, + out_features: int, + device: Optional[torch.device] = None, ): """Create LoRA A and B parameters in (G, K, N) format. A: [A_experts, in_features, r] B: [B_experts, r, out_features] """ - setattr(self, f"{name}_lora_A", nn.Parameter( - torch.empty(A_experts, in_features, r, dtype=torch.float32, device=device) - )) - setattr(self, f"{name}_lora_B", nn.Parameter( - torch.empty(B_experts, r, out_features, dtype=torch.float32, device=device) - )) + setattr( + self, + f"{name}_lora_A", + nn.Parameter(torch.empty(A_experts, in_features, r, dtype=torch.float32, device=device)), + ) + setattr( + self, + f"{name}_lora_B", + nn.Parameter(torch.empty(B_experts, r, out_features, dtype=torch.float32, device=device)), + ) def reset_lora_parameters(self): """Initialize LoRA weights: kaiming_uniform for A, zeros for B.""" @@ -278,12 +304,11 @@ def _recover_tensor(self, buf: Tensor, original_dtype: torch.dtype) -> Tensor: return buf return buf.contiguous().view(original_dtype) - def _quantize_proj(self, proj_name: str, w3d: Tensor, - global_amax_per_expert: Optional[Tensor] = None) -> None: + def _quantize_proj(self, proj_name: str, w3d: Tensor, global_amax_per_expert: Optional[Tensor] = None) -> None: """Quantize a 3D [num_local_experts, K, N] tensor and store as stacked buffers.""" results = [] for i in range(w3d.shape[0]): - ga = global_amax_per_expert[i:i+1] if global_amax_per_expert is not None else None + ga = global_amax_per_expert[i : i + 1] if global_amax_per_expert is not None else None packed, scales_dict = self._quantize_2d(w3d[i], global_amax=ga) results.append((packed, scales_dict)) @@ -349,7 +374,9 @@ def dequantize_expert(self, proj_name: str, expert_idx: int, K: int, N: int) -> # Weight loading from checkpoint (only local experts) # ------------------------------------------------------------------ - def load_and_quantize_weights(self, weights_path: str, weight_map: Optional[dict] = None, shard_cache: Optional[dict] = None) -> None: + def load_and_quantize_weights( + self, weights_path: str, weight_map: Optional[dict] = None, shard_cache: Optional[dict] = None + ) -> None: """Load LOCAL expert weights from checkpoint and quantize directly. Each subclass loads its own format via _load_experts(). @@ -425,12 +452,11 @@ def forward( # Check EP -- use unified dispatch/compute/combine path from xorl.distributed.parallel_state import get_parallel_state + parallel_state = get_parallel_state() if parallel_state.ep_enabled: - return self._ep_forward( - hidden_states, routing_weights, selected_experts, parallel_state - ) + return self._ep_forward(hidden_states, routing_weights, selected_experts, parallel_state) # Local path -- registry-based from xorl.models.layers.moe.backend import MOE_EXPERT_BACKENDS_LORA @@ -467,7 +493,7 @@ def _ep_forward( Uses the same dispatch/combine as MoEExperts._ep_forward() but routes to the LoRA-aware EP compute registry. """ - from xorl.models.layers.moe.backend import EP_DISPATCH, EP_COMBINE, EP_EXPERT_COMPUTE_LORA + from xorl.models.layers.moe.backend import EP_COMBINE, EP_DISPATCH, EP_EXPERT_COMPUTE_LORA if self.moe_implementation not in EP_EXPERT_COMPUTE_LORA: raise ValueError( @@ -476,8 +502,7 @@ def _ep_forward( ) if self.ep_dispatch not in EP_DISPATCH: raise ValueError( - f"ep_dispatch={self.ep_dispatch!r} is not available. " - f"Available: {list(EP_DISPATCH.keys())}" + f"ep_dispatch={self.ep_dispatch!r} is not available. Available: {list(EP_DISPATCH.keys())}" ) dispatch_fn = EP_DISPATCH[self.ep_dispatch] @@ -485,28 +510,28 @@ def _ep_forward( compute_fn = EP_EXPERT_COMPUTE_LORA[self.moe_implementation] # Step 1: Dispatch tokens to expert-owning ranks - dispatch_kwargs = self._build_dispatch_kwargs( - hidden_states, routing_weights, selected_experts, parallel_state - ) + dispatch_kwargs = self._build_dispatch_kwargs(hidden_states, routing_weights, selected_experts, parallel_state) permute_tokens, cumsum, ctx = dispatch_fn(**dispatch_kwargs) # Step 2: Expert computation with dequantized base + LoRA compute_dtype = permute_tokens.dtype expert_output = compute_fn( - permute_tokens, cumsum, + permute_tokens, + cumsum, self.gate_proj.to(compute_dtype), self.up_proj.to(compute_dtype), self.down_proj.to(compute_dtype), - self.gate_proj_lora_A, self.gate_proj_lora_B, - self.up_proj_lora_A, self.up_proj_lora_B, - self.down_proj_lora_A, self.down_proj_lora_B, + self.gate_proj_lora_A, + self.gate_proj_lora_B, + self.up_proj_lora_A, + self.up_proj_lora_B, + self.down_proj_lora_A, + self.down_proj_lora_B, self.scaling, ) # Step 3: Combine expert outputs back to original ranks - combine_kwargs = self._build_combine_kwargs( - expert_output, ctx, dispatch_kwargs, parallel_state - ) + combine_kwargs = self._build_combine_kwargs(expert_output, ctx, dispatch_kwargs, parallel_state) return combine_fn(**combine_kwargs) def _build_dispatch_kwargs(self, hidden_states, routing_weights, selected_experts, parallel_state): @@ -521,6 +546,7 @@ def _build_dispatch_kwargs(self, hidden_states, routing_weights, selected_expert kwargs["ep_group"] = parallel_state.ep_group elif self.ep_dispatch == "deepep": from xorl.distributed.moe.deepep import get_default_buffer + kwargs["buffer"] = get_default_buffer( ep_group=parallel_state.ep_group, buffer_size_gb=self.deepep_buffer_size_gb, @@ -551,7 +577,7 @@ def _eager_lora_forward(self, hidden_states: Tensor, expert_idx: int) -> Tensor: # Convert global expert index to local local_idx = expert_idx - self.expert_offset if local_idx < 0 or local_idx >= self.num_local_experts: - return torch.zeros_like(hidden_states[:, :self.hidden_size]) + return torch.zeros_like(hidden_states[:, : self.hidden_size]) compute_dtype = hidden_states.dtype @@ -629,12 +655,21 @@ def __init__( deepep_async_combine: bool = False, ): super().__init__( - num_local_experts, num_experts, intermediate_size, hidden_size, - r=r, lora_alpha=lora_alpha, - quant_format="nvfp4", quant_group_size=16, - act_fn=act_fn, expert_offset=expert_offset, device=device, - moe_implementation=moe_implementation, hybrid_shared=hybrid_shared, - use_rslora=use_rslora, ep_dispatch=ep_dispatch, + num_local_experts, + num_experts, + intermediate_size, + hidden_size, + r=r, + lora_alpha=lora_alpha, + quant_format="nvfp4", + quant_group_size=16, + act_fn=act_fn, + expert_offset=expert_offset, + device=device, + moe_implementation=moe_implementation, + hybrid_shared=hybrid_shared, + use_rslora=use_rslora, + ep_dispatch=ep_dispatch, deepep_buffer_size_gb=deepep_buffer_size_gb, deepep_num_sms=deepep_num_sms, deepep_async_combine=deepep_async_combine, @@ -680,15 +715,15 @@ def _load_experts(self, _load_tensor, _shard_cache) -> None: amax_list: list = [] if is_prequantized: - packed_hf_list: list = [] # [N, K//2] uint8 per expert - bs_u8_list: list = [] # [N, K//BS] uint8 (fp8 bits reinterpreted) per expert - gs_list: list = [] # scalar f32 per expert + packed_hf_list: list = [] # [N, K//2] uint8 per expert + bs_u8_list: list = [] # [N, K//BS] uint8 (fp8 bits reinterpreted) per expert + gs_list: list = [] # scalar f32 per expert for i in expert_range: fqn_prefix = f"{self._source_fqn}.{i}.{hf_name}" - packed = _load_tensor(f"{fqn_prefix}.weight") # [N, K//2] uint8 - bs = _load_tensor(f"{fqn_prefix}.weight_scale") # [N, K//BS] fp8 - gs = _load_tensor(f"{fqn_prefix}.weight_scale_2") # scalar f32 + packed = _load_tensor(f"{fqn_prefix}.weight") # [N, K//2] uint8 + bs = _load_tensor(f"{fqn_prefix}.weight_scale") # [N, K//BS] fp8 + gs = _load_tensor(f"{fqn_prefix}.weight_scale_2") # scalar f32 gs_val = gs.float().item() packed_hf_list.append(packed) @@ -699,21 +734,21 @@ def _load_experts(self, _load_tensor, _shard_cache) -> None: # Stack in HF layout — no per-expert transpose on CPU packed_hf = torch.stack(packed_hf_list) # [E, N, K//2] uint8 - bs_u8_hf = torch.stack(bs_u8_list) # [E, N, K//BS] uint8 (fp8 bits) - gs_cpu = torch.tensor(gs_list, dtype=torch.float32) # [E] + bs_u8_hf = torch.stack(bs_u8_list) # [E, N, K//BS] uint8 (fp8 bits) + gs_cpu = torch.tensor(gs_list, dtype=torch.float32) # [E] # Bulk H2D copy - packed_gpu = packed_hf.to(device) # [E, N, K//2] uint8 - bs_u8_gpu = bs_u8_hf.to(device) # [E, N, K//BS] uint8 - gs_gpu = gs_cpu.to(device) # [E] f32 + packed_gpu = packed_hf.to(device) # [E, N, K//2] uint8 + bs_u8_gpu = bs_u8_hf.to(device) # [E, N, K//BS] uint8 + gs_gpu = gs_cpu.to(device) # [E] f32 # GPU: reinterpret uint8 → fp8, upcast to f32, absorb global scale bs_f32 = bs_u8_gpu.view(torch.float8_e4m3fn).float() # [E, N, K//BS] f32 - bs_absorbed = bs_f32 * gs_gpu[:, None, None] # [E, N, K//BS] f32 + bs_absorbed = bs_f32 * gs_gpu[:, None, None] # [E, N, K//BS] f32 # GPU permute HF [E, N, *] → GKN [E, *, N] - packed_gkn = packed_gpu.permute(0, 2, 1).contiguous() # [E, K//2, N] uint8 - bs_gkn = bs_absorbed.permute(0, 2, 1).contiguous() # [E, K//BS, N] f32 + packed_gkn = packed_gpu.permute(0, 2, 1).contiguous() # [E, K//2, N] uint8 + bs_gkn = bs_absorbed.permute(0, 2, 1).contiguous() # [E, K//BS, N] f32 else: # BF16 checkpoint: load weight, move to GPU, quantize per expert. @@ -740,13 +775,13 @@ def _load_experts(self, _load_tensor, _shard_cache) -> None: _shard_cache.clear() packed_gkn = torch.stack(packed_gkn_list) # [E, K//2, N] uint8 - bs_gkn = torch.stack(bs_gkn_list) # [E, K//BS, N] f32 (in uint8) + bs_gkn = torch.stack(bs_gkn_list) # [E, K//BS, N] f32 (in uint8) - setattr(self, f"{proj_name}_packed", self._to_uint8(packed_gkn)) + setattr(self, f"{proj_name}_packed", self._to_uint8(packed_gkn)) setattr(self, f"{proj_name}_block_scales", self._to_uint8(bs_gkn)) - setattr(self, f"{proj_name}_global_scale", self._to_uint8( - torch.ones(E, 1, dtype=torch.float32, device=device) - )) + setattr( + self, f"{proj_name}_global_scale", self._to_uint8(torch.ones(E, 1, dtype=torch.float32, device=device)) + ) self._scale_dtypes[proj_name] = { "weight_block_scales": torch.float32, "weight_global_scale": torch.float32, @@ -758,9 +793,7 @@ def _load_experts(self, _load_tensor, _shard_cache) -> None: def _quantize_2d(self, w: Tensor, global_amax=None): """Quantize a 2D [K, N] weight using NVFP4. Groups along K (contraction dim).""" - packed, block_scales, global_scale = nvfp4_quantize_gkn( - w, self.quant_group_size, global_amax=global_amax - ) + packed, block_scales, global_scale = nvfp4_quantize_gkn(w, self.quant_group_size, global_amax=global_amax) return self._to_uint8(packed), { "weight_block_scales": (self._to_uint8(block_scales), block_scales.dtype), "weight_global_scale": (self._to_uint8(global_scale), global_scale.dtype), @@ -785,13 +818,12 @@ def merge_weights(self, ema_decay: float = 0.1) -> None: w_merged = w + delta fresh_amax = w_merged.float().abs().amax(dim=(1, 2)) # [E] if self._ema_amax[proj_name] is not None: - self._ema_amax[proj_name].lerp_( - fresh_amax.to(self._ema_amax[proj_name].device), ema_decay - ) + self._ema_amax[proj_name].lerp_(fresh_amax.to(self._ema_amax[proj_name].device), ema_decay) else: self._ema_amax[proj_name] = fresh_amax self._quantize_proj( - proj_name, w_merged, + proj_name, + w_merged, global_amax_per_expert=self._ema_amax[proj_name], ) self.reset_lora_parameters() @@ -825,12 +857,21 @@ def __init__( deepep_async_combine: bool = False, ): super().__init__( - num_local_experts, num_experts, intermediate_size, hidden_size, - r=r, lora_alpha=lora_alpha, - quant_format="block_fp8", quant_group_size=128, - act_fn=act_fn, expert_offset=expert_offset, device=device, - moe_implementation=moe_implementation, hybrid_shared=hybrid_shared, - use_rslora=use_rslora, ep_dispatch=ep_dispatch, + num_local_experts, + num_experts, + intermediate_size, + hidden_size, + r=r, + lora_alpha=lora_alpha, + quant_format="block_fp8", + quant_group_size=128, + act_fn=act_fn, + expert_offset=expert_offset, + device=device, + moe_implementation=moe_implementation, + hybrid_shared=hybrid_shared, + use_rslora=use_rslora, + ep_dispatch=ep_dispatch, deepep_buffer_size_gb=deepep_buffer_size_gb, deepep_num_sms=deepep_num_sms, deepep_async_combine=deepep_async_combine, @@ -860,12 +901,12 @@ def _load_experts(self, _load_tensor, _shard_cache) -> None: expert_range = range(self.expert_offset, self.expert_offset + E) for proj_name, hf_name in [("gate", "gate_proj"), ("up", "up_proj"), ("down", "down_proj")]: - fp8_u8_list: list = [] # [N, K] uint8 (fp8 bits reinterpreted) per expert - scales_list: list = [] # [N//BS, K//BS] f32 per expert + fp8_u8_list: list = [] # [N, K] uint8 (fp8 bits reinterpreted) per expert + scales_list: list = [] # [N//BS, K//BS] f32 per expert for i in expert_range: fqn_prefix = f"{self._source_fqn}.{i}.{hf_name}" - fp8_w = _load_tensor(f"{fqn_prefix}.weight") # [N, K] fp8 + fp8_w = _load_tensor(f"{fqn_prefix}.weight") # [N, K] fp8 scales = _load_tensor(f"{fqn_prefix}.weight_scale_inv") # [N//BS, K//BS] bf16/f32 # Reinterpret fp8 bits as uint8 (both 1-byte, zero-copy on CPU) @@ -873,18 +914,18 @@ def _load_experts(self, _load_tensor, _shard_cache) -> None: scales_list.append(scales.float()) # Stack in HF layout — no per-expert transpose on CPU - fp8_u8_hf = torch.stack(fp8_u8_list) # [E, N, K] uint8 - scales_hf = torch.stack(scales_list) # [E, N//BS, K//BS] f32 + fp8_u8_hf = torch.stack(fp8_u8_list) # [E, N, K] uint8 + scales_hf = torch.stack(scales_list) # [E, N//BS, K//BS] f32 # Bulk H2D copy - fp8_u8_gpu = fp8_u8_hf.to(device) # [E, N, K] uint8 - scales_gpu = scales_hf.to(device) # [E, N//BS, K//BS] f32 + fp8_u8_gpu = fp8_u8_hf.to(device) # [E, N, K] uint8 + scales_gpu = scales_hf.to(device) # [E, N//BS, K//BS] f32 # GPU permute HF [E, N, *] → GKN [E, *, N] - fp8_gkn = fp8_u8_gpu.permute(0, 2, 1).contiguous() # [E, K, N] uint8 - scales_gkn = scales_gpu.permute(0, 2, 1).contiguous() # [E, K//BS, N//BS] f32 + fp8_gkn = fp8_u8_gpu.permute(0, 2, 1).contiguous() # [E, K, N] uint8 + scales_gkn = scales_gpu.permute(0, 2, 1).contiguous() # [E, K//BS, N//BS] f32 - setattr(self, f"{proj_name}_packed", self._to_uint8(fp8_gkn)) + setattr(self, f"{proj_name}_packed", self._to_uint8(fp8_gkn)) setattr(self, f"{proj_name}_block_scales", self._to_uint8(scales_gkn)) self._scale_dtypes[proj_name] = { "weight_block_scales": torch.float32, @@ -949,12 +990,21 @@ def __init__( deepep_async_combine: bool = False, ): super().__init__( - num_local_experts, num_experts, intermediate_size, hidden_size, - r=r, lora_alpha=lora_alpha, - quant_format="nf4", quant_group_size=64, - act_fn=act_fn, expert_offset=expert_offset, device=device, - moe_implementation=moe_implementation, hybrid_shared=hybrid_shared, - use_rslora=use_rslora, ep_dispatch=ep_dispatch, + num_local_experts, + num_experts, + intermediate_size, + hidden_size, + r=r, + lora_alpha=lora_alpha, + quant_format="nf4", + quant_group_size=64, + act_fn=act_fn, + expert_offset=expert_offset, + device=device, + moe_implementation=moe_implementation, + hybrid_shared=hybrid_shared, + use_rslora=use_rslora, + ep_dispatch=ep_dispatch, deepep_buffer_size_gb=deepep_buffer_size_gb, deepep_num_sms=deepep_num_sms, deepep_async_combine=deepep_async_combine, diff --git a/src/xorl/qlora/modules/nf4_linear.py b/src/xorl/qlora/modules/nf4_linear.py index 96999e0c..8ca36838 100644 --- a/src/xorl/qlora/modules/nf4_linear.py +++ b/src/xorl/qlora/modules/nf4_linear.py @@ -7,7 +7,7 @@ import torch.nn as nn from torch import Tensor -from xorl.ops.quantize import nf4_quantize, nf4_dequantize +from xorl.ops.quantize import nf4_dequantize, nf4_quantize from xorl.ops.quantize.nf4_codec import NF4_MIN_STEP from xorl.qlora.modules.linear import QLoRALinear @@ -37,9 +37,16 @@ def __init__( aqn_alpha: float = 1.0, ): super().__init__( - in_features, out_features, r=r, lora_alpha=lora_alpha, - quant_format="nf4", quant_group_size=64, - bias=bias, device=device, enable_aqn=enable_aqn, aqn_alpha=aqn_alpha, + in_features, + out_features, + r=r, + lora_alpha=lora_alpha, + quant_format="nf4", + quant_group_size=64, + bias=bias, + device=device, + enable_aqn=enable_aqn, + aqn_alpha=aqn_alpha, ) # nf4: 2 codes per byte -> in_features // 2 bytes -> in_features // 8 float32 elements pw_cols = in_features // 8 @@ -56,8 +63,15 @@ def __init__( self.reset_lora_parameters() @classmethod - def from_module(cls, module: nn.Module, r: int = 16, lora_alpha: int = 16, - enable_aqn: bool = False, aqn_alpha: float = 1.0, **kwargs) -> "NF4QLoRALinear": + def from_module( + cls, + module: nn.Module, + r: int = 16, + lora_alpha: int = 16, + enable_aqn: bool = False, + aqn_alpha: float = 1.0, + **kwargs, + ) -> "NF4QLoRALinear": """Create from a bf16 nn.Linear by quantizing its weight to NF4. On meta device: defers quantization — keeps ``weight`` as a parameter @@ -65,10 +79,14 @@ def from_module(cls, module: nn.Module, r: int = 16, lora_alpha: int = 16, On real device: quantizes immediately and discards ``weight``. """ qlora = cls( - in_features=module.in_features, out_features=module.out_features, - r=r, lora_alpha=lora_alpha, - bias=module.bias is not None, device=module.weight.device, - enable_aqn=enable_aqn, aqn_alpha=aqn_alpha, + in_features=module.in_features, + out_features=module.out_features, + r=r, + lora_alpha=lora_alpha, + bias=module.bias is not None, + device=module.weight.device, + enable_aqn=enable_aqn, + aqn_alpha=aqn_alpha, ) if module.weight.device.type == "meta": qlora.weight = nn.Parameter(module.weight.detach(), requires_grad=False) @@ -96,9 +114,7 @@ def _compute_aqn_step(self) -> Tensor: """nf4: 0.5 * NF4_MIN_STEP * per_group_scale (minimum quantization resolution).""" M, K = self.out_features, self.in_features gs = self.quant_group_size - scales = self._recover_tensor( - self.weight_scales, self._scale_dtypes["weight_scales"] - ).float() + scales = self._recover_tensor(self.weight_scales, self._scale_dtypes["weight_scales"]).float() step = scales.reshape(M, K // gs).repeat_interleave(gs, dim=1) return (0.5 * NF4_MIN_STEP * step).contiguous() diff --git a/src/xorl/qlora/modules/nvfp4_linear.py b/src/xorl/qlora/modules/nvfp4_linear.py index b14200c7..cb0055a3 100644 --- a/src/xorl/qlora/modules/nvfp4_linear.py +++ b/src/xorl/qlora/modules/nvfp4_linear.py @@ -7,7 +7,7 @@ import torch.nn as nn from torch import Tensor -from xorl.ops.quantize import nvfp4_quantize, nvfp4_dequantize +from xorl.ops.quantize import nvfp4_dequantize, nvfp4_quantize from xorl.qlora.modules.linear import QLoRALinear @@ -33,9 +33,16 @@ def __init__( aqn_alpha: float = 1.0, ): super().__init__( - in_features, out_features, r=r, lora_alpha=lora_alpha, - quant_format="nvfp4", quant_group_size=16, - bias=bias, device=device, enable_aqn=enable_aqn, aqn_alpha=aqn_alpha, + in_features, + out_features, + r=r, + lora_alpha=lora_alpha, + quant_format="nvfp4", + quant_group_size=16, + bias=bias, + device=device, + enable_aqn=enable_aqn, + aqn_alpha=aqn_alpha, ) # nvfp4: 2 fp4 values per byte -> in_features // 2 bytes -> in_features // 8 float32 elements pw_cols = in_features // 8 @@ -55,14 +62,25 @@ def __init__( self.reset_lora_parameters() @classmethod - def from_module(cls, module: nn.Module, r: int = 16, lora_alpha: int = 16, - enable_aqn: bool = False, aqn_alpha: float = 1.0, **kwargs) -> "NvFP4QLoRALinear": + def from_module( + cls, + module: nn.Module, + r: int = 16, + lora_alpha: int = 16, + enable_aqn: bool = False, + aqn_alpha: float = 1.0, + **kwargs, + ) -> "NvFP4QLoRALinear": """Create from a bf16 nn.Linear by quantizing its weight to nvfp4.""" qlora = cls( - in_features=module.in_features, out_features=module.out_features, - r=r, lora_alpha=lora_alpha, - bias=module.bias is not None, device=module.weight.device, - enable_aqn=enable_aqn, aqn_alpha=aqn_alpha, + in_features=module.in_features, + out_features=module.out_features, + r=r, + lora_alpha=lora_alpha, + bias=module.bias is not None, + device=module.weight.device, + enable_aqn=enable_aqn, + aqn_alpha=aqn_alpha, ) w = module.weight.detach() qlora._ema_amax = w.float().abs().max().reshape(1).to(w.device) @@ -71,7 +89,9 @@ def from_module(cls, module: nn.Module, r: int = 16, lora_alpha: int = 16, def _quantize_and_store(self, w: Tensor, global_amax: Optional[Tensor] = None) -> None: packed, block_scales, global_scale = nvfp4_quantize( - w, self.quant_group_size, global_amax=global_amax, + w, + self.quant_group_size, + global_amax=global_amax, ) uint8_data = self._to_uint8(packed) self._write_packed_weight(uint8_data) @@ -87,12 +107,8 @@ def _quantize_and_store(self, w: Tensor, global_amax: Optional[Tensor] = None) - def _dequantize_weight(self) -> Tensor: M, K = self.out_features, self.in_features uint8_data = self._read_packed_weight_uint8() - block_scales = self._recover_tensor( - self.weight_block_scales, self._scale_dtypes["weight_block_scales"] - ) - global_scale = self._recover_tensor( - self.weight_global_scale, self._scale_dtypes["weight_global_scale"] - ) + block_scales = self._recover_tensor(self.weight_block_scales, self._scale_dtypes["weight_block_scales"]) + global_scale = self._recover_tensor(self.weight_global_scale, self._scale_dtypes["weight_global_scale"]) w = nvfp4_dequantize(uint8_data, block_scales, global_scale, M * K, self.quant_group_size) return w.reshape(M, K) @@ -101,12 +117,8 @@ def _compute_aqn_step(self) -> Tensor: """nvfp4: 0.5 * block_scale * global_scale (FP4 min linear step).""" M, K = self.out_features, self.in_features bs = self.quant_group_size - block_scales = self._recover_tensor( - self.weight_block_scales, self._scale_dtypes["weight_block_scales"] - ).float() - global_scale = self._recover_tensor( - self.weight_global_scale, self._scale_dtypes["weight_global_scale"] - ).float() + block_scales = self._recover_tensor(self.weight_block_scales, self._scale_dtypes["weight_block_scales"]).float() + global_scale = self._recover_tensor(self.weight_global_scale, self._scale_dtypes["weight_global_scale"]).float() effective = block_scales * global_scale step = effective.reshape(M, K // bs).repeat_interleave(bs, dim=1) return (0.5 * step).contiguous() diff --git a/src/xorl/qlora/utils.py b/src/xorl/qlora/utils.py index 5f9a9880..3be7331a 100644 --- a/src/xorl/qlora/utils.py +++ b/src/xorl/qlora/utils.py @@ -7,17 +7,18 @@ import json import logging import os -from typing import Collection, Dict, List, Optional, Set, Tuple +from typing import Collection, List, Optional, Tuple import torch import torch.nn as nn from safetensors.torch import save_file -from xorl.qlora.modules.linear import QLoRALinear -from xorl.qlora.modules.nvfp4_linear import NvFP4QLoRALinear from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear -from xorl.qlora.modules.nf4_linear import NF4QLoRALinear +from xorl.qlora.modules.linear import QLoRALinear from xorl.qlora.modules.moe_experts import QLoRAMoeExperts +from xorl.qlora.modules.nf4_linear import NF4QLoRALinear +from xorl.qlora.modules.nvfp4_linear import NvFP4QLoRALinear + logger = logging.getLogger(__name__) @@ -96,9 +97,7 @@ def inject_qlora_into_model( aqn_alpha: Scale factor for AQN noise magnitude (default 1.0). """ if quant_format not in ("nvfp4", "block_fp8", "nf4"): - raise ValueError( - f"Supported QLoRA formats: 'nvfp4', 'block_fp8', 'nf4'. Got quant_format={quant_format!r}" - ) + raise ValueError(f"Supported QLoRA formats: 'nvfp4', 'block_fp8', 'nf4'. Got quant_format={quant_format!r}") if target_modules is None: # Pre-quantized checkpoints: target the fused module names in the model. # Separate HF projections (q/k/v, gate/up) are merged during loading. @@ -128,15 +127,11 @@ def inject_qlora_into_model( if exclude_modules: all_short_names = {p.split(".")[-1] for p in target_paths} before = len(target_paths) - target_paths = [ - p for p in target_paths - if p.split(".")[-1] not in exclude_modules - ] + target_paths = [p for p in target_paths if p.split(".")[-1] not in exclude_modules] excluded = before - len(target_paths) if excluded > 0: logger.info( - f"Excluded {excluded} modules from QLoRA injection " - f"(not quantized in checkpoint): {exclude_modules}" + f"Excluded {excluded} modules from QLoRA injection (not quantized in checkpoint): {exclude_modules}" ) # Warn about exclude_modules names that didn't match any target module unmatched = exclude_modules - all_short_names @@ -169,8 +164,11 @@ def inject_qlora_into_model( if quant_format == "nf4": # NF4: quantize bf16 weights on-the-fly (no pre-quantized checkpoint) qlora_module = NF4QLoRALinear.from_module( - original_module, r=r, lora_alpha=lora_alpha, - enable_aqn=enable_aqn, aqn_alpha=aqn_alpha, + original_module, + r=r, + lora_alpha=lora_alpha, + enable_aqn=enable_aqn, + aqn_alpha=aqn_alpha, ) else: # nvfp4/block_fp8: create empty shell, load pre-quantized weights later @@ -213,6 +211,7 @@ def inject_qlora_into_model( return model from xorl.distributed.parallel_state import get_parallel_state + try: parallel_state = get_parallel_state() ep_size = parallel_state.ep_size if parallel_state.ep_enabled else 1 @@ -228,9 +227,12 @@ def inject_qlora_into_model( continue experts = module.experts has_3d_weights = ( - hasattr(experts, "gate_proj") and isinstance(experts.gate_proj, nn.Parameter) - and hasattr(experts, "up_proj") and isinstance(experts.up_proj, nn.Parameter) - and hasattr(experts, "down_proj") and isinstance(experts.down_proj, nn.Parameter) + hasattr(experts, "gate_proj") + and isinstance(experts.gate_proj, nn.Parameter) + and hasattr(experts, "up_proj") + and isinstance(experts.up_proj, nn.Parameter) + and hasattr(experts, "down_proj") + and isinstance(experts.down_proj, nn.Parameter) and experts.gate_proj.dim() == 3 and not isinstance(experts, QLoRAMoeExperts) ) @@ -266,7 +268,7 @@ def inject_qlora_into_model( ("up", experts.up_proj), ("down", experts.down_proj), ]: - local_w = src_param[expert_offset:expert_offset + num_local_experts] + local_w = src_param[expert_offset : expert_offset + num_local_experts] qlora_experts._quantize_proj(proj_name, local_w.float()) qlora_experts._weights_loaded = True else: @@ -337,10 +339,7 @@ def maybe_quantize_qlora(model: nn.Module) -> int: Returns: Number of modules quantized. """ - modules_to_quantize = [ - m for m in model.modules() - if isinstance(m, QLoRALinear) and m.weight is not None - ] + modules_to_quantize = [m for m in model.modules() if isinstance(m, QLoRALinear) and m.weight is not None] if not modules_to_quantize: return 0 @@ -378,12 +377,8 @@ def maybe_load_and_quantize_moe_qlora( Number of MoE modules loaded. """ import json - import torch.distributed as dist - needs_moe = any( - isinstance(m, QLoRAMoeExperts) and not m._weights_loaded - for m in model.modules() - ) + needs_moe = any(isinstance(m, QLoRAMoeExperts) and not m._weights_loaded for m in model.modules()) if not needs_moe: return 0 @@ -391,7 +386,8 @@ def maybe_load_and_quantize_moe_qlora( weight_map = None if weights_path: try: - from transformers.utils import cached_file, SAFE_WEIGHTS_INDEX_NAME + from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, cached_file + index_path = cached_file(weights_path, SAFE_WEIGHTS_INDEX_NAME) if index_path: with open(index_path) as f: @@ -412,10 +408,13 @@ def maybe_load_and_quantize_moe_qlora( for module in model.modules(): if isinstance(module, QLoRAMoeExperts) and not module._weights_loaded: module.load_and_quantize_weights( - weights_path, weight_map=weight_map, shard_cache=shard_cache, + weights_path, + weight_map=weight_map, + shard_cache=shard_cache, ) # Materialize LoRA params from meta device to GPU from torch.distributed._tensor import DTensor + for name, param in module.named_parameters(): if param.device.type == "meta": if isinstance(param, DTensor): @@ -423,19 +422,27 @@ def maybe_load_and_quantize_moe_qlora( placement = param.placements mesh = param.device_mesh local_data = torch.zeros( - local_shape, dtype=param.dtype, device="cuda", + local_shape, + dtype=param.dtype, + device="cuda", ) materialized = nn.Parameter( DTensor.from_local( - local_data, mesh, placement, run_check=False, + local_data, + mesh, + placement, + run_check=False, ), requires_grad=param.requires_grad, ) else: import math as _math + materialized = nn.Parameter( torch.zeros( - param.shape, dtype=param.dtype, device="cuda", + param.shape, + dtype=param.dtype, + device="cuda", ), requires_grad=param.requires_grad, ) @@ -443,7 +450,8 @@ def maybe_load_and_quantize_moe_qlora( if parts[-1] == "A": for i in range(materialized.shape[0]): nn.init.kaiming_uniform_( - materialized.data[i], a=_math.sqrt(5), + materialized.data[i], + a=_math.sqrt(5), ) setattr(module, name, materialized) moe_count += 1 @@ -451,9 +459,7 @@ def maybe_load_and_quantize_moe_qlora( shard_cache.clear() if moe_count > 0: - logger.info( - f"Loaded and quantized {moe_count} MoE expert modules from bf16 checkpoint" - ) + logger.info(f"Loaded and quantized {moe_count} MoE expert modules from bf16 checkpoint") return moe_count @@ -498,7 +504,8 @@ def _deregister_qlora_weights_from_fsdp( if pg is None: continue to_remove = [ - fp for fp in pg.fsdp_params + fp + for fp in pg.fsdp_params if hasattr(fp, "_module_info") and fp._module_info.param_name in param_name_set and id(fp._module_info.module) in qlora_modules @@ -511,7 +518,6 @@ def _deregister_qlora_weights_from_fsdp( return removed - def detect_prequantized_nvfp4(weights_path: str) -> bool: """Detect whether a checkpoint contains pre-quantized NVFP4 weights (modelopt format). @@ -524,6 +530,7 @@ def detect_prequantized_nvfp4(weights_path: str) -> bool: True if the checkpoint is pre-quantized NVFP4. """ from xorl.models.checkpoint_handlers.buffers import detect_prequantized_checkpoint + return detect_prequantized_checkpoint(weights_path) @@ -539,6 +546,7 @@ def detect_prequantized_block_fp8(weights_path: str) -> bool: True if the checkpoint is pre-quantized block FP8. """ from xorl.models.checkpoint_handlers.buffers import detect_prequantized_block_fp8_checkpoint + return detect_prequantized_block_fp8_checkpoint(weights_path) @@ -555,10 +563,11 @@ def _broadcast_shard_cache( Returns: shard_cache: dict mapping shard filename → {tensor_name: tensor (CPU)} """ - import torch.distributed as dist from concurrent.futures import ThreadPoolExecutor - from transformers.utils import cached_file + import safetensors.torch + import torch.distributed as dist + from transformers.utils import cached_file rank = dist.get_rank() device = torch.device("cuda") @@ -615,10 +624,7 @@ def _read_shard(shard_file): # Free GPU memory between shards torch.cuda.empty_cache() if rank == 0: - logger.info( - f"Broadcast shard {i + 1}/{len(needed_shards)}: {shard_file} " - f"({len(shard_dict)} tensors)" - ) + logger.info(f"Broadcast shard {i + 1}/{len(needed_shards)}: {shard_file} ({len(shard_dict)} tensors)") if executor is not None: executor.shutdown(wait=False) @@ -650,7 +656,8 @@ def maybe_load_prequantized_qlora(model: nn.Module, weights_path: str, load_mode # Load weight map from index file weight_map = None try: - from transformers.utils import cached_file, SAFE_WEIGHTS_INDEX_NAME + from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, cached_file + index_path = cached_file(weights_path, SAFE_WEIGHTS_INDEX_NAME) if index_path: with open(index_path) as f: @@ -669,8 +676,7 @@ def maybe_load_prequantized_qlora(model: nn.Module, weights_path: str, load_mode # "broadcast" (default): rank 0 reads, broadcasts via NCCL. Best for shared/NFS filesystems. # "all_ranks": every rank reads from disk independently. Best for local SSDs. use_broadcast = ( - load_mode == "broadcast" - and dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1 + load_mode == "broadcast" and dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1 ) if use_broadcast: needed_shards = sorted(set(weight_map.values())) @@ -704,6 +710,7 @@ def maybe_load_prequantized_qlora(model: nn.Module, weights_path: str, load_mode # we initialize ALL DTensor params to zeros and let kaiming_uniform # only apply to non-DTensor (non-EP) params. from torch.distributed._tensor import DTensor + _rank = dist.get_rank() if (dist.is_available() and dist.is_initialized()) else 0 for name, param in module.named_parameters(): is_meta = param.device.type == "meta" @@ -721,11 +728,12 @@ def maybe_load_prequantized_qlora(model: nn.Module, weights_path: str, load_mode mesh = param.device_mesh # DTensor: initialize to zeros (preserves Replicate consistency) local_data = torch.zeros( - local_shape, dtype=param.dtype, device="cuda", + local_shape, + dtype=param.dtype, + device="cuda", ) materialized = nn.Parameter( - DTensor.from_local(local_data, mesh, placement, - run_check=False), + DTensor.from_local(local_data, mesh, placement, run_check=False), requires_grad=param.requires_grad, ) else: @@ -737,6 +745,7 @@ def maybe_load_prequantized_qlora(model: nn.Module, weights_path: str, load_mode parts = name.split("_") if parts[-1] == "A": import math + for i in range(materialized.shape[0]): nn.init.kaiming_uniform_(materialized.data[i], a=math.sqrt(5)) setattr(module, name, materialized) @@ -750,7 +759,8 @@ def maybe_load_prequantized_qlora(model: nn.Module, weights_path: str, load_mode # Deregister packed_weight_f32 from FSDP2 to prevent mixed-precision # bf16 cast that corrupts packed uint8 byte patterns. removed = _deregister_qlora_weights_from_fsdp( - model, param_names=("packed_weight_f32",), + model, + param_names=("packed_weight_f32",), ) torch.cuda.empty_cache() logger.info( @@ -766,7 +776,7 @@ def maybe_load_prequantized_qlora(model: nn.Module, weights_path: str, load_mode _diag_fqn = None _diag_mod = None for fqn, mod in model.named_modules(): - if hasattr(mod, '_dequantize_weight') and hasattr(mod, 'weight_block_scales'): + if hasattr(mod, "_dequantize_weight") and hasattr(mod, "weight_block_scales"): _diag_fqn = fqn _diag_mod = mod break diff --git a/src/xorl/server/api_server/__init__.py b/src/xorl/server/api_server/__init__.py index b6c510ad..b63dd862 100644 --- a/src/xorl/server/api_server/__init__.py +++ b/src/xorl/server/api_server/__init__.py @@ -5,29 +5,30 @@ communicating with the training engine backend. """ -from xorl.server.api_server.orchestrator_client import OrchestratorClient -from xorl.server.api_server.server import APIServer from xorl.server.api_server.api_types import ( - TensorData, + AdamParams, Datum, DatumInput, - ForwardRequest, - ForwardResponse, + ErrorResponse, ForwardBackwardRequest, ForwardBackwardResponse, + ForwardRequest, + ForwardResponse, + HealthCheckResponse, + LoadWeightsRequest, + LoadWeightsResponse, LossFnOutput, - AdamParams, OptimStepRequest, OptimStepResponse, - SaveWeightsRequest, - SaveWeightsResponse, - LoadWeightsRequest, - LoadWeightsResponse, SaveWeightsForSamplerRequest, SaveWeightsForSamplerResponse, - HealthCheckResponse, - ErrorResponse, + SaveWeightsRequest, + SaveWeightsResponse, + TensorData, ) +from xorl.server.api_server.orchestrator_client import OrchestratorClient +from xorl.server.api_server.server import APIServer + __all__ = [ "OrchestratorClient", diff --git a/src/xorl/server/api_server/__main__.py b/src/xorl/server/api_server/__main__.py index bcfcc980..3c759b8f 100644 --- a/src/xorl/server/api_server/__main__.py +++ b/src/xorl/server/api_server/__main__.py @@ -19,63 +19,43 @@ import asyncio import logging import sys -import uvicorn from pathlib import Path +import uvicorn + + # Add src to path repo_root = Path(__file__).parent.parent.parent.parent sys.path.insert(0, str(repo_root)) -from xorl.server.api_server.server import APIServer, app +from xorl.server.api_server.server import app from xorl.server.backend import DummyBackend from xorl.server.orchestrator.orchestrator import Orchestrator -logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" -) +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) def parse_args(): - parser = argparse.ArgumentParser( - description="Launch API server with optional mock engine" - ) - parser.add_argument( - "--host", - type=str, - default="127.0.0.1", - help="API server host (default: 127.0.0.1)" - ) - parser.add_argument( - "--port", - type=int, - default=20000, - help="API server port (default: 20000)" - ) - parser.add_argument( - "--mock", - action="store_true", - help="Launch with mock engine for testing" - ) + parser = argparse.ArgumentParser(description="Launch API server with optional mock engine") + parser.add_argument("--host", type=str, default="127.0.0.1", help="API server host (default: 127.0.0.1)") + parser.add_argument("--port", type=int, default=20000, help="API server port (default: 20000)") + parser.add_argument("--mock", action="store_true", help="Launch with mock engine for testing") parser.add_argument( "--engine-input", type=str, default="tcp://127.0.0.1:6000", - help="Engine input address (default: tcp://127.0.0.1:6000)" + help="Engine input address (default: tcp://127.0.0.1:6000)", ) parser.add_argument( "--engine-output", type=str, default="tcp://127.0.0.1:6001", - help="Engine output address (default: tcp://127.0.0.1:6001)" + help="Engine output address (default: tcp://127.0.0.1:6001)", ) parser.add_argument( - "--timeout", - type=float, - default=120.0, - help="Default timeout for engine operations (default: 120.0)" + "--timeout", type=float, default=120.0, help="Default timeout for engine operations (default: 120.0)" ) return parser.parse_args() diff --git a/src/xorl/server/api_server/_state.py b/src/xorl/server/api_server/_state.py index 2260a45f..aa59c76e 100644 --- a/src/xorl/server/api_server/_state.py +++ b/src/xorl/server/api_server/_state.py @@ -1,9 +1,10 @@ """Global API server state. Shared by api_server.py and endpoints.py to avoid circular imports.""" -from typing import Optional, TYPE_CHECKING +from typing import TYPE_CHECKING, Optional from fastapi import HTTPException, status + if TYPE_CHECKING: from xorl.server.api_server.server import APIServer diff --git a/src/xorl/server/api_server/api_types.py b/src/xorl/server/api_server/api_types.py index d1fd8429..7e7c85ff 100644 --- a/src/xorl/server/api_server/api_types.py +++ b/src/xorl/server/api_server/api_types.py @@ -4,7 +4,7 @@ Pydantic type definitions for FastAPI endpoints in the unified API server. """ -from typing import Any, Callable, Dict, List, Literal, Optional, Union +from typing import Any, Dict, List, Literal, Optional, Union from pydantic import BaseModel, ConfigDict, Field, field_serializer, model_validator @@ -352,44 +352,41 @@ class CreateModelResponse(BaseModel): type: Literal["create_model"] = Field(default="create_model", description="Response type identifier") - class UnloadModelRequest(BaseModel): """API request for unloading a model (Tinker-compatible).""" + model_id: str = Field(..., description="Model identifier to unload") type: Literal["unload_model"] = Field(default="unload_model", description="Request type identifier") class UnloadModelResponse(BaseModel): """API response for unloading a model (Tinker-compatible).""" + model_id: str = Field(..., description="Model identifier that was unloaded") type: Optional[Literal["unload_model"]] = Field(default=None, description="Response type identifier") class SessionInfoResponse(BaseModel): """API response with session information for monitoring.""" + registered_models: List[str] = Field(..., description="List of registered model IDs (all ever registered)") active_sessions: int = Field(..., description="Number of active sessions") session_activity: Dict[str, float] = Field( - default_factory=dict, - description="Map of model_id to last activity timestamp" + default_factory=dict, description="Map of model_id to last activity timestamp" ) idle_timeout_seconds: float = Field(..., description="Idle session timeout in seconds") loaded_sampling_adapters: Dict[str, int] = Field( - default_factory=dict, - description="Map of model_id to number of loaded sampling adapters" + default_factory=dict, description="Map of model_id to number of loaded sampling adapters" ) # Training adapter info (from worker's adapter manager) loaded_training_adapters: List[str] = Field( - default_factory=list, - description="List of training adapters currently loaded in GPU memory" + default_factory=list, description="List of training adapters currently loaded in GPU memory" ) max_training_adapters: int = Field( - default=0, - description="Maximum number of training adapters that can be loaded (LRU eviction threshold)" + default=0, description="Maximum number of training adapters that can be loaded (LRU eviction threshold)" ) current_training_adapter: Optional[str] = Field( - default=None, - description="Currently active training adapter (if any)" + default=None, description="Currently active training adapter (if any)" ) @@ -403,13 +400,9 @@ class SaveAdapterStateRequest(BaseModel): model_id: str = Field(..., description="Adapter/session identifier to save") path: Optional[str] = Field( - default=None, - description="Directory to save adapter state to. Auto-generated if not specified." - ) - save_optimizer: bool = Field( - default=True, - description="Whether to save optimizer state for resuming training" + default=None, description="Directory to save adapter state to. Auto-generated if not specified." ) + save_optimizer: bool = Field(default=True, description="Whether to save optimizer state for resuming training") seq_id: Optional[int] = Field(default=None, description="Sequence ID for request ordering") @@ -422,7 +415,6 @@ class SaveAdapterStateResponse(BaseModel): step: int = Field(..., description="Global step at save time") - class RegisterWorkersRequest(BaseModel): """API request for registering inference workers for a model.""" @@ -449,8 +441,9 @@ class SaveWeightsForSamplerRequest(BaseModel): class SaveWeightsForSamplerResponse(BaseModel): """API response for saving weights for sampler.""" - path: str = Field(..., description="Xorl URI for the saved checkpoint (e.g., 'xorl://model-0/sampler_weights/step-100')") - + path: str = Field( + ..., description="Xorl URI for the saved checkpoint (e.g., 'xorl://model-0/sampler_weights/step-100')" + ) # ============================================================================ @@ -461,7 +454,9 @@ class SaveWeightsForSamplerResponse(BaseModel): class CheckpointInfo(BaseModel): """Information about a single checkpoint.""" - checkpoint_id: str = Field(..., description="The checkpoint ID (e.g., 'weights/model_id/name' or 'sampler_weights/name')") + checkpoint_id: str = Field( + ..., description="The checkpoint ID (e.g., 'weights/model_id/name' or 'sampler_weights/name')" + ) checkpoint_type: Literal["training", "sampler"] = Field(..., description="The type of checkpoint") time: str = Field(..., description="ISO format timestamp when the checkpoint was created") path: str = Field(..., description="The xorl:// path to the checkpoint") @@ -528,7 +523,9 @@ class KillSessionResponse(BaseModel): success: bool = Field(..., description="Whether the session was killed successfully") message: str = Field(..., description="Status message") - checkpoint_path: Optional[str] = Field(default=None, description="Path to saved checkpoint (if save_checkpoint=True)") + checkpoint_path: Optional[str] = Field( + default=None, description="Path to saved checkpoint (if save_checkpoint=True)" + ) # ============================================================================ @@ -563,7 +560,9 @@ class InferenceEndpointServerInfo(BaseModel): model_path: Optional[str] = Field(default=None, description="Model path loaded on the server") served_model_name: Optional[str] = Field(default=None, description="Served model name") tp_size: Optional[int] = Field(default=None, description="Tensor parallelism size") - quantization: Optional[str] = Field(default=None, description="Quantization method reported by SGLang (e.g., 'fp8')") + quantization: Optional[str] = Field( + default=None, description="Quantization method reported by SGLang (e.g., 'fp8')" + ) quantization_config: Optional[Dict[str, Any]] = Field( default=None, description="Full HF quantization_config dict auto-detected from model's config.json", diff --git a/src/xorl/server/api_server/endpoints.py b/src/xorl/server/api_server/endpoints.py index e69d4d5f..b92ecd0c 100644 --- a/src/xorl/server/api_server/endpoints.py +++ b/src/xorl/server/api_server/endpoints.py @@ -50,12 +50,12 @@ from xorl.server.protocol.api_orchestrator import OrchestratorRequest from xorl.server.protocol.operations import KillSessionData + logger = logging.getLogger(__name__) router = APIRouter() - # ============================================================================ # Training Endpoints (Two-Phase Pattern) # ============================================================================ @@ -441,19 +441,13 @@ async def unload_model_endpoint(request: UnloadModelRequest, server=Depends(requ UntypedAPIFuture with request_id for polling """ if not server.future_store: - raise HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Future store not initialized" - ) + raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Future store not initialized") model_id = request.model_id # Check if the session exists if model_id not in server.registered_model_ids: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail=f"Model {model_id} not found" - ) + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Model {model_id} not found") async def process_unload_model(request_data: Dict[str, Any]) -> Dict[str, Any]: """Process unload_model request and return result dict.""" @@ -544,10 +538,7 @@ async def kill_session_endpoint(request: KillSessionRequest, server=Depends(requ raise except Exception as e: logger.error(f"Error killing session {model_id}: {e}") - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"Failed to kill session: {e}" - ) + raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to kill session: {e}") @router.get( @@ -572,10 +563,7 @@ async def session_info_endpoint(server=Depends(require_api_server)): session_activity = dict(server.session_last_activity) # Build adapter count map for sampling - loaded_adapters = { - model_id: len(adapters) - for model_id, adapters in server.loaded_sampling_loras.items() - } + loaded_adapters = {model_id: len(adapters) for model_id, adapters in server.loaded_sampling_loras.items()} # Query worker for training adapter info loaded_training_adapters: List[str] = [] @@ -687,7 +675,9 @@ async def list_inference_endpoints_endpoint(server=Depends(require_api_server)): }, tags=["Inference Endpoint Management"], ) -async def remove_inference_endpoint_endpoint(request: RemoveInferenceEndpointRequest, server=Depends(require_api_server)): +async def remove_inference_endpoint_endpoint( + request: RemoveInferenceEndpointRequest, server=Depends(require_api_server) +): """ Remove an inference endpoint from the registry. diff --git a/src/xorl/server/api_server/future_store.py b/src/xorl/server/api_server/future_store.py index a9ea4f5a..780c7073 100644 --- a/src/xorl/server/api_server/future_store.py +++ b/src/xorl/server/api_server/future_store.py @@ -27,6 +27,7 @@ from enum import Enum from typing import Any, Callable, Coroutine, Dict, List, Optional, Set + logger = logging.getLogger(__name__) @@ -176,10 +177,7 @@ def __init__( self._queue_state = "active" self._queue_state_reason: Optional[str] = None - logger.info( - f"FutureStore initialized: default_ttl={default_ttl}s, " - f"max_concurrent={max_concurrent}" - ) + logger.info(f"FutureStore initialized: default_ttl={default_ttl}s, max_concurrent={max_concurrent}") async def start(self): """Start the store and background cleanup task.""" @@ -246,9 +244,7 @@ async def create( self._by_model[model_id].add(request_id) # Start processing in background - asyncio.create_task( - self._process(request_id, process_fn, request_data) - ) + asyncio.create_task(self._process(request_id, process_fn, request_data)) logger.debug( f"Future created: request_id={request_id}, model_id={model_id}, " @@ -500,7 +496,8 @@ async def _cleanup_expired(self): """Remove expired entries from the store.""" async with self._lock: to_delete = [ - request_id for request_id, entry in list(self._entries.items()) + request_id + for request_id, entry in list(self._entries.items()) if entry.is_terminal() and entry.is_expired() ] for request_id in to_delete: @@ -528,10 +525,7 @@ def get_stats(self) -> Dict[str, Any]: "failed": 0, "expired": 0, }, - "by_model": { - model_id: len(request_ids) - for model_id, request_ids in self._by_model.items() - }, + "by_model": {model_id: len(request_ids) for model_id, request_ids in self._by_model.items()}, } for entry in self._entries.values(): diff --git a/src/xorl/server/api_server/health.py b/src/xorl/server/api_server/health.py index b7433b0e..317e35c0 100644 --- a/src/xorl/server/api_server/health.py +++ b/src/xorl/server/api_server/health.py @@ -11,6 +11,7 @@ from xorl.server.api_server.api_types import HealthCheckResponse from xorl.server.protocol.api_orchestrator import OrchestratorRequest, RequestType + logger = logging.getLogger(__name__) diff --git a/src/xorl/server/api_server/inference_endpoints.py b/src/xorl/server/api_server/inference_endpoints.py index 3f7d6129..1afbea0b 100644 --- a/src/xorl/server/api_server/inference_endpoints.py +++ b/src/xorl/server/api_server/inference_endpoints.py @@ -30,6 +30,7 @@ from xorl.server.protocol.api_orchestrator import OrchestratorRequest from xorl.server.protocol.operations import SyncWeightsData + logger = logging.getLogger(__name__) @@ -63,6 +64,7 @@ def _detect_quantization_from_hf_config(model_path: str) -> dict | None: if config_dict is None: try: from huggingface_hub import hf_hub_download + cached_path = hf_hub_download(model_path, "config.json") with open(cached_path) as f: config_dict = json.load(f) @@ -283,7 +285,11 @@ async def add_inference_endpoint(self, request: AddInferenceEndpointRequest) -> existing_info = self.inference_endpoints[0].server_info if existing_info is not None: mismatches = [] - if existing_info.model_path and server_info.model_path and existing_info.model_path != server_info.model_path: + if ( + existing_info.model_path + and server_info.model_path + and existing_info.model_path != server_info.model_path + ): mismatches.append(f"model_path: {existing_info.model_path} vs {server_info.model_path}") if existing_info.quantization != server_info.quantization: mismatches.append(f"quantization: {existing_info.quantization} vs {server_info.quantization}") @@ -564,7 +570,9 @@ async def sync_inference_weights(self, request: SyncInferenceWeightsRequest) -> raise except Exception as e: logger.error(f"Sync inference weights failed: {e}", exc_info=True) - raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Sync inference weights failed: {e}") + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Sync inference weights failed: {e}" + ) # ========================================================================= # Sampling Session Management (LoRA Adapter Loading) @@ -602,11 +610,11 @@ def _resolve_model_path(self, model_path: str) -> tuple[str, str]: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid xorl:// URI format for sampler weights: {model_path}. " - f"Expected: xorl://model_id/sampler_weights/adapter_name" + f"Expected: xorl://model_id/sampler_weights/adapter_name", ) elif model_path.startswith("sampler_weights/"): # Format: sampler_weights/adapter_name - lora_name = model_path[len("sampler_weights/"):] + lora_name = model_path[len("sampler_weights/") :] else: # Just the adapter name lora_name = model_path @@ -872,10 +880,7 @@ def _track_adapter(self, lora_name: str, lora_path: str, model_id: str = "defaul logger.info(f"LoRA adapter '{lora_name}' added to tracking list (count={len(adapters)})") return False - async def create_sampling_session( - self, - request: CreateSamplingSessionRequest - ) -> CreateSamplingSessionResponse: + async def create_sampling_session(self, request: CreateSamplingSessionRequest) -> CreateSamplingSessionResponse: """ Create a sampling session by loading a LoRA adapter on inference workers. @@ -910,7 +915,9 @@ async def create_sampling_session( if len(adapters) > self.max_adapters_per_model: # We just added one, so if we're over capacity, remove the oldest (index 0) oldest_name, oldest_path = adapters[0] - logger.info(f"Max LoRA adapters exceeded ({self.max_adapters_per_model}), unloading oldest: {oldest_name}") + logger.info( + f"Max LoRA adapters exceeded ({self.max_adapters_per_model}), unloading oldest: {oldest_name}" + ) await self._unload_lora_on_inference_endpoints(oldest_name) adapters.pop(0) diff --git a/src/xorl/server/api_server/orchestrator_client.py b/src/xorl/server/api_server/orchestrator_client.py index 39418fe0..993c7599 100644 --- a/src/xorl/server/api_server/orchestrator_client.py +++ b/src/xorl/server/api_server/orchestrator_client.py @@ -30,9 +30,10 @@ import logging import time from typing import Optional + import zmq.asyncio -from xorl.server.protocol.api_orchestrator import OrchestratorRequest, OrchestratorOutputs, RequestType +from xorl.server.protocol.api_orchestrator import OrchestratorOutputs, OrchestratorRequest, RequestType from xorl.server.protocol.operations import AbortData from xorl.server.utils.zmq_channels import AsyncPullChannel, AsyncRouterChannel @@ -87,9 +88,7 @@ def __init__( self._last_request_time = 0.0 self._last_output_time = 0.0 - logger.info( - f"OrchestratorClient initialized: input={input_addr}, output={output_addr}" - ) + logger.info(f"OrchestratorClient initialized: input={input_addr}, output={output_addr}") async def start(self): """Start the client and connect sockets.""" @@ -219,10 +218,7 @@ async def get_output( try: if timeout is not None: - output = await asyncio.wait_for( - self.output_queue.get(), - timeout=timeout - ) + output = await asyncio.wait_for(self.output_queue.get(), timeout=timeout) else: output = await self.output_queue.get() @@ -283,8 +279,7 @@ async def process_outputs_socket(self): logger.warning(f"Future for request {output.request_id} already done") else: logger.warning( - f"Received output for unknown request {output.request_id} " - f"(may have timed out or was cancelled)" + f"Received output for unknown request {output.request_id} (may have timed out or was cancelled)" ) # Still put in output queue as fallback await self.output_queue.put(output) @@ -385,5 +380,3 @@ async def __aenter__(self): async def __aexit__(self, exc_type, exc_val, exc_tb): """Async context manager exit.""" await self.stop() - - diff --git a/src/xorl/server/api_server/server.py b/src/xorl/server/api_server/server.py index 78e28a56..2ee23132 100644 --- a/src/xorl/server/api_server/server.py +++ b/src/xorl/server/api_server/server.py @@ -38,28 +38,27 @@ import logging import os import time -import uvicorn from contextlib import asynccontextmanager from typing import Any, Dict, List, Optional +import uvicorn from fastapi import FastAPI, HTTPException, status +from xorl.server.api_server.api_types import ( + InferenceEndpoint, + LossFnOutput, + TensorData, +) from xorl.server.api_server.endpoints import router +from xorl.server.api_server.future_store import FutureStore from xorl.server.api_server.health import HealthMixin from xorl.server.api_server.inference_endpoints import InferenceEndpointsMixin +from xorl.server.api_server.orchestrator_client import OrchestratorClient from xorl.server.api_server.training_ops import TrainingOpsMixin from xorl.server.api_server.utils import ( MODEL_ID_NOT_REGISTERED_ERROR, - validate_model_id, ) from xorl.server.api_server.weights import WeightsMixin -from xorl.server.api_server.api_types import ( - InferenceEndpoint, - LossFnOutput, - TensorData, -) -from xorl.server.api_server.orchestrator_client import OrchestratorClient -from xorl.server.api_server.future_store import FutureStore from xorl.server.protocol.api_orchestrator import OrchestratorRequest from xorl.server.protocol.operations import KillSessionData @@ -248,7 +247,9 @@ def _build_loss_fn_outputs(result: Dict[str, Any]): outputs.append( LossFnOutput( logprobs=TensorData(data=logprobs, dtype="float32", shape=[len(logprobs)]), - elementwise_loss=TensorData(data=elementwise_loss, dtype="float32", shape=[len(elementwise_loss)]), + elementwise_loss=TensorData( + data=elementwise_loss, dtype="float32", shape=[len(elementwise_loss)] + ), k3=k3_val, ) ) @@ -359,10 +360,7 @@ async def _cleanup_idle_sessions(self) -> None: ] for model_id, last_activity in idle_model_ids: - logger.info( - f"Cleaning up idle session: {model_id} " - f"(idle for {current_time - last_activity:.0f}s)" - ) + logger.info(f"Cleaning up idle session: {model_id} (idle for {current_time - last_activity:.0f}s)") await self._cleanup_session(model_id) except asyncio.CancelledError: diff --git a/src/xorl/server/api_server/training_ops.py b/src/xorl/server/api_server/training_ops.py index 05971488..76b1f1a0 100644 --- a/src/xorl/server/api_server/training_ops.py +++ b/src/xorl/server/api_server/training_ops.py @@ -5,13 +5,10 @@ import logging import math import time -from typing import Any, Dict, List +from typing import Any, Dict from fastapi import HTTPException, status -from xorl.server.api_server.utils import ( - validate_model_id, -) from xorl.server.api_server.api_types import ( ForwardBackwardRequest, ForwardBackwardResponse, @@ -30,9 +27,13 @@ from xorl.server.api_server.future_store import ( FutureStatus, ) -from xorl.server.protocol.api_orchestrator import OrchestratorRequest, RequestType +from xorl.server.api_server.utils import ( + validate_model_id, +) +from xorl.server.protocol.api_orchestrator import OrchestratorRequest from xorl.server.protocol.operations import ModelPassData, OptimStepData + logger = logging.getLogger(__name__) @@ -78,19 +79,13 @@ async def retrieve_future(self, request: FutureRetrieveRequest, timeout: float = HTTPException: 404 if request_id not found, 503 if store not initialized """ if not self.future_store: - raise HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Future store not initialized" - ) + raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Future store not initialized") # Use long polling: wait for result with timeout entry = await self.future_store.wait_for_result(request.request_id, timeout=timeout) if entry is None: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail=f"Request {request.request_id} not found" - ) + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Request {request.request_id} not found") # Return appropriate response based on status if entry.status == FutureStatus.PENDING: @@ -151,10 +146,7 @@ async def _submit_async(self, request, request_type: str, handler_method: str) - """ self._require_engine() if not self.future_store: - raise HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Future store not initialized" - ) + raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Future store not initialized") model_id = validate_model_id(request.model_id) self.validate_model_id(model_id) @@ -307,7 +299,9 @@ async def forward_backward(self, request: ForwardBackwardRequest) -> ForwardBack raise except Exception as e: logger.error(f"Forward-backward failed: {e}", exc_info=True) - raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Forward-backward failed: {e}") + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Forward-backward failed: {e}" + ) async def forward(self, request: ForwardRequest) -> ForwardResponse: """ @@ -455,4 +449,3 @@ async def optim_step(self, request: OptimStepRequest) -> OptimStepResponse: except Exception as e: logger.error(f"Optimizer step failed: {e}", exc_info=True) raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Optimizer step failed: {e}") - diff --git a/src/xorl/server/api_server/weights.py b/src/xorl/server/api_server/weights.py index bc679578..34605044 100644 --- a/src/xorl/server/api_server/weights.py +++ b/src/xorl/server/api_server/weights.py @@ -7,11 +7,10 @@ import shutil import time from datetime import datetime -from typing import Any, Dict, List, Optional +from typing import List from fastapi import HTTPException, status -from xorl.server.api_server.utils import validate_model_id from xorl.server.api_server.api_types import ( CheckpointInfo, Cursor, @@ -28,6 +27,7 @@ TrainingRun, TrainingRunsResponse, ) +from xorl.server.api_server.utils import validate_model_id from xorl.server.protocol.api_orchestrator import OrchestratorRequest from xorl.server.protocol.operations import LoadStateData, SaveFullWeightsData, SaveLoraOnlyData, SaveStateData from xorl.server.utils.storage import ( @@ -35,6 +35,7 @@ check_storage_limit, ) + logger = logging.getLogger(__name__) @@ -214,7 +215,10 @@ async def save_weights(self, request: SaveWeightsRequest) -> SaveWeightsResponse weights_model_dir = os.path.normpath(os.path.abspath(os.path.join(self.output_dir, "weights", model_id))) saved_path_normalized = os.path.normpath(os.path.abspath(saved_path)) - if saved_path_normalized.startswith(weights_model_dir + os.sep) or saved_path_normalized == weights_model_dir: + if ( + saved_path_normalized.startswith(weights_model_dir + os.sep) + or saved_path_normalized == weights_model_dir + ): # Extract just the checkpoint name (relative to model_id directory) checkpoint_name_from_path = os.path.relpath(saved_path_normalized, weights_model_dir) else: @@ -713,8 +717,10 @@ async def save_weights_for_sampler(self, request: SaveWeightsForSamplerRequest) engine_request = OrchestratorRequest( operation="save_full_weights", payload=SaveFullWeightsData( - output_path=save_path, dtype="bfloat16", - base_model_path=base_model_path, model_id=request.model_id, + output_path=save_path, + dtype="bfloat16", + base_model_path=base_model_path, + model_id=request.model_id, ), ) @@ -762,5 +768,6 @@ async def save_weights_for_sampler(self, request: SaveWeightsForSamplerRequest) raise except Exception as e: logger.error(f"Save weights for sampler failed: {e}", exc_info=True) - raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Save weights for sampler failed: {e}") - + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Save weights for sampler failed: {e}" + ) diff --git a/src/xorl/server/backend/__init__.py b/src/xorl/server/backend/__init__.py index 5961d1ac..7efa1e49 100644 --- a/src/xorl/server/backend/__init__.py +++ b/src/xorl/server/backend/__init__.py @@ -1,5 +1,6 @@ from xorl.server.backend.base import Backend -from xorl.server.backend.remote import RemoteBackend from xorl.server.backend.dummy import DummyBackend +from xorl.server.backend.remote import RemoteBackend + __all__ = ["Backend", "RemoteBackend", "DummyBackend"] diff --git a/src/xorl/server/backend/dummy.py b/src/xorl/server/backend/dummy.py index 5f6d1b43..6d33cc6c 100644 --- a/src/xorl/server/backend/dummy.py +++ b/src/xorl/server/backend/dummy.py @@ -7,11 +7,10 @@ import logging import random -import time -from typing import Any, Dict, List, Optional from xorl.server.backend.base import Backend + logger = logging.getLogger(__name__) @@ -39,32 +38,29 @@ def _maybe_fail(self, operation: str): if random.random() < self.failure_rate: raise RuntimeError(f"Simulated failure in {operation} (failure_rate={self.failure_rate})") - async def forward_backward(self, batches, loss_fn="causallm_loss", loss_fn_params=None, - model_id=None, routed_experts=None, request_id=None): + async def forward_backward( + self, batches, loss_fn="causallm_loss", loss_fn_params=None, model_id=None, routed_experts=None, request_id=None + ): self._maybe_fail("forward_backward") - valid_tokens = sum( - len(batch.get("input_ids", [])) for batch in (batches or []) - ) + valid_tokens = sum(len(batch.get("input_ids", [])) for batch in (batches or [])) return { "total_loss": random.uniform(0.5, 5.0), "global_valid_tokens": valid_tokens, "num_batches": len(batches or []), } - async def forward(self, batches, loss_fn="causallm_loss", loss_fn_params=None, - model_id=None, request_id=None): + async def forward(self, batches, loss_fn="causallm_loss", loss_fn_params=None, model_id=None, request_id=None): self._maybe_fail("forward") - valid_tokens = sum( - len(batch.get("input_ids", [])) for batch in (batches or []) - ) + valid_tokens = sum(len(batch.get("input_ids", [])) for batch in (batches or [])) return { "total_loss": random.uniform(0.5, 5.0), "global_valid_tokens": valid_tokens, "num_batches": len(batches or []), } - async def optim_step(self, lr, gradient_clip=None, beta1=None, beta2=None, - eps=None, model_id=None, request_id=None): + async def optim_step( + self, lr, gradient_clip=None, beta1=None, beta2=None, eps=None, model_id=None, request_id=None + ): self._maybe_fail("optim_step") self._step += 1 return { @@ -73,13 +69,13 @@ async def optim_step(self, lr, gradient_clip=None, beta1=None, beta2=None, "learning_rate": lr, } - async def save_state(self, checkpoint_path=None, save_optimizer=True, - use_timestamp=False, model_id=None, request_id=None): + async def save_state( + self, checkpoint_path=None, save_optimizer=True, use_timestamp=False, model_id=None, request_id=None + ): self._maybe_fail("save_state") return {"checkpoint_path": checkpoint_path or "/tmp/dummy_ckpt", "success": True} - async def load_state(self, checkpoint_path=None, load_optimizer=True, - model_id=None, request_id=None): + async def load_state(self, checkpoint_path=None, load_optimizer=True, model_id=None, request_id=None): self._maybe_fail("load_state") return {"checkpoint_path": checkpoint_path or "/tmp/dummy_ckpt", "success": True} @@ -87,8 +83,9 @@ async def save_lora_only(self, lora_path=None, model_id=None, request_id=None): self._maybe_fail("save_lora_only") return {"lora_path": lora_path or "/tmp/dummy_lora", "success": True} - async def save_full_weights(self, output_path=None, dtype="bfloat16", - base_model_path=None, model_id=None, request_id=None): + async def save_full_weights( + self, output_path=None, dtype="bfloat16", base_model_path=None, model_id=None, request_id=None + ): self._maybe_fail("save_full_weights") return {"output_path": output_path or "/tmp/dummy_weights", "success": True, "num_shards": 1} @@ -98,20 +95,26 @@ async def sleep(self, request_id=None): async def wake_up(self, request_id=None): return {"status": "awake", "load_time": 0.0} - async def sync_inference_weights(self, endpoints, master_address="localhost", - master_port=29600, request_id=None, **kwargs): - return {"success": True, "message": "dummy sync", "transfer_time": 0.0, - "total_bytes": 0, "num_parameters": 0, "num_buckets": 0, "endpoint_results": []} + async def sync_inference_weights( + self, endpoints, master_address="localhost", master_port=29600, request_id=None, **kwargs + ): + return { + "success": True, + "message": "dummy sync", + "transfer_time": 0.0, + "total_bytes": 0, + "num_parameters": 0, + "num_buckets": 0, + "endpoint_results": [], + } async def register_adapter(self, model_id="default", lr=1e-5, request_id=None): return {"model_id": model_id, "lr": lr, "registered": True} - async def save_adapter_state(self, model_id="default", path=None, - save_optimizer=True, request_id=None): + async def save_adapter_state(self, model_id="default", path=None, save_optimizer=True, request_id=None): return {"success": True} - async def load_adapter_state(self, model_id="default", path=None, - load_optimizer=True, lr=None, request_id=None): + async def load_adapter_state(self, model_id="default", path=None, load_optimizer=True, lr=None, request_id=None): return {"success": True} async def get_adapter_info(self, request_id=None): diff --git a/src/xorl/server/backend/remote.py b/src/xorl/server/backend/remote.py index f46ca3ed..51a90e0d 100644 --- a/src/xorl/server/backend/remote.py +++ b/src/xorl/server/backend/remote.py @@ -7,17 +7,12 @@ import asyncio import logging -from typing import Any, Dict, List, Optional +from typing import Any, Dict, Optional -from xorl.server.protocol.orchestrator_runner import ( - RunnerDispatchCommand, - RunnerAck, - RunnerReady, - RunnerResponse, - deserialize_message, - serialize_message, -) +from xorl.server.backend.base import Backend from xorl.server.protocol.operations import ( + LOAD_STATE_TIMEOUT, + SAVE_STATE_TIMEOUT, AdapterStateData, EmptyData, KillSessionData, @@ -30,10 +25,17 @@ SaveStateData, SyncWeightsData, ) -from xorl.server.backend.base import Backend -from xorl.server.protocol.operations import LOAD_STATE_TIMEOUT, SAVE_STATE_TIMEOUT +from xorl.server.protocol.orchestrator_runner import ( + RunnerAck, + RunnerDispatchCommand, + RunnerReady, + RunnerResponse, + deserialize_message, + serialize_message, +) from xorl.server.utils.zmq_channels import AsyncDealerChannel + logger = logging.getLogger(__name__) @@ -198,7 +200,9 @@ async def _send_and_receive( f"success={response.success}, time={response.execution_time:.3f}s" ) except asyncio.TimeoutError: - raise RuntimeError(f"Timeout waiting for response (request {request.message_id}, timeout={response_timeout}s)") + raise RuntimeError( + f"Timeout waiting for response (request {request.message_id}, timeout={response_timeout}s)" + ) if not response.success: raise RuntimeError(f"Operation failed: {response.error or 'Unknown error'}") @@ -209,8 +213,9 @@ async def _send_and_receive( # Backend Operations # ======================================================================== - async def _execute(self, operation: str, payload, request_id: Optional[str] = None, - timeout: Optional[float] = None) -> Dict[str, Any]: + async def _execute( + self, operation: str, payload, request_id: Optional[str] = None, timeout: Optional[float] = None + ) -> Dict[str, Any]: """Generic execute: create request, send, return result dict.""" request = RunnerDispatchCommand.create(operation, payload, request_id=request_id) response = await self._send_and_receive(request, timeout=timeout) @@ -219,56 +224,101 @@ async def _execute(self, operation: str, payload, request_id: Optional[str] = No result["execution_time"] = response.execution_time return result - async def forward_backward(self, batches, loss_fn="causallm_loss", loss_fn_params=None, - model_id=None, routed_experts=None, request_id=None): - return await self._execute("forward_backward", ModelPassData( - batches=batches, loss_fn=loss_fn, loss_fn_params=loss_fn_params, - model_id=model_id, routed_experts=routed_experts, - ), request_id=request_id) - - async def forward(self, batches, loss_fn="causallm_loss", loss_fn_params=None, - model_id=None, request_id=None): - return await self._execute("forward", ModelPassData( - batches=batches, loss_fn=loss_fn, loss_fn_params=loss_fn_params, - model_id=model_id, - ), request_id=request_id) - - async def optim_step(self, lr, gradient_clip=None, beta1=None, beta2=None, - eps=None, model_id=None, request_id=None): - return await self._execute("optim_step", OptimStepData( - lr=lr, gradient_clip=gradient_clip, beta1=beta1, beta2=beta2, - eps=eps, model_id=model_id, - ), request_id=request_id) - - async def save_state(self, checkpoint_path=None, save_optimizer=True, use_timestamp=False, - model_id=None, request_id=None): - - return await self._execute("save_state", SaveStateData( - checkpoint_path=checkpoint_path, save_optimizer=save_optimizer, - use_timestamp=use_timestamp, model_id=model_id, - ), request_id=request_id, timeout=SAVE_STATE_TIMEOUT) - - async def load_state(self, checkpoint_path=None, load_optimizer=True, - model_id=None, request_id=None): - - return await self._execute("load_state", LoadStateData( - checkpoint_path=checkpoint_path, load_optimizer=load_optimizer, - model_id=model_id, - ), request_id=request_id, timeout=LOAD_STATE_TIMEOUT) + async def forward_backward( + self, batches, loss_fn="causallm_loss", loss_fn_params=None, model_id=None, routed_experts=None, request_id=None + ): + return await self._execute( + "forward_backward", + ModelPassData( + batches=batches, + loss_fn=loss_fn, + loss_fn_params=loss_fn_params, + model_id=model_id, + routed_experts=routed_experts, + ), + request_id=request_id, + ) + + async def forward(self, batches, loss_fn="causallm_loss", loss_fn_params=None, model_id=None, request_id=None): + return await self._execute( + "forward", + ModelPassData( + batches=batches, + loss_fn=loss_fn, + loss_fn_params=loss_fn_params, + model_id=model_id, + ), + request_id=request_id, + ) + + async def optim_step( + self, lr, gradient_clip=None, beta1=None, beta2=None, eps=None, model_id=None, request_id=None + ): + return await self._execute( + "optim_step", + OptimStepData( + lr=lr, + gradient_clip=gradient_clip, + beta1=beta1, + beta2=beta2, + eps=eps, + model_id=model_id, + ), + request_id=request_id, + ) + + async def save_state( + self, checkpoint_path=None, save_optimizer=True, use_timestamp=False, model_id=None, request_id=None + ): + return await self._execute( + "save_state", + SaveStateData( + checkpoint_path=checkpoint_path, + save_optimizer=save_optimizer, + use_timestamp=use_timestamp, + model_id=model_id, + ), + request_id=request_id, + timeout=SAVE_STATE_TIMEOUT, + ) + + async def load_state(self, checkpoint_path=None, load_optimizer=True, model_id=None, request_id=None): + return await self._execute( + "load_state", + LoadStateData( + checkpoint_path=checkpoint_path, + load_optimizer=load_optimizer, + model_id=model_id, + ), + request_id=request_id, + timeout=LOAD_STATE_TIMEOUT, + ) async def save_lora_only(self, lora_path=None, model_id=None, request_id=None): - - return await self._execute("save_lora_only", SaveLoraOnlyData( - lora_path=lora_path, model_id=model_id, - ), request_id=request_id, timeout=SAVE_STATE_TIMEOUT) - - async def save_full_weights(self, output_path=None, dtype="bfloat16", - base_model_path=None, model_id=None, request_id=None): - - return await self._execute("save_full_weights", SaveFullWeightsData( - output_path=output_path, dtype=dtype, base_model_path=base_model_path, - model_id=model_id, - ), request_id=request_id, timeout=SAVE_STATE_TIMEOUT) + return await self._execute( + "save_lora_only", + SaveLoraOnlyData( + lora_path=lora_path, + model_id=model_id, + ), + request_id=request_id, + timeout=SAVE_STATE_TIMEOUT, + ) + + async def save_full_weights( + self, output_path=None, dtype="bfloat16", base_model_path=None, model_id=None, request_id=None + ): + return await self._execute( + "save_full_weights", + SaveFullWeightsData( + output_path=output_path, + dtype=dtype, + base_model_path=base_model_path, + model_id=model_id, + ), + request_id=request_id, + timeout=SAVE_STATE_TIMEOUT, + ) async def sleep(self, request_id=None): return await self._execute("sleep", EmptyData(), request_id=request_id, timeout=30.0) @@ -276,43 +326,87 @@ async def sleep(self, request_id=None): async def wake_up(self, request_id=None): return await self._execute("wake_up", EmptyData(), request_id=request_id, timeout=30.0) - async def sync_inference_weights(self, endpoints, master_address="localhost", - master_port=29600, group_name="weight_sync_group", - buffer_size_mb=1024, sync_method="nccl_broadcast", - flush_cache=True, pause_mode="retract", - weight_version=None, quantization=None, - request_id=None): - return await self._execute("sync_inference_weights", SyncWeightsData( - endpoints=endpoints, master_address=master_address, master_port=master_port, - group_name=group_name, buffer_size_mb=buffer_size_mb, sync_method=sync_method, - flush_cache=flush_cache, pause_mode=pause_mode, weight_version=weight_version, - quantization=quantization, - ), request_id=request_id, timeout=600.0) + async def sync_inference_weights( + self, + endpoints, + master_address="localhost", + master_port=29600, + group_name="weight_sync_group", + buffer_size_mb=1024, + sync_method="nccl_broadcast", + flush_cache=False, + pause_mode="retract", + weight_version=None, + quantization=None, + request_id=None, + ): + return await self._execute( + "sync_inference_weights", + SyncWeightsData( + endpoints=endpoints, + master_address=master_address, + master_port=master_port, + group_name=group_name, + buffer_size_mb=buffer_size_mb, + sync_method=sync_method, + flush_cache=flush_cache, + pause_mode=pause_mode, + weight_version=weight_version, + quantization=quantization, + ), + request_id=request_id, + timeout=600.0, + ) async def register_adapter(self, model_id="default", lr=1e-5, request_id=None): - return await self._execute("register_adapter", RegisterAdapterData( - model_id=model_id, lr=lr, - ), request_id=request_id, timeout=60.0) - - async def save_adapter_state(self, model_id="default", path=None, - save_optimizer=True, request_id=None): - return await self._execute("save_adapter_state", AdapterStateData( - model_id=model_id, path=path, save_optimizer=save_optimizer, - ), request_id=request_id, timeout=600.0) - - async def load_adapter_state(self, model_id="default", path=None, - load_optimizer=True, lr=None, request_id=None): - return await self._execute("load_adapter_state", AdapterStateData( - model_id=model_id, path=path, load_optimizer=load_optimizer, lr=lr, - ), request_id=request_id, timeout=600.0) + return await self._execute( + "register_adapter", + RegisterAdapterData( + model_id=model_id, + lr=lr, + ), + request_id=request_id, + timeout=60.0, + ) + + async def save_adapter_state(self, model_id="default", path=None, save_optimizer=True, request_id=None): + return await self._execute( + "save_adapter_state", + AdapterStateData( + model_id=model_id, + path=path, + save_optimizer=save_optimizer, + ), + request_id=request_id, + timeout=600.0, + ) + + async def load_adapter_state(self, model_id="default", path=None, load_optimizer=True, lr=None, request_id=None): + return await self._execute( + "load_adapter_state", + AdapterStateData( + model_id=model_id, + path=path, + load_optimizer=load_optimizer, + lr=lr, + ), + request_id=request_id, + timeout=600.0, + ) async def get_adapter_info(self, request_id=None): return await self._execute("get_adapter_info", EmptyData(), request_id=request_id, timeout=60.0) async def kill_session(self, model_id="default", save_checkpoint=True, request_id=None): - return await self._execute("kill_session", KillSessionData( - model_id=model_id, save_checkpoint=save_checkpoint, - ), request_id=request_id, timeout=120.0) + return await self._execute( + "kill_session", + KillSessionData( + model_id=model_id, + save_checkpoint=save_checkpoint, + ), + request_id=request_id, + timeout=120.0, + ) async def health_check(self, request_id=None): return await self._execute("health_check", EmptyData(), request_id=request_id, timeout=5.0) diff --git a/src/xorl/server/launcher.py b/src/xorl/server/launcher.py index 49e1d55f..5eaf5c58 100644 --- a/src/xorl/server/launcher.py +++ b/src/xorl/server/launcher.py @@ -45,9 +45,7 @@ # Setup logging logging.basicConfig( - level=logging.INFO, - format="[%(levelname)s][%(name)s] %(asctime)s >> %(message)s", - datefmt="%H:%M:%S" + level=logging.INFO, format="[%(levelname)s][%(name)s] %(asctime)s >> %(message)s", datefmt="%H:%M:%S" ) logger = logging.getLogger(__name__) @@ -63,11 +61,12 @@ class RetrieveFutureFilter(logging.Filter): every second until results are ready. This filter suppresses those logs to avoid cluttering the output. """ + def filter(self, record: logging.LogRecord) -> bool: # Filter out retrieve_future access log entries - if hasattr(record, 'getMessage'): + if hasattr(record, "getMessage"): msg = record.getMessage() - if '/api/v1/retrieve_future' in msg: + if "/api/v1/retrieve_future" in msg: return False return True @@ -83,6 +82,7 @@ def configure_uvicorn_logging(): # Port Finding Utilities # ============================================================================ + def find_free_port(start_port: int = 50000, max_attempts: int = 10000) -> int: """ Find a free port by randomly picking from a range. @@ -98,13 +98,14 @@ def find_free_port(start_port: int = 50000, max_attempts: int = 10000) -> int: RuntimeError: If no free port found """ import random + end_port = min(start_port + max_attempts, 60000) ports = list(range(start_port, end_port)) random.shuffle(ports) for port in ports: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock: try: - sock.bind(('', port)) + sock.bind(("", port)) sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) return port except OSError: @@ -138,6 +139,7 @@ def find_free_ports(count: int, start_port: int = 50000) -> List[int]: # Engine Core Process # ============================================================================ + def run_orchestrator( input_addr: str, output_addr: str, @@ -175,11 +177,8 @@ def run_orchestrator( level=getattr(logging, log_level), format="[%(levelname)s][ENGINE] %(asctime)s >> %(message)s", datefmt="%H:%M:%S", - handlers=[ - logging.FileHandler(log_file, mode='w'), - logging.StreamHandler(sys.stdout) - ], - force=True + handlers=[logging.FileHandler(log_file, mode="w"), logging.StreamHandler(sys.stdout)], + force=True, ) logger = logging.getLogger("Orchestrator") logger.info(f"Engine Core logging to: {os.path.abspath(log_file)}") @@ -224,7 +223,7 @@ def run_orchestrator( except Exception as e: logger.error(f"Error in engine core: {e}", exc_info=True) finally: - if 'engine' in locals(): + if "engine" in locals(): engine.stop() logger.info("Engine Core stopped") @@ -233,6 +232,7 @@ def run_orchestrator( # API Server Process # ============================================================================ + def run_api_server( host: str, port: int, @@ -268,9 +268,7 @@ def run_api_server( # Setup logging for this process logging.basicConfig( - level=getattr(logging, log_level), - format="[%(levelname)s][API] %(asctime)s >> %(message)s", - datefmt="%H:%M:%S" + level=getattr(logging, log_level), format="[%(levelname)s][API] %(asctime)s >> %(message)s", datefmt="%H:%M:%S" ) logger = logging.getLogger("APIServer") @@ -287,10 +285,9 @@ def run_api_server( try: # Import the FastAPI app from api_server module - from xorl.server.api_server.server import app - # Update the global state shared between api_server.py and endpoints.py import xorl.server.api_server._state as _state_module + from xorl.server.api_server.server import app # Override the lifespan to use our addresses @asynccontextmanager @@ -345,6 +342,7 @@ async def custom_lifespan(app): # Configuration Parsing # ============================================================================ + def load_server_arguments(config_path: str, overrides: Optional[Dict[str, any]] = None) -> ServerArguments: """ Load ServerArguments from a YAML configuration file. @@ -360,73 +358,78 @@ def load_server_arguments(config_path: str, overrides: Optional[Dict[str, any]] Returns: ServerArguments instance with all fields populated """ - with open(config_path, 'r') as f: + with open(config_path, "r") as f: config = yaml.safe_load(f) if not config: raise ValueError(f"Empty config file: {config_path}") from dataclasses import fields + valid_fields = {f.name for f in fields(ServerArguments)} # Check if this is a nested config (has model/train/worker sections) - if 'model' in config and isinstance(config['model'], dict): + if "model" in config and isinstance(config["model"], dict): # Nested config - flatten it flat_config = {} # Model section - model_config = config.get('model', {}) - flat_config['model_path'] = model_config.get('model_path') - flat_config['model_name'] = model_config.get('model_name') - flat_config['config_path'] = model_config.get('config_path') - flat_config['tokenizer_path'] = model_config.get('tokenizer_path') - flat_config['attn_implementation'] = model_config.get('attn_implementation', 'flash_attention_3') - flat_config['moe_implementation'] = model_config.get('moe_implementation') + model_config = config.get("model", {}) + flat_config["model_path"] = model_config.get("model_path") + flat_config["model_name"] = model_config.get("model_name") + flat_config["config_path"] = model_config.get("config_path") + flat_config["tokenizer_path"] = model_config.get("tokenizer_path") + flat_config["attn_implementation"] = model_config.get("attn_implementation", "flash_attention_3") + flat_config["moe_implementation"] = model_config.get("moe_implementation") # Train section - train_config = config.get('train', {}) - flat_config['data_parallel_mode'] = train_config.get('data_parallel_mode', 'fsdp2') - flat_config['data_parallel_shard_size'] = train_config.get('data_parallel_shard_size', 1) - flat_config['data_parallel_replicate_size'] = train_config.get('data_parallel_replicate_size', 1) - flat_config['ulysses_parallel_size'] = train_config.get('ulysses_parallel_size', 1) - flat_config['expert_parallel_size'] = train_config.get('expert_parallel_size', 1) - flat_config['enable_mixed_precision'] = train_config.get('enable_mixed_precision', True) - flat_config['enable_gradient_checkpointing'] = train_config.get('enable_gradient_checkpointing', True) - flat_config['enable_full_shard'] = train_config.get('enable_full_shard', True) - flat_config['enable_activation_offload'] = train_config.get('enable_activation_offload', False) - flat_config['init_device'] = train_config.get('init_device', 'meta') - flat_config['load_checkpoint_path'] = train_config.get('load_checkpoint_path', '') - flat_config['ckpt_manager'] = train_config.get('ckpt_manager', 'dcp') - flat_config['log_level'] = train_config.get('log_level', 'INFO') + train_config = config.get("train", {}) + flat_config["data_parallel_mode"] = train_config.get("data_parallel_mode", "fsdp2") + flat_config["data_parallel_shard_size"] = train_config.get("data_parallel_shard_size", 1) + flat_config["data_parallel_replicate_size"] = train_config.get("data_parallel_replicate_size", 1) + flat_config["ulysses_parallel_size"] = train_config.get("ulysses_parallel_size", 1) + flat_config["expert_parallel_size"] = train_config.get("expert_parallel_size", 1) + flat_config["enable_mixed_precision"] = train_config.get("enable_mixed_precision", True) + flat_config["enable_gradient_checkpointing"] = train_config.get("enable_gradient_checkpointing", True) + flat_config["enable_full_shard"] = train_config.get("enable_full_shard", True) + flat_config["enable_activation_offload"] = train_config.get("enable_activation_offload", False) + flat_config["init_device"] = train_config.get("init_device", "meta") + flat_config["load_checkpoint_path"] = train_config.get("load_checkpoint_path", "") + flat_config["ckpt_manager"] = train_config.get("ckpt_manager", "dcp") + flat_config["log_level"] = train_config.get("log_level", "INFO") # Data processing section (can be in train or data section) - data_config = config.get('data', {}) - flat_config['sample_packing_sequence_len'] = train_config.get('sample_packing_sequence_len') or data_config.get('sample_packing_sequence_len', 32000) - flat_config['enable_packing'] = train_config.get('enable_packing', data_config.get('enable_packing', True)) + data_config = config.get("data", {}) + flat_config["sample_packing_sequence_len"] = train_config.get("sample_packing_sequence_len") or data_config.get( + "sample_packing_sequence_len", 32000 + ) + flat_config["enable_packing"] = train_config.get("enable_packing", data_config.get("enable_packing", True)) # Output directory (can be in train section or top-level) - used for checkpoints, sampler weights, logs # Note: This uses output_dir from config, which should be on shared filesystem for multi-node - flat_config['output_dir'] = train_config.get('output_dir', config.get('output_dir', 'outputs')) + flat_config["output_dir"] = train_config.get("output_dir", config.get("output_dir", "outputs")) # Storage limit (can be in train section or top-level) - limits disk usage for output_dir - flat_config['storage_limit'] = train_config.get('storage_limit', config.get('storage_limit')) + flat_config["storage_limit"] = train_config.get("storage_limit", config.get("storage_limit")) # Idle session timeout (can be in train section or top-level) - sessions inactive for this duration are cleaned up - flat_config['idle_session_timeout'] = train_config.get('idle_session_timeout', config.get('idle_session_timeout', 7200.0)) + flat_config["idle_session_timeout"] = train_config.get( + "idle_session_timeout", config.get("idle_session_timeout", 7200.0) + ) # Training flags - flat_config['skip_initial_checkpoint'] = train_config.get('skip_initial_checkpoint', False) - flat_config['log_gradient_norms'] = train_config.get('log_gradient_norms', True) - flat_config['log_router_stats'] = train_config.get('log_router_stats', True) - flat_config['freeze_router'] = train_config.get('freeze_router', True) + flat_config["skip_initial_checkpoint"] = train_config.get("skip_initial_checkpoint", False) + flat_config["log_gradient_norms"] = train_config.get("log_gradient_norms", True) + flat_config["log_router_stats"] = train_config.get("log_router_stats", True) + flat_config["freeze_router"] = train_config.get("freeze_router", True) # Worker section - worker_config = config.get('worker', {}) - flat_config['worker_bind_address'] = worker_config.get('bind_address', 'auto') - flat_config['worker_bind_host'] = worker_config.get('bind_host', '0.0.0.0') - flat_config['worker_bind_port'] = worker_config.get('bind_port', 5556) - flat_config['engine_connect_host'] = worker_config.get('engine_connect_host') - flat_config['worker_connection_timeout'] = worker_config.get('connection_timeout', 120.0) - flat_config['worker_max_retries'] = worker_config.get('max_retries', 3) + worker_config = config.get("worker", {}) + flat_config["worker_bind_address"] = worker_config.get("bind_address", "auto") + flat_config["worker_bind_host"] = worker_config.get("bind_host", "0.0.0.0") + flat_config["worker_bind_port"] = worker_config.get("bind_port", 5556) + flat_config["engine_connect_host"] = worker_config.get("engine_connect_host") + flat_config["worker_connection_timeout"] = worker_config.get("connection_timeout", 120.0) + flat_config["worker_max_retries"] = worker_config.get("max_retries", 3) filtered_config = {k: v for k, v in flat_config.items() if k in valid_fields and v is not None} else: @@ -436,7 +439,7 @@ def load_server_arguments(config_path: str, overrides: Optional[Dict[str, any]] # Handle None values for Optional fields for key, value in list(filtered_config.items()): if value is None: - if key in ['config_path', 'tokenizer_path']: + if key in ["config_path", "tokenizer_path"]: del filtered_config[key] # Apply CLI overrides on top of YAML config @@ -470,23 +473,23 @@ def calculate_world_size_from_config(config_path: str) -> int: Returns: Total number of GPUs needed """ - with open(config_path, 'r') as f: + with open(config_path, "r") as f: config = yaml.safe_load(f) # Support both flat config (ServerArguments style) and nested config (train: section) - if 'train' in config: - train_config = config.get('train', {}) + if "train" in config: + train_config = config.get("train", {}) else: train_config = config # Get parallelism sizes (with defaults matching model_runner.py) - ep_size = train_config.get('expert_parallel_size', 1) - ulysses_size = train_config.get('ulysses_parallel_size', 1) - ringattn_size = train_config.get('ringattn_parallel_size', 1) + ep_size = train_config.get("expert_parallel_size", 1) + ulysses_size = train_config.get("ulysses_parallel_size", 1) + ringattn_size = train_config.get("ringattn_parallel_size", 1) # Data parallel sizes - dp_replicate_size = train_config.get('data_parallel_replicate_size', 1) - dp_shard_size = train_config.get('data_parallel_shard_size', 1) + dp_replicate_size = train_config.get("data_parallel_replicate_size", 1) + dp_shard_size = train_config.get("data_parallel_shard_size", 1) # Calculate world size # EP creates a separate 2D mesh where: world_size = ep_size * ep_fsdp_size @@ -494,7 +497,7 @@ def calculate_world_size_from_config(config_path: str) -> int: world_size = dp_replicate_size * dp_shard_size * ulysses_size * ringattn_size ep_fsdp_size = world_size // ep_size - logger.info(f"Calculated world size from config:") + logger.info("Calculated world size from config:") logger.info(f" expert_parallel_size: {ep_size}") logger.info(f" ulysses_parallel_size: {ulysses_size}") logger.info(f" ringattn_parallel_size: {ringattn_size}") @@ -511,6 +514,7 @@ def calculate_world_size_from_config(config_path: str) -> int: # Main Launcher # ============================================================================ + class Launcher: """Main launcher for all server components.""" @@ -599,7 +603,9 @@ def __init__( total_world_size = calculate_world_size_from_config(config_path) # For multi-node, nproc_per_node is total world size divided by number of nodes self.nproc_per_node = total_world_size // self.nnodes - logger.info(f"Total world size = {total_world_size}, nnodes = {self.nnodes}, nproc_per_node = {self.nproc_per_node}") + logger.info( + f"Total world size = {total_world_size}, nnodes = {self.nnodes}, nproc_per_node = {self.nproc_per_node}" + ) else: self.nproc_per_node = 1 # Not used in connect mode @@ -646,7 +652,9 @@ def __init__( if self.server_args: self.sample_packing_sequence_len = self.server_args.sample_packing_sequence_len self.enable_packing = self.server_args.enable_packing - logger.info(f"Using packing config: seq_len={self.sample_packing_sequence_len}, enabled={self.enable_packing}") + logger.info( + f"Using packing config: seq_len={self.sample_packing_sequence_len}, enabled={self.enable_packing}" + ) # Output directory - prefer from ServerArguments if available if self.server_args: @@ -851,9 +859,7 @@ def _save_initial_checkpoint(self, max_retries: int = 3, retry_delay: float = 5. ) except requests.exceptions.RequestException as e: - logger.warning( - f"Failed to save initial checkpoint (attempt {attempt + 1}/{max_retries}): {e}" - ) + logger.warning(f"Failed to save initial checkpoint (attempt {attempt + 1}/{max_retries}): {e}") if attempt < max_retries - 1: logger.info(f"Retrying in {retry_delay} seconds...") @@ -867,8 +873,9 @@ def _save_initial_checkpoint(self, max_retries: int = 3, retry_delay: float = 5. def _find_free_port() -> int: """Find and return a free TCP port.""" import socket + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(('', 0)) + s.bind(("", 0)) return s.getsockname()[1] def _launch_workers_with_torchrun(self): @@ -897,7 +904,8 @@ def _launch_workers_with_torchrun(self): f"--nproc-per-node={self.nproc_per_node}", f"--master-addr={self.master_addr}", f"--master-port={self.master_port}", - "-m", "xorl.server.runner.runner_dispatcher", + "-m", + "xorl.server.runner.runner_dispatcher", self.config_path, f"--worker.bind_address={self.worker_address}", ] @@ -1010,7 +1018,7 @@ def start(self): logger.info(" Process object created successfully") self.engine_process.start() - logger.info(f" Process.start() called") + logger.info(" Process.start() called") logger.info(f" Engine process PID: {self.engine_process.pid}") logger.info(f" Engine process alive: {self.engine_process.is_alive()}") except Exception as e: @@ -1119,7 +1127,6 @@ def start(self): def wait(self): """Wait for all processes to finish. Exit cleanly if any process dies.""" - import select try: while True: @@ -1219,6 +1226,7 @@ def stop(self): # CLI # ============================================================================ + def parse_server_overrides(argv: List[str]) -> Tuple[List[str], Dict[str, any]]: """ Parse --server.* arguments from command line. @@ -1323,7 +1331,7 @@ def main(): # => world_size = 2 * 4 = 8 GPUs Set these values in your config file under the 'train' section to control world size. - """ + """, ) # Mode selection @@ -1332,54 +1340,31 @@ def main(): type=str, choices=["auto", "connect"], default="auto", - help="Launch mode: 'auto' (launch workers with torchrun) or 'connect' (connect to external workers)" + help="Launch mode: 'auto' (launch workers with torchrun) or 'connect' (connect to external workers)", ) # Worker configuration + parser.add_argument("--config", type=str, help="Path to training config YAML (required for auto mode)") parser.add_argument( - "--config", - type=str, - help="Path to training config YAML (required for auto mode)" - ) - parser.add_argument( - "--worker-address", - type=str, - default=None, - help="Worker ZMQ address (default: tcp://127.0.0.1:)" + "--worker-address", type=str, default=None, help="Worker ZMQ address (default: tcp://127.0.0.1:)" ) # API Server configuration - parser.add_argument( - "--api-host", - type=str, - default="0.0.0.0", - help="API server host (default: 0.0.0.0)" - ) - parser.add_argument( - "--api-port", - type=int, - default=None, - help="API server port (default: auto-find free port)" - ) + parser.add_argument("--api-host", type=str, default="0.0.0.0", help="API server host (default: 0.0.0.0)") + parser.add_argument("--api-port", type=int, default=None, help="API server port (default: auto-find free port)") # Engine configuration parser.add_argument( - "--max-running-requests", - type=int, - default=2, - help="Maximum concurrent running requests (default: 2)" + "--max-running-requests", type=int, default=2, help="Maximum concurrent running requests (default: 2)" ) parser.add_argument( - "--max-pending-requests", - type=int, - default=100, - help="Maximum pending requests in queue (default: 100)" + "--max-pending-requests", type=int, default=100, help="Maximum pending requests in queue (default: 100)" ) parser.add_argument( "--operation-timeout", type=float, default=1800.0, - help="Timeout for engine operations in seconds (default: 1800.0)" + help="Timeout for engine operations in seconds (default: 1800.0)", ) # Logging parser.add_argument( @@ -1387,27 +1372,19 @@ def main(): type=str, default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"], - help="Logging level (default: INFO)" + help="Logging level (default: INFO)", ) # Distributed training (auto mode only) - parser.add_argument( - "--nnodes", - type=int, - default=1, - help="Number of nodes (auto mode, default: 1)" - ) + parser.add_argument("--nnodes", type=int, default=1, help="Number of nodes (auto mode, default: 1)") parser.add_argument( "--master-addr", type=str, default="127.0.0.1", - help="Master address for torch distributed (auto mode, default: 127.0.0.1)" + help="Master address for torch distributed (auto mode, default: 127.0.0.1)", ) parser.add_argument( - "--master-port", - type=int, - default=29500, - help="Master port for torch distributed (auto mode, default: 29500)" + "--master-port", type=int, default=29500, help="Master port for torch distributed (auto mode, default: 29500)" ) # Parse only the remaining args (after extracting --server.* overrides) @@ -1434,4 +1411,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/src/xorl/server/orchestrator/__init__.py b/src/xorl/server/orchestrator/__init__.py index 5ed9908c..e86b9df2 100644 --- a/src/xorl/server/orchestrator/__init__.py +++ b/src/xorl/server/orchestrator/__init__.py @@ -1,6 +1,7 @@ from xorl.server.orchestrator.orchestrator import Orchestrator -from xorl.server.orchestrator.scheduler import Scheduler -from xorl.server.orchestrator.request_processor import RequestProcessor from xorl.server.orchestrator.packing import Packer +from xorl.server.orchestrator.request_processor import RequestProcessor +from xorl.server.orchestrator.scheduler import Scheduler + __all__ = ["Orchestrator", "Scheduler", "RequestProcessor", "Packer"] diff --git a/src/xorl/server/orchestrator/orchestrator.py b/src/xorl/server/orchestrator/orchestrator.py index 33fec23c..33999b89 100644 --- a/src/xorl/server/orchestrator/orchestrator.py +++ b/src/xorl/server/orchestrator/orchestrator.py @@ -173,15 +173,15 @@ import zmq +from xorl.server.backend import Backend, RemoteBackend +from xorl.server.orchestrator.request_processor import RequestProcessor +from xorl.server.orchestrator.scheduler import Scheduler, SeqIdAwareFIFOPolicy from xorl.server.protocol.api_orchestrator import ( OrchestratorOutputs, OrchestratorRequest, OutputType, RequestType, ) -from xorl.server.backend import Backend, RemoteBackend -from xorl.server.orchestrator.request_processor import RequestProcessor -from xorl.server.orchestrator.scheduler import Scheduler, SeqIdAwareFIFOPolicy from xorl.server.utils.zmq_channels import SyncDealerChannel, SyncPushChannel @@ -693,7 +693,11 @@ def _process_engine_step(self) -> bool: request.payload.save_optimizer = False processor_method = getattr(self.request_processor, processor_method_name) - log_level = logging.DEBUG if operation in ("forward", "forward_backward", "optim_step", "get_adapter_info") else logging.INFO + log_level = ( + logging.DEBUG + if operation in ("forward", "forward_backward", "optim_step", "get_adapter_info") + else logging.INFO + ) logger.log(log_level, f"Executing {operation} for request {request.request_id}") output = self._call_async(processor_method(request)) diff --git a/src/xorl/server/orchestrator/packing.py b/src/xorl/server/orchestrator/packing.py index a8acd2e4..148f8733 100644 --- a/src/xorl/server/orchestrator/packing.py +++ b/src/xorl/server/orchestrator/packing.py @@ -46,11 +46,14 @@ import logging from abc import ABC, abstractmethod from typing import Any, Dict, List, Optional, TypedDict, Union + import torch + from xorl.data.constants import IGNORE_INDEX from xorl.distributed.parallel_state import get_parallel_state from xorl.utils.seqlen_pos_transform_utils import pos2culen + logger = logging.getLogger(__name__) @@ -83,8 +86,7 @@ def apply_weights_to_labels( if len(weights) != len(labels): raise ValueError( - f"Sample {sample_idx}: weights length ({len(weights)}) doesn't match " - f"labels length ({len(labels)})" + f"Sample {sample_idx}: weights length ({len(weights)}) doesn't match labels length ({len(labels)})" ) # Validate weights are only 0.0 or 1.0 @@ -97,10 +99,7 @@ def apply_weights_to_labels( ) # Apply mask: set labels to IGNORE_INDEX where weights=0 - masked_labels = [ - IGNORE_INDEX if w == 0 or w == 0.0 else label - for label, w in zip(labels, weights) - ] + masked_labels = [IGNORE_INDEX if w == 0 or w == 0.0 else label for label, w in zip(labels, weights)] return masked_labels @@ -133,8 +132,7 @@ def apply_advantages_to_labels( if len(advantages) != len(labels): raise ValueError( - f"Sample {sample_idx}: advantages length ({len(advantages)}) doesn't match " - f"labels length ({len(labels)})" + f"Sample {sample_idx}: advantages length ({len(advantages)}) doesn't match labels length ({len(labels)})" ) # Use torch for faster processing with long sequences (e.g., 128k tokens) @@ -205,7 +203,7 @@ def unpack_per_token_outputs( # Calculate expected output length if shifted # Each sample loses 1 token, so shifted_len = pos_len - num_samples expected_shifted_len = pos_len - num_samples - is_shifted = (output_len == expected_shifted_len) + is_shifted = output_len == expected_shifted_len if output_len != pos_len and not is_shifted: logger.warning( @@ -229,7 +227,7 @@ def unpack_per_token_outputs( sample_output_len = sample_len # Extract this sample's output - sample_output = packed_output[output_offset:output_offset + sample_output_len] + sample_output = packed_output[output_offset : output_offset + sample_output_len] results.append(sample_output.cpu().tolist()) output_offset += sample_output_len @@ -241,17 +239,19 @@ def unpack_per_token_outputs( # Type Definitions for Input/Output Contracts # ============================================================================ + class Datum(TypedDict, total=False): """ Input format for a single training sample. - + Required fields: - input_ids: List[int] - + Optional fields: - labels: List[int] - position_ids: List[int] """ + input_ids: List[int] # Required labels: List[int] # Optional position_ids: List[int] # Optional @@ -260,9 +260,9 @@ class Datum(TypedDict, total=False): class MicroBatch(TypedDict, total=False): """ Output format for a packed micro-batch. - + Contains packed sequences with identity tracking. - + Fields: - input_ids: Packed input token sequences - labels: Packed label sequences (if present) @@ -270,6 +270,7 @@ class MicroBatch(TypedDict, total=False): - request_id: Request ID to track source request - batch_id: Batch index (multiple batches per request) """ + input_ids: List[List[int]] labels: List[List[int]] position_ids: List[List[int]] @@ -281,20 +282,21 @@ class MicroBatch(TypedDict, total=False): # Abstract Base Class for Data Processors # ============================================================================ + class Packer(ABC): """ Abstract base class for data packing strategies. - + Subclasses should implement the `pack()` method to define their own packing algorithm (sequential, sorted by length, optimal bin packing, etc.) - + Example: class MyCustomPacker(Packer): def pack(self, datum_list, max_seq_len): # Your custom packing logic here return micro_batches """ - + @abstractmethod def pack( self, @@ -304,17 +306,17 @@ def pack( ) -> List[Dict[str, Any]]: """ Pack datum into micro-batches. - + Args: datum_list: List of datum (see Datum TypedDict for format) max_seq_len: Maximum sequence length per micro-batch request_id: Request ID to track source (optional) - + Returns: List of micro-batches (see MicroBatch TypedDict for format) """ pass - + def get_name(self) -> str: """Return the name of this packing strategy.""" return self.__class__.__name__ @@ -324,18 +326,19 @@ def get_name(self) -> str: # Sequential Packer Implementation # ============================================================================ + class SequentialPacker(Packer): """ Sequential greedy packing strategy. - + Packs samples sequentially into micro-batches using a greedy first-fit algorithm. Samples are added to the current batch until it reaches capacity, then a new batch is started. - + Args: enable_packing: If False, each sample becomes its own micro-batch (no packing) log_stats: If True, log detailed packing statistics after packing - + Example: >>> packer = SequentialPacker(enable_packing=True) >>> datum_list = [ @@ -348,7 +351,7 @@ class SequentialPacker(Packer): >>> batches[0]["num_samples"] 2 """ - + def __init__( self, enable_packing: bool = True, @@ -358,7 +361,7 @@ def __init__( self.enable_packing = enable_packing self.log_stats = log_stats self.pad_to_multiple_of = pad_to_multiple_of - + def pack( self, datum_list: List[Dict[str, Any]], @@ -367,45 +370,45 @@ def pack( ) -> List[Dict[str, Any]]: """ Pack samples sequentially into micro-batches with concatenation. - + When packing is enabled, samples are CONCATENATED into a single sequence per micro-batch with shape [1, total_packed_length]. Position IDs are generated to mark sample boundaries (reset to 0 for each sample). - + This format is compatible with TextSequenceShardCollator for sequence parallelism. - + Args: datum_list: List of datum dictionaries (see Datum TypedDict) max_seq_len: Maximum sequence length per micro-batch request_id: Request ID to track source request - + Returns: List of micro-batch dictionaries with concatenated sequences """ if not datum_list: logger.warning("Empty datum_list provided") return [] - + logger.debug( f"[{self.get_name()}] Packing {len(datum_list)} samples with " f"max_seq_len={max_seq_len}, packing={'enabled' if self.enable_packing else 'disabled'}, " f"request_id={request_id}" ) - + # If packing disabled, create one batch per sample (not concatenated) if not self.enable_packing: return self._pack_without_packing(datum_list, request_id) - + # Pack samples sequentially with CONCATENATION # Output format: [1, total_packed_length] with position_ids marking boundaries micro_batches = [] current_batch = self._create_empty_packed_batch(request_id, batch_id=0) current_tokens = 0 - + # Track skipped samples for better error messages skipped_samples = [] - + for sample_idx, datum in enumerate(datum_list): # Handle nested structure to check for input_ids if "input_ids" in datum: @@ -419,24 +422,24 @@ def pack( continue if not isinstance(input_ids, list): - input_ids = input_ids.tolist() if hasattr(input_ids, 'tolist') else list(input_ids) + input_ids = input_ids.tolist() if hasattr(input_ids, "tolist") else list(input_ids) seq_len = len(input_ids) - + # Skip samples that exceed max_seq_len if seq_len > max_seq_len: reason = f"has {seq_len} tokens, exceeding max_seq_len {max_seq_len}" skipped_samples.append((sample_idx, reason)) logger.warning(f"Sample {sample_idx} {reason}. Skipping.") continue - + # Skip empty samples if seq_len == 0: reason = "has empty input_ids (0 tokens)" skipped_samples.append((sample_idx, reason)) logger.warning(f"Sample {sample_idx} {reason}, skipping") continue - + # Check if sample fits in current batch if current_tokens + seq_len > max_seq_len: # Current batch is full, finalize and start a new one @@ -445,26 +448,26 @@ def pack( micro_batches.append(current_batch) current_batch = self._create_empty_packed_batch(request_id, batch_id=len(micro_batches)) current_tokens = 0 - + # Add sample to current batch (concatenate) self._add_sample_to_packed_batch(current_batch, datum, sample_idx) current_tokens += seq_len - + # Add the last batch if it has samples if current_batch["_num_samples"] > 0: self._finalize_packed_batch(current_batch) micro_batches.append(current_batch) - + # If all samples were skipped, raise a descriptive error if not micro_batches and skipped_samples: self._raise_all_samples_skipped_error(len(datum_list), skipped_samples, max_seq_len) - + # Log statistics if self.log_stats: self._log_packing_stats(datum_list, micro_batches, max_seq_len) - + return micro_batches - + def _create_empty_batch(self, request_id: str, batch_id: int) -> Dict[str, Any]: """Create an empty micro-batch structure (list of lists format).""" return { @@ -478,7 +481,7 @@ def _create_empty_batch(self, request_id: str, batch_id: int) -> Dict[str, Any]: def _create_empty_packed_batch(self, request_id: str, batch_id: int) -> Dict[str, Any]: """ Create an empty packed micro-batch structure. - + Packed batches concatenate all samples into single flat lists. After finalization, the batch will have shape [1, total_len]. """ @@ -539,17 +542,17 @@ def _add_sample_to_packed_batch( if input_ids is None: raise ValueError(f"Sample {sample_idx} missing 'input_ids'") if not isinstance(input_ids, list): - input_ids = input_ids.tolist() if hasattr(input_ids, 'tolist') else list(input_ids) + input_ids = input_ids.tolist() if hasattr(input_ids, "tolist") else list(input_ids) # Extract labels/target_tokens if "labels" in flattened_datum: labels = flattened_datum["labels"] if not isinstance(labels, list): - labels = labels.tolist() if hasattr(labels, 'tolist') else list(labels) + labels = labels.tolist() if hasattr(labels, "tolist") else list(labels) elif "target_tokens" in flattened_datum: labels = flattened_datum["target_tokens"] if not isinstance(labels, list): - labels = labels.tolist() if hasattr(labels, 'tolist') else list(labels) + labels = labels.tolist() if hasattr(labels, "tolist") else list(labels) else: # No labels - use IGNORE_INDEX for this sample's tokens labels = [IGNORE_INDEX] * len(input_ids) @@ -558,21 +561,18 @@ def _add_sample_to_packed_batch( weights = flattened_datum.get("weights") if weights is not None: if not isinstance(weights, list): - weights = weights.tolist() if hasattr(weights, 'tolist') else list(weights) + weights = weights.tolist() if hasattr(weights, "tolist") else list(weights) # Extract advantages if present (for RL losses like importance_sampling) advantages = flattened_datum.get("advantages") if advantages is not None: if not isinstance(advantages, list): - advantages = advantages.tolist() if hasattr(advantages, 'tolist') else list(advantages) + advantages = advantages.tolist() if hasattr(advantages, "tolist") else list(advantages) # Detect if tokens are already shifted (xorl_client API format) # xorl_client format: len(input_ids) == len(target_tokens) and target_tokens field exists # HF format: len(input_ids) == len(labels) but need shifting - is_already_shifted = ( - "target_tokens" in flattened_datum and - len(input_ids) == len(labels) - ) + is_already_shifted = "target_tokens" in flattened_datum and len(input_ids) == len(labels) if not is_already_shifted and len(input_ids) == len(labels): # HF format: shift tokens here @@ -624,7 +624,7 @@ def _add_sample_to_packed_batch( if key not in ["input_ids", "position_ids", "labels", "weights"]: if key not in batch: batch[key] = [] - if hasattr(value, 'tolist'): + if hasattr(value, "tolist"): value = value.tolist() if isinstance(value, list) and len(value) == seq_len: # Sequence field - concatenate @@ -677,8 +677,15 @@ def _finalize_packed_batch(self, batch: Dict[str, Any]) -> None: # Pad other sequence fields that were wrapped as [[...]] for key, value in batch.items(): - if key in ("input_ids", "labels", "position_ids", "request_id", - "batch_id", "num_samples", "_shifted"): + if key in ( + "input_ids", + "labels", + "position_ids", + "request_id", + "batch_id", + "num_samples", + "_shifted", + ): continue if isinstance(value, list) and len(value) == 1 and isinstance(value[0], list): if len(value[0]) == seq_len: @@ -703,7 +710,7 @@ def _finalize_packed_batch(self, batch: Dict[str, Any]) -> None: f"Finalized packed batch {batch['batch_id']}: " f"total_len={len(batch['input_ids'][0])}, num_samples={num_samples}" ) - + def _add_sample_to_batch( self, batch: Dict[str, Any], @@ -731,7 +738,7 @@ def _add_sample_to_batch( if input_ids is None: raise ValueError(f"Sample {sample_idx} missing 'input_ids'") if not isinstance(input_ids, list): - input_ids = input_ids.tolist() if hasattr(input_ids, 'tolist') else list(input_ids) + input_ids = input_ids.tolist() if hasattr(input_ids, "tolist") else list(input_ids) seq_len = len(input_ids) batch["input_ids"].append(input_ids) @@ -740,7 +747,7 @@ def _add_sample_to_batch( if "position_ids" in flattened_datum: position_ids = flattened_datum["position_ids"] if not isinstance(position_ids, list): - position_ids = position_ids.tolist() if hasattr(position_ids, 'tolist') else list(position_ids) + position_ids = position_ids.tolist() if hasattr(position_ids, "tolist") else list(position_ids) else: # Auto-generate position_ids: [0, 1, 2, ..., seq_len-1] position_ids = list(range(seq_len)) @@ -751,12 +758,12 @@ def _add_sample_to_batch( if "labels" in flattened_datum: labels = flattened_datum["labels"] if not isinstance(labels, list): - labels = labels.tolist() if hasattr(labels, 'tolist') else list(labels) + labels = labels.tolist() if hasattr(labels, "tolist") else list(labels) elif "target_tokens" in flattened_datum: # For RL datums, use target_tokens as labels labels = flattened_datum["target_tokens"] if not isinstance(labels, list): - labels = labels.tolist() if hasattr(labels, 'tolist') else list(labels) + labels = labels.tolist() if hasattr(labels, "tolist") else list(labels) else: # No labels for this sample labels = [] @@ -767,7 +774,7 @@ def _add_sample_to_batch( weights = flattened_datum.get("weights") if weights is not None: if not isinstance(weights, list): - weights = weights.tolist() if hasattr(weights, 'tolist') else list(weights) + weights = weights.tolist() if hasattr(weights, "tolist") else list(weights) labels = apply_weights_to_labels(labels, weights, sample_idx) # Apply advantages mask to labels if advantages field is present @@ -776,7 +783,7 @@ def _add_sample_to_batch( advantages = flattened_datum.get("advantages") if advantages is not None: if not isinstance(advantages, list): - advantages = advantages.tolist() if hasattr(advantages, 'tolist') else list(advantages) + advantages = advantages.tolist() if hasattr(advantages, "tolist") else list(advantages) labels = apply_advantages_to_labels(labels, advantages, sample_idx) batch["labels"].append(labels) @@ -792,9 +799,15 @@ def _add_sample_to_batch( # Append value if not isinstance(value, list): - value = value.tolist() if hasattr(value, 'tolist') else list(value) if hasattr(value, '__iter__') and not isinstance(value, str) else value + value = ( + value.tolist() + if hasattr(value, "tolist") + else list(value) + if hasattr(value, "__iter__") and not isinstance(value, str) + else value + ) batch[key].append(value) - + def _pack_without_packing( self, datum_list: List[Dict[str, Any]], @@ -805,11 +818,10 @@ def _pack_without_packing( micro_batches = [] skipped_samples = [] - + for batch_id, datum in enumerate(datum_list): # Handle nested structure to check for input_ids - has_input_ids = ("input_ids" in datum) or \ - ("model_input" in datum and "input_ids" in datum["model_input"]) + has_input_ids = ("input_ids" in datum) or ("model_input" in datum and "input_ids" in datum["model_input"]) if not has_input_ids: reason = "missing 'input_ids' field" skipped_samples.append((batch_id, reason)) @@ -825,7 +837,7 @@ def _pack_without_packing( self._raise_all_samples_skipped_error(len(datum_list), skipped_samples) return micro_batches - + def _log_packing_stats( self, datum_list: List[Dict[str, Any]], @@ -836,15 +848,15 @@ def _log_packing_stats( if not micro_batches: logger.warning("No batches created") return - + total_samples = len(datum_list) total_batches = len(micro_batches) - + # Calculate statistics from actual data # Handle both packed format (list of one list) and unpacked format (list of lists) samples_per_batch = [] tokens_per_batch = [] - + for batch in micro_batches: input_ids = batch["input_ids"] if "num_samples" in batch: @@ -855,17 +867,19 @@ def _log_packing_stats( # Unpacked format: list of separate sequences samples_per_batch.append(len(input_ids)) tokens_per_batch.append(sum(len(seq) for seq in input_ids)) - + total_tokens = sum(tokens_per_batch) total_capacity = total_batches * max_seq_len utilization = (total_tokens / total_capacity * 100) if total_capacity > 0 else 0 - + is_packed = any("num_samples" in batch for batch in micro_batches) - + logger.debug("=" * 70) logger.debug(f"[{self.get_name()}] Packing Statistics:") logger.debug("-" * 70) - logger.debug(f" Packing mode: {'CONCATENATED (batch_size=1)' if is_packed else 'BATCHED (separate sequences)'}") + logger.debug( + f" Packing mode: {'CONCATENATED (batch_size=1)' if is_packed else 'BATCHED (separate sequences)'}" + ) logger.debug(f" Total samples: {total_samples}") logger.debug(f" Total micro-batches: {total_batches}") logger.debug(f" Total tokens: {total_tokens}") @@ -873,18 +887,20 @@ def _log_packing_stats( logger.debug(f" Total capacity: {total_capacity}") logger.debug(f" Utilization: {utilization:.1f}%") logger.debug("-" * 70) - logger.debug(f" Samples per batch:") + logger.debug(" Samples per batch:") logger.debug(f" Min: {min(samples_per_batch)}") logger.debug(f" Max: {max(samples_per_batch)}") logger.debug(f" Avg: {sum(samples_per_batch) / len(samples_per_batch):.1f}") logger.debug("-" * 70) - logger.debug(f" Tokens per batch:") + logger.debug(" Tokens per batch:") logger.debug(f" Min: {min(tokens_per_batch)}") logger.debug(f" Max: {max(tokens_per_batch)}") logger.debug(f" Avg: {sum(tokens_per_batch) / len(tokens_per_batch):.1f}") logger.debug("=" * 70) - logger.info(f"[{self.get_name()}] Packed {total_samples} samples into {total_batches} batches ({utilization:.1f}% utilization, {total_tokens} tokens)") - + logger.info( + f"[{self.get_name()}] Packed {total_samples} samples into {total_batches} batches ({utilization:.1f}% utilization, {total_tokens} tokens)" + ) + def _raise_all_samples_skipped_error( self, total_samples: int, @@ -893,7 +909,7 @@ def _raise_all_samples_skipped_error( ) -> None: """ Raise a descriptive ValueError when all samples were skipped during packing. - + Args: total_samples: Total number of samples submitted skipped_samples: List of (sample_idx, reason) tuples @@ -905,32 +921,32 @@ def _raise_all_samples_skipped_error( if reason not in reasons_count: reasons_count[reason] = [] reasons_count[reason].append(idx) - + # Build detailed error message error_lines = [ f"All {total_samples} samples were skipped - no valid batches could be created.", ] - + if max_seq_len is not None: error_lines.append(f"Server max_seq_len is {max_seq_len} tokens.") - + error_lines.append("") error_lines.append("Skipped samples by reason:") - + for reason, indices in reasons_count.items(): if len(indices) <= 5: indices_str = ", ".join(str(i) for i in indices) else: indices_str = f"{', '.join(str(i) for i in indices[:5])}, ... ({len(indices)} total)" error_lines.append(f" - {reason}: sample(s) {indices_str}") - + error_lines.append("") error_lines.append("Please check your input data and ensure:") error_lines.append(" 1. Each sample has an 'input_ids' field (list of token IDs)") if max_seq_len is not None: error_lines.append(f" 2. Sequence lengths do not exceed {max_seq_len} tokens") error_lines.append(" 3. Samples are not empty (at least 1 token)") - + error_msg = "\n".join(error_lines) raise ValueError(error_msg) @@ -939,6 +955,7 @@ def _raise_all_samples_skipped_error( # Function-based API (backward compatible) # ============================================================================ + def pack_samples( datum_list: List[Dict[str, Any]], max_seq_len: int = 2048, @@ -987,11 +1004,12 @@ def pack_samples( # Validation Utilities # ============================================================================ + def validate_micro_batches(micro_batches: List[Dict[str, Any]]) -> bool: """ Validate micro-batch structure. - Supports both packed format (batch_size=1, concatenated) and + Supports both packed format (batch_size=1, concatenated) and unpacked format (batch_size>1, separate sequences). Checks: @@ -1015,41 +1033,37 @@ def validate_micro_batches(micro_batches: List[Dict[str, Any]]) -> bool: if field not in batch: logger.error(f"Batch {batch_idx} missing required field: {field}") return False - + # Check that input_ids is a list if not isinstance(batch["input_ids"], list): logger.error(f"Batch {batch_idx}: input_ids should be a list") return False - + num_sequences = len(batch["input_ids"]) if num_sequences == 0: logger.error(f"Batch {batch_idx}: input_ids is empty") return False - + # Check if this is packed format (list of one list) or unpacked (list of multiple lists) is_packed = "num_samples" in batch - + # Check lengths are consistent if len(batch["position_ids"]) != num_sequences: logger.error( - f"Batch {batch_idx}: position_ids length mismatch " - f"({len(batch['position_ids'])} != {num_sequences})" + f"Batch {batch_idx}: position_ids length mismatch ({len(batch['position_ids'])} != {num_sequences})" ) return False - + if len(batch["labels"]) != num_sequences: - logger.error( - f"Batch {batch_idx}: labels length mismatch " - f"({len(batch['labels'])} != {num_sequences})" - ) + logger.error(f"Batch {batch_idx}: labels length mismatch ({len(batch['labels'])} != {num_sequences})") return False - + # Check each sequence has matching lengths for i, input_seq in enumerate(batch["input_ids"]): if not isinstance(input_seq, list): logger.error(f"Batch {batch_idx}, sequence {i}: input_ids should be a list of lists") return False - + pos_seq = batch["position_ids"][i] if len(input_seq) != len(pos_seq): logger.error( @@ -1057,7 +1071,7 @@ def validate_micro_batches(micro_batches: List[Dict[str, Any]]) -> bool: f"!= position_ids length ({len(pos_seq)})" ) return False - + labels_seq = batch["labels"][i] if len(input_seq) != len(labels_seq): logger.error( @@ -1065,7 +1079,7 @@ def validate_micro_batches(micro_batches: List[Dict[str, Any]]) -> bool: f"!= labels length ({len(labels_seq)})" ) return False - + # Check batch_id is an integer if not isinstance(batch["batch_id"], int): logger.error(f"Batch {batch_idx}: batch_id should be an integer") diff --git a/src/xorl/server/orchestrator/request_processor.py b/src/xorl/server/orchestrator/request_processor.py index 8f045117..60de2448 100644 --- a/src/xorl/server/orchestrator/request_processor.py +++ b/src/xorl/server/orchestrator/request_processor.py @@ -41,7 +41,7 @@ import logging import math import time -from typing import Any, Callable, Dict, Optional, Union +from typing import Any, Callable, Dict, Union import torch @@ -49,8 +49,6 @@ from xorl.server.orchestrator.packing import pack_samples, unpack_per_token_outputs, validate_micro_batches from xorl.server.protocol.api_orchestrator import OrchestratorOutputs, OrchestratorRequest, OutputType from xorl.server.protocol.operations import ( - LOAD_STATE_TIMEOUT, - SAVE_STATE_TIMEOUT, AdapterStateData, KillSessionData, LoadStateData, @@ -181,14 +179,11 @@ async def _execute_model_pass( if not batches: raise ValueError( - f"No batches created from {len(data)} samples. " - "The packer did not produce any valid batches." + f"No batches created from {len(data)} samples. The packer did not produce any valid batches." ) if not validate_micro_batches(batches): - raise ValueError( - "Invalid batch structure after packing. This may indicate a bug in the packing logic." - ) + raise ValueError("Invalid batch structure after packing. This may indicate a bug in the packing logic.") t_packed = time.perf_counter() logger.debug(f"Packed {len(data)} samples into {len(batches)} batches") @@ -311,13 +306,17 @@ def _unpack_per_sample_outputs(result: Dict, batches: list) -> list: async def execute_forward_backward(self, request: OrchestratorRequest) -> OrchestratorOutputs: """Execute forward-backward pass on workers.""" return await self._execute_model_pass( - request, "forward_backward", OutputType.FORWARD_BACKWARD, + request, + "forward_backward", + OutputType.FORWARD_BACKWARD, ) async def execute_forward(self, request: OrchestratorRequest) -> OrchestratorOutputs: """Execute forward pass on workers (no gradient computation).""" return await self._execute_model_pass( - request, "forward", OutputType.FORWARD, + request, + "forward", + OutputType.FORWARD, ) async def _execute_operation( @@ -382,13 +381,19 @@ def build_output(result): return [output_dict] return await self._execute_operation( - request, "optim_step", + request, + "optim_step", self.backend.optim_step( - lr=p.lr, gradient_clip=p.gradient_clip, - beta1=p.beta1, beta2=p.beta2, eps=p.eps, - model_id=p.model_id, request_id=request.request_id, + lr=p.lr, + gradient_clip=p.gradient_clip, + beta1=p.beta1, + beta2=p.beta2, + eps=p.eps, + model_id=p.model_id, + request_id=request.request_id, ), - OutputType.OPTIM_STEP, build_output, + OutputType.OPTIM_STEP, + build_output, ) async def execute_save_state(self, request: OrchestratorRequest) -> OrchestratorOutputs: @@ -399,20 +404,27 @@ async def execute_save_state(self, request: OrchestratorRequest) -> Orchestrator def build_output(result): actual_path = result.get("checkpoint_path", checkpoint_path) success = result.get("success", False) - return [{ - "checkpoint_path": actual_path, "success": success, - "execution_time": result.get("execution_time", 0.0), - "message": "Checkpoint saved successfully" if success else "Save failed", - }] + return [ + { + "checkpoint_path": actual_path, + "success": success, + "execution_time": result.get("execution_time", 0.0), + "message": "Checkpoint saved successfully" if success else "Save failed", + } + ] return await self._execute_operation( - request, "save_state", + request, + "save_state", self.backend.save_state( - checkpoint_path=p.checkpoint_path, save_optimizer=p.save_optimizer, - use_timestamp=p.use_timestamp, model_id=p.model_id, + checkpoint_path=p.checkpoint_path, + save_optimizer=p.save_optimizer, + use_timestamp=p.use_timestamp, + model_id=p.model_id, request_id=request.request_id, ), - OutputType.SAVE_STATE, build_output, + OutputType.SAVE_STATE, + build_output, ) async def execute_save_lora_only(self, request: OrchestratorRequest) -> OrchestratorOutputs: @@ -425,19 +437,25 @@ async def execute_save_lora_only(self, request: OrchestratorRequest) -> Orchestr def build_output(result): actual_path = result.get("lora_path", lora_path) success = result.get("success", False) - return [{ - "lora_path": actual_path, "success": success, - "execution_time": result.get("execution_time", 0.0), - "message": "LoRA adapter saved successfully (PEFT format)" if success else "Save failed", - }] + return [ + { + "lora_path": actual_path, + "success": success, + "execution_time": result.get("execution_time", 0.0), + "message": "LoRA adapter saved successfully (PEFT format)" if success else "Save failed", + } + ] return await self._execute_operation( - request, "save_lora_only", + request, + "save_lora_only", self.backend.save_lora_only( - lora_path=p.lora_path, model_id=p.model_id, + lora_path=p.lora_path, + model_id=p.model_id, request_id=request.request_id, ), - OutputType.SAVE_LORA_ONLY, build_output, + OutputType.SAVE_LORA_ONLY, + build_output, ) async def execute_save_full_weights(self, request: OrchestratorRequest) -> OrchestratorOutputs: @@ -451,21 +469,29 @@ async def execute_save_full_weights(self, request: OrchestratorRequest) -> Orche def build_output(result): success = result.get("success", False) num_shards = result.get("num_shards", 1) - return [{ - "output_path": output_path, "dtype": dtype, - "num_shards": num_shards, "success": success, - "execution_time": result.get("execution_time", 0.0), - "message": f"Full weights saved as safetensors ({num_shards} shards)" if success else "Save failed", - }] + return [ + { + "output_path": output_path, + "dtype": dtype, + "num_shards": num_shards, + "success": success, + "execution_time": result.get("execution_time", 0.0), + "message": f"Full weights saved as safetensors ({num_shards} shards)" if success else "Save failed", + } + ] return await self._execute_operation( - request, "save_full_weights", + request, + "save_full_weights", self.backend.save_full_weights( - output_path=p.output_path, dtype=p.dtype, - base_model_path=p.base_model_path, model_id=p.model_id, + output_path=p.output_path, + dtype=p.dtype, + base_model_path=p.base_model_path, + model_id=p.model_id, request_id=request.request_id, ), - OutputType.SAVE_STATE, build_output, + OutputType.SAVE_STATE, + build_output, ) async def execute_load_state(self, request: OrchestratorRequest) -> OrchestratorOutputs: @@ -477,45 +503,66 @@ async def execute_load_state(self, request: OrchestratorRequest) -> Orchestrator def build_output(result): success = result.get("success", False) - return [{ - "checkpoint_path": checkpoint_path, "success": success, - "execution_time": result.get("execution_time", 0.0), - "message": "Checkpoint loaded successfully" if success else "Load failed", - }] + return [ + { + "checkpoint_path": checkpoint_path, + "success": success, + "execution_time": result.get("execution_time", 0.0), + "message": "Checkpoint loaded successfully" if success else "Load failed", + } + ] return await self._execute_operation( - request, "load_state", + request, + "load_state", self.backend.load_state( - checkpoint_path=p.checkpoint_path, load_optimizer=p.load_optimizer, - model_id=p.model_id, request_id=request.request_id, + checkpoint_path=p.checkpoint_path, + load_optimizer=p.load_optimizer, + model_id=p.model_id, + request_id=request.request_id, ), - OutputType.LOAD_STATE, build_output, + OutputType.LOAD_STATE, + build_output, ) async def execute_sleep(self, request: OrchestratorRequest) -> OrchestratorOutputs: """Execute sleep operation (offload model and optimizer to CPU).""" + def build_output(result): - return [{"status": result.get("status", "sleeping"), - "offload_time": result.get("offload_time", 0.0), - "execution_time": result.get("execution_time", 0.0)}] + return [ + { + "status": result.get("status", "sleeping"), + "offload_time": result.get("offload_time", 0.0), + "execution_time": result.get("execution_time", 0.0), + } + ] return await self._execute_operation( - request, "sleep", + request, + "sleep", self.backend.sleep(request_id=request.request_id), - OutputType.SLEEP, build_output, + OutputType.SLEEP, + build_output, ) async def execute_wake_up(self, request: OrchestratorRequest) -> OrchestratorOutputs: """Execute wake_up operation (load model and optimizer to GPU).""" + def build_output(result): - return [{"status": result.get("status", "awake"), - "load_time": result.get("load_time", 0.0), - "execution_time": result.get("execution_time", 0.0)}] + return [ + { + "status": result.get("status", "awake"), + "load_time": result.get("load_time", 0.0), + "execution_time": result.get("execution_time", 0.0), + } + ] return await self._execute_operation( - request, "wake_up", + request, + "wake_up", self.backend.wake_up(request_id=request.request_id), - OutputType.WAKE_UP, build_output, + OutputType.WAKE_UP, + build_output, ) async def execute_sync_inference_weights(self, request: OrchestratorRequest) -> OrchestratorOutputs: @@ -525,28 +572,37 @@ async def execute_sync_inference_weights(self, request: OrchestratorRequest) -> raise ValueError("inference endpoints must be provided") def build_output(result): - return [{ - "success": result.get("success", False), - "message": result.get("message", ""), - "transfer_time": result.get("transfer_time", 0.0), - "total_bytes": result.get("total_bytes", 0), - "num_parameters": result.get("num_parameters", 0), - "num_buckets": result.get("num_buckets", 0), - "endpoint_results": result.get("endpoint_results", []), - "execution_time": result.get("execution_time", 0.0), - }] + return [ + { + "success": result.get("success", False), + "message": result.get("message", ""), + "transfer_time": result.get("transfer_time", 0.0), + "total_bytes": result.get("total_bytes", 0), + "num_parameters": result.get("num_parameters", 0), + "num_buckets": result.get("num_buckets", 0), + "endpoint_results": result.get("endpoint_results", []), + "execution_time": result.get("execution_time", 0.0), + } + ] return await self._execute_operation( - request, "sync_inference_weights", + request, + "sync_inference_weights", self.backend.sync_inference_weights( - endpoints=p.endpoints, master_address=p.master_address, - master_port=p.master_port, group_name=p.group_name, - buffer_size_mb=p.buffer_size_mb, sync_method=p.sync_method, - flush_cache=p.flush_cache, pause_mode=p.pause_mode, - weight_version=p.weight_version, quantization=p.quantization, + endpoints=p.endpoints, + master_address=p.master_address, + master_port=p.master_port, + group_name=p.group_name, + buffer_size_mb=p.buffer_size_mb, + sync_method=p.sync_method, + flush_cache=p.flush_cache, + pause_mode=p.pause_mode, + weight_version=p.weight_version, + quantization=p.quantization, request_id=request.request_id, ), - OutputType.SYNC_INFERENCE_WEIGHTS, build_output, + OutputType.SYNC_INFERENCE_WEIGHTS, + build_output, ) async def execute_register_adapter(self, request: OrchestratorRequest) -> OrchestratorOutputs: @@ -557,11 +613,15 @@ def build_output(result): return {"result": result} return await self._execute_operation( - request, "register_adapter", + request, + "register_adapter", self.backend.register_adapter( - model_id=p.model_id, lr=p.lr, request_id=request.request_id, + model_id=p.model_id, + lr=p.lr, + request_id=request.request_id, ), - OutputType.REGISTER_ADAPTER, build_output, + OutputType.REGISTER_ADAPTER, + build_output, ) async def execute_save_adapter_state(self, request: OrchestratorRequest) -> OrchestratorOutputs: @@ -572,12 +632,16 @@ def build_output(result): return {"result": result} return await self._execute_operation( - request, "save_adapter_state", + request, + "save_adapter_state", self.backend.save_adapter_state( - model_id=p.model_id, path=p.path, - save_optimizer=p.save_optimizer, request_id=request.request_id, + model_id=p.model_id, + path=p.path, + save_optimizer=p.save_optimizer, + request_id=request.request_id, ), - OutputType.SAVE_ADAPTER_STATE, build_output, + OutputType.SAVE_ADAPTER_STATE, + build_output, ) async def execute_load_adapter_state(self, request: OrchestratorRequest) -> OrchestratorOutputs: @@ -590,24 +654,31 @@ def build_output(result): return {"result": result} return await self._execute_operation( - request, "load_adapter_state", + request, + "load_adapter_state", self.backend.load_adapter_state( - model_id=p.model_id, path=p.path, - load_optimizer=p.load_optimizer, lr=p.lr, + model_id=p.model_id, + path=p.path, + load_optimizer=p.load_optimizer, + lr=p.lr, request_id=request.request_id, ), - OutputType.LOAD_ADAPTER_STATE, build_output, + OutputType.LOAD_ADAPTER_STATE, + build_output, ) async def execute_get_adapter_info(self, request: OrchestratorRequest) -> OrchestratorOutputs: """Execute get adapter info on workers.""" + def build_output(result): return [result] return await self._execute_operation( - request, "get_adapter_info", + request, + "get_adapter_info", self.backend.get_adapter_info(request_id=request.request_id), - OutputType.GET_ADAPTER_INFO, build_output, + OutputType.GET_ADAPTER_INFO, + build_output, ) async def execute_kill_session(self, request: OrchestratorRequest) -> OrchestratorOutputs: @@ -615,20 +686,25 @@ async def execute_kill_session(self, request: OrchestratorRequest) -> Orchestrat p: KillSessionData = request.payload def build_output(result): - return [{ - "success": result.get("success", False), - "message": result.get("message", ""), - "checkpoint_path": result.get("checkpoint_path"), - "execution_time": result.get("execution_time", 0.0), - }] + return [ + { + "success": result.get("success", False), + "message": result.get("message", ""), + "checkpoint_path": result.get("checkpoint_path"), + "execution_time": result.get("execution_time", 0.0), + } + ] return await self._execute_operation( - request, "kill_session", + request, + "kill_session", self.backend.kill_session( - model_id=p.model_id, save_checkpoint=p.save_checkpoint, + model_id=p.model_id, + save_checkpoint=p.save_checkpoint, request_id=request.request_id, ), - OutputType.KILL_SESSION, build_output, + OutputType.KILL_SESSION, + build_output, ) # ======================================================================== diff --git a/src/xorl/server/orchestrator/scheduler.py b/src/xorl/server/orchestrator/scheduler.py index a27ee58a..76958e63 100644 --- a/src/xorl/server/orchestrator/scheduler.py +++ b/src/xorl/server/orchestrator/scheduler.py @@ -290,7 +290,7 @@ def get_next_request(self) -> Optional[ScheduledRequest]: while scheduler.has_pending_requests(): # Get next request (respects max_running_requests) scheduled_req = scheduler.get_next_request() - + if not scheduled_req: # At max capacity, wait for a slot to free up break @@ -346,11 +346,10 @@ def get_next_request(self) -> Optional[ScheduledRequest]: from abc import ABC, abstractmethod from collections import deque from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Deque +from typing import Any, Deque, Dict, Optional from xorl.server.protocol.api_orchestrator import ( OrchestratorRequest, - RequestType, ) @@ -361,6 +360,7 @@ def get_next_request(self) -> Optional[ScheduledRequest]: # Request Metadata # ============================================================================ + @dataclass class ScheduledRequest: """ @@ -368,6 +368,7 @@ class ScheduledRequest: Tracks timing and execution state for scheduling decisions. """ + request: OrchestratorRequest arrival_time: float = field(default_factory=time.time) retries: int = 0 @@ -420,6 +421,7 @@ def mark_failed(self, error: str): # Scheduling Policies # ============================================================================ + class SchedulingPolicy(ABC): """ Abstract base class for scheduling policies. @@ -625,6 +627,7 @@ def get_policy_name(self) -> str: # Scheduler # ============================================================================ + class Scheduler: """ Request scheduler for Orchestrator. @@ -742,8 +745,7 @@ def get_next_request(self) -> Optional[ScheduledRequest]: current_running = self.get_running_count() if current_running >= self.max_running_requests: logger.debug( - f"Cannot dispatch request: at max running capacity " - f"({current_running}/{self.max_running_requests})" + f"Cannot dispatch request: at max running capacity ({current_running}/{self.max_running_requests})" ) return None @@ -844,10 +846,7 @@ def mark_failed(self, request_id: str, error: str): self.completed_requests.append(scheduled_req) self.total_failed += 1 - logger.error( - f"Request {request_id} failed " - f"(operation={scheduled_req.operation}, error={error})" - ) + logger.error(f"Request {request_id} failed (operation={scheduled_req.operation}, error={error})") # ======================================================================== # Query Methods @@ -976,6 +975,7 @@ def __repr__(self) -> str: # Seq_id-Aware Scheduling Policy # ============================================================================ + class SeqIdAwareFIFOPolicy(SchedulingPolicy): """ Scheduling policy that sorts requests by seq_id per model_id. @@ -1180,10 +1180,7 @@ def clear(self): self.fifo_queue.clear() self.request_map.clear() - logger.info( - f"Cleared SeqIdAwareFIFO queue " - f"({seq_id_count} seq_id + {fifo_count} FIFO requests removed)" - ) + logger.info(f"Cleared SeqIdAwareFIFO queue ({seq_id_count} seq_id + {fifo_count} FIFO requests removed)") def get_policy_name(self) -> str: """Get policy name.""" @@ -1196,10 +1193,7 @@ def get_stats(self) -> Dict[str, Any]: Returns: Dictionary with scheduler statistics """ - seq_id_by_model = { - model_id: list(sorted(seq_map.keys())) - for model_id, seq_map in self.seq_id_requests.items() - } + seq_id_by_model = {model_id: sorted(seq_map.keys()) for model_id, seq_map in self.seq_id_requests.items()} return { "seq_id_requests": sum(len(seq_map) for seq_map in self.seq_id_requests.values()), "fifo_queue_size": len(self.fifo_queue), @@ -1211,4 +1205,3 @@ def get_stats(self) -> Dict[str, Any]: # ============================================================================ # Future Scheduling Policies (Placeholder Implementations) # ============================================================================ - diff --git a/src/xorl/server/protocol/__init__.py b/src/xorl/server/protocol/__init__.py index 3f90d8bc..ce3b9f45 100644 --- a/src/xorl/server/protocol/__init__.py +++ b/src/xorl/server/protocol/__init__.py @@ -11,24 +11,6 @@ """ # Typed operation payloads -from xorl.server.protocol.operations import ( # noqa: F401 - AbortData, - AdapterStateData, - EmptyData, - KillSessionData, - LoadStateData, - ModelPassData, - OperationPayload, - OptimStepData, - RegisterAdapterData, - SaveFullWeightsData, - SaveLoraOnlyData, - SaveStateData, - SyncWeightsData, - payload_from_dict, - payload_to_dict, -) - from xorl.server.protocol.api_orchestrator import ( # noqa: F401 OrchestratorOutputs, OrchestratorRequest, @@ -51,11 +33,28 @@ validate_output, validate_request, ) +from xorl.server.protocol.operations import ( # noqa: F401 + AbortData, + AdapterStateData, + EmptyData, + KillSessionData, + LoadStateData, + ModelPassData, + OperationPayload, + OptimStepData, + RegisterAdapterData, + SaveFullWeightsData, + SaveLoraOnlyData, + SaveStateData, + SyncWeightsData, + payload_from_dict, + payload_to_dict, +) from xorl.server.protocol.orchestrator_runner import ( # noqa: F401 BaseMessage, - RunnerDispatchCommand, MessageType, RunnerAck, + RunnerDispatchCommand, RunnerReady, RunnerResponse, create_ack_for_request, diff --git a/src/xorl/server/protocol/api_orchestrator.py b/src/xorl/server/protocol/api_orchestrator.py index 56dbffd2..310b4921 100644 --- a/src/xorl/server/protocol/api_orchestrator.py +++ b/src/xorl/server/protocol/api_orchestrator.py @@ -196,9 +196,7 @@ def from_msgpack(cls, data: bytes) -> "OrchestratorOutputs": def __repr__(self) -> str: status = "finished" if self.finished else "streaming" error_str = f", error='{self.error}'" if self.error else "" - return ( - f"OrchestratorOutputs(id={self.request_id[:8]}..., type={self.output_type.value}, status={status}{error_str})" - ) + return f"OrchestratorOutputs(id={self.request_id[:8]}..., type={self.output_type.value}, status={status}{error_str})" # ============================================================================ @@ -243,8 +241,11 @@ def create_forward_backward_output( if additional_metrics: outputs_data.update(additional_metrics) return OrchestratorOutputs( - request_id=request_id, output_type=OutputType.FORWARD_BACKWARD, - outputs=[outputs_data], finished=True, error=error, + request_id=request_id, + output_type=OutputType.FORWARD_BACKWARD, + outputs=[outputs_data], + finished=True, + error=error, ) @@ -268,8 +269,11 @@ def create_optim_step_output( if additional_metrics: outputs_data.update(additional_metrics) return OrchestratorOutputs( - request_id=request_id, output_type=OutputType.OPTIM_STEP, - outputs=[outputs_data], finished=True, error=error, + request_id=request_id, + output_type=OutputType.OPTIM_STEP, + outputs=[outputs_data], + finished=True, + error=error, ) @@ -281,8 +285,11 @@ def create_save_state_output( ) -> OrchestratorOutputs: """Create a save checkpoint output response.""" return _build_output( - request_id, OutputType.SAVE_STATE, error=error, - success=success, checkpoint_path=checkpoint_path, + request_id, + OutputType.SAVE_STATE, + error=error, + success=success, + checkpoint_path=checkpoint_path, ) @@ -294,8 +301,11 @@ def create_save_lora_only_output( ) -> OrchestratorOutputs: """Create a save LoRA-only output response.""" return _build_output( - request_id, OutputType.SAVE_LORA_ONLY, error=error, - success=success, lora_path=lora_path, + request_id, + OutputType.SAVE_LORA_ONLY, + error=error, + success=success, + lora_path=lora_path, ) @@ -307,8 +317,11 @@ def create_load_state_output( ) -> OrchestratorOutputs: """Create a load checkpoint output response.""" return _build_output( - request_id, OutputType.LOAD_STATE, error=error, - success=success, checkpoint_path=checkpoint_path, + request_id, + OutputType.LOAD_STATE, + error=error, + success=success, + checkpoint_path=checkpoint_path, ) @@ -322,8 +335,13 @@ def create_save_adapter_state_output( ) -> OrchestratorOutputs: """Create a save adapter state output response.""" return _build_output( - request_id, OutputType.SAVE_ADAPTER_STATE, error=error, - success=success, model_id=model_id, path=path, step=step, + request_id, + OutputType.SAVE_ADAPTER_STATE, + error=error, + success=success, + model_id=model_id, + path=path, + step=step, ) @@ -337,8 +355,13 @@ def create_load_adapter_state_output( ) -> OrchestratorOutputs: """Create a load adapter state output response.""" return _build_output( - request_id, OutputType.LOAD_ADAPTER_STATE, error=error, - success=success, model_id=model_id, path=path, step=step, + request_id, + OutputType.LOAD_ADAPTER_STATE, + error=error, + success=success, + model_id=model_id, + path=path, + step=step, ) @@ -359,8 +382,11 @@ def create_health_check_output( if additional_info: outputs_data.update(additional_info) return OrchestratorOutputs( - request_id=request_id, output_type=OutputType.HEALTH_CHECK, - outputs=[outputs_data], finished=True, error=error, + request_id=request_id, + output_type=OutputType.HEALTH_CHECK, + outputs=[outputs_data], + finished=True, + error=error, ) @@ -372,8 +398,11 @@ def create_sleep_output( ) -> OrchestratorOutputs: """Create a sleep output response.""" return _build_output( - request_id, OutputType.SLEEP, error=error, - status=status, offload_time=offload_time, + request_id, + OutputType.SLEEP, + error=error, + status=status, + offload_time=offload_time, ) @@ -385,8 +414,11 @@ def create_wake_up_output( ) -> OrchestratorOutputs: """Create a wake_up output response.""" return _build_output( - request_id, OutputType.WAKE_UP, error=error, - status=status, load_time=load_time, + request_id, + OutputType.WAKE_UP, + error=error, + status=status, + load_time=load_time, ) @@ -403,10 +435,16 @@ def create_sync_weights_output( ) -> OrchestratorOutputs: """Create a sync inference weights output response.""" return _build_output( - request_id, OutputType.SYNC_INFERENCE_WEIGHTS, error=error, - success=success, message=message, transfer_time=transfer_time, - total_bytes=total_bytes, num_parameters=num_parameters, - num_buckets=num_buckets, endpoint_results=endpoint_results or [], + request_id, + OutputType.SYNC_INFERENCE_WEIGHTS, + error=error, + success=success, + message=message, + transfer_time=transfer_time, + total_bytes=total_bytes, + num_parameters=num_parameters, + num_buckets=num_buckets, + endpoint_results=endpoint_results or [], ) @@ -417,8 +455,11 @@ def create_error_output( ) -> OrchestratorOutputs: """Create an error output response.""" return OrchestratorOutputs( - request_id=request_id, output_type=operation_type, - outputs=[], finished=True, error=error_message, + request_id=request_id, + output_type=operation_type, + outputs=[], + finished=True, + error=error_message, ) diff --git a/src/xorl/server/protocol/operations.py b/src/xorl/server/protocol/operations.py index a870f260..cc370a4a 100644 --- a/src/xorl/server/protocol/operations.py +++ b/src/xorl/server/protocol/operations.py @@ -12,9 +12,11 @@ """ import os -from dataclasses import asdict, dataclass, field, fields as dc_fields +from dataclasses import asdict, dataclass, field +from dataclasses import fields as dc_fields from typing import Any, Dict, List, Optional, Union + # ============================================================================ # Timeout Constants (shared by engine/executor and backend/remote) # ============================================================================ diff --git a/src/xorl/server/runner/__init__.py b/src/xorl/server/runner/__init__.py index d6d8b72c..563195c8 100644 --- a/src/xorl/server/runner/__init__.py +++ b/src/xorl/server/runner/__init__.py @@ -1,4 +1,5 @@ -from xorl.server.runner.runner_dispatcher import RunnerDispatcher from xorl.server.runner.model_runner import ModelRunner +from xorl.server.runner.runner_dispatcher import RunnerDispatcher + __all__ = ["RunnerDispatcher", "ModelRunner"] diff --git a/src/xorl/server/runner/adapters/__init__.py b/src/xorl/server/runner/adapters/__init__.py index cc22aa91..924cbc10 100644 --- a/src/xorl/server/runner/adapters/__init__.py +++ b/src/xorl/server/runner/adapters/__init__.py @@ -1,4 +1,5 @@ -from xorl.server.runner.adapters.manager import LoRAAdapterManager from xorl.server.runner.adapters.adapter_coordinator import AdapterCoordinator +from xorl.server.runner.adapters.manager import LoRAAdapterManager + __all__ = ["LoRAAdapterManager", "AdapterCoordinator"] diff --git a/src/xorl/server/runner/adapters/adapter_coordinator.py b/src/xorl/server/runner/adapters/adapter_coordinator.py index fdc7d04b..d9230653 100644 --- a/src/xorl/server/runner/adapters/adapter_coordinator.py +++ b/src/xorl/server/runner/adapters/adapter_coordinator.py @@ -22,6 +22,7 @@ ) from xorl.server.runner.model_runner import ModelRunner + logger = logging.getLogger(__name__) @@ -69,11 +70,13 @@ def broadcast_adapter_state(self, model_id: str, default_lr: float) -> None: # Broadcast metadata metadata = [None] if self.rank == 0: - metadata = [{ - "global_step": adapter_state.global_step, - "global_forward_backward_step": adapter_state.global_forward_backward_step, - "lr": adapter_state.lr, - }] + metadata = [ + { + "global_step": adapter_state.global_step, + "global_forward_backward_step": adapter_state.global_forward_backward_step, + "lr": adapter_state.lr, + } + ] dist.broadcast_object_list(metadata, src=0, group=self.cpu_group) # Update metadata on non-rank-0 workers @@ -156,18 +159,12 @@ def auto_load_if_evicted(self, model_id: str) -> Tuple[bool, Optional[str]]: if checkpoint_path is None: # No checkpoint — register fresh adapter - logger.debug( - f"Rank {self.rank}: Auto-registering new adapter '{model_id}' " - f"(no previous checkpoint found)" - ) + logger.debug(f"Rank {self.rank}: Auto-registering new adapter '{model_id}' (no previous checkpoint found)") self._register_fresh_adapter(model_id) return True, None # Auto-load from checkpoint - logger.debug( - f"Rank {self.rank}: Auto-loading evicted adapter '{model_id}' " - f"from checkpoint: {checkpoint_path}" - ) + logger.debug(f"Rank {self.rank}: Auto-loading evicted adapter '{model_id}' from checkpoint: {checkpoint_path}") try: effective_lr = 1e-5 # Default, will be overwritten from checkpoint @@ -191,8 +188,7 @@ def auto_load_if_evicted(self, model_id: str) -> Tuple[bool, Optional[str]]: except Exception as e: logger.error( - f"Rank {self.rank}: Failed to auto-load adapter '{model_id}' " - f"from {checkpoint_path}: {e}", + f"Rank {self.rank}: Failed to auto-load adapter '{model_id}' from {checkpoint_path}: {e}", exc_info=True, ) return False, None @@ -320,7 +316,9 @@ async def handle_load_adapter_state(self, command_dict: Dict[str, Any]) -> Dict[ lr=lr, ) - logger.debug(f"Rank {self.rank}: load_adapter_state completed: model_id={model_id}, step={result.get('step', 0)}") + logger.debug( + f"Rank {self.rank}: load_adapter_state completed: model_id={model_id}, step={result.get('step', 0)}" + ) # Step 3: Broadcast loaded weights from rank 0 to all other ranks self.broadcast_adapter_state(model_id, effective_lr) @@ -401,7 +399,9 @@ async def handle_kill_session(self, command_dict: Dict[str, Any]) -> Dict[str, A model_id = p.model_id save_checkpoint = p.save_checkpoint - logger.debug(f"Rank {self.rank}: Handling kill_session for model_id={model_id}, save_checkpoint={save_checkpoint}") + logger.debug( + f"Rank {self.rank}: Handling kill_session for model_id={model_id}, save_checkpoint={save_checkpoint}" + ) try: result = self.trainer.kill_session(model_id=model_id, save_checkpoint=save_checkpoint) diff --git a/src/xorl/server/runner/adapters/manager.py b/src/xorl/server/runner/adapters/manager.py index e8445197..4aa1b4ea 100644 --- a/src/xorl/server/runner/adapters/manager.py +++ b/src/xorl/server/runner/adapters/manager.py @@ -25,9 +25,11 @@ from safetensors.torch import load_file as safetensors_load_file from safetensors.torch import save_file as safetensors_save_file + try: from torch.distributed._tensor import DTensor from torch.distributed._tensor.placement_types import Shard + _HAS_DTENSOR = True except ImportError: _HAS_DTENSOR = False @@ -47,7 +49,7 @@ class AdapterState: model_id: str lora_params: Dict[str, nn.Parameter] # Actual Parameters with own .grad - optimizer: torch.optim.Optimizer # Per-adapter optimizer + optimizer: torch.optim.Optimizer # Per-adapter optimizer global_step: int = 0 global_forward_backward_step: int = 0 lr: float = 1e-5 @@ -369,7 +371,7 @@ def optim_step( # Update learning rate state.lr = lr for pg in state.optimizer.param_groups: - pg['lr'] = lr + pg["lr"] = lr # Deferred gradient normalization: scale raw gradients by 1/accumulated_valid_tokens if accumulated_valid_tokens > 0: @@ -381,10 +383,7 @@ def optim_step( # Always use clip_grad_norm_ for correct grad norm computation # Using a large clip value (10000.0) effectively means no clipping clip_value = gradient_clip if (gradient_clip is not None and gradient_clip > 0) else 10000.0 - grad_norm = torch.nn.utils.clip_grad_norm_( - list(state.lora_params.values()), - clip_value - ) + grad_norm = torch.nn.utils.clip_grad_norm_(list(state.lora_params.values()), clip_value) if hasattr(grad_norm, "item"): grad_norm = grad_norm.item() diff --git a/src/xorl/server/runner/checkpoint/__init__.py b/src/xorl/server/runner/checkpoint/__init__.py index e8b286e8..595132f6 100644 --- a/src/xorl/server/runner/checkpoint/__init__.py +++ b/src/xorl/server/runner/checkpoint/__init__.py @@ -1,3 +1,4 @@ from xorl.server.runner.checkpoint.manager import CheckpointManager + __all__ = ["CheckpointManager"] diff --git a/src/xorl/server/runner/checkpoint/manager.py b/src/xorl/server/runner/checkpoint/manager.py index 40cbbaf0..486ae1e1 100644 --- a/src/xorl/server/runner/checkpoint/manager.py +++ b/src/xorl/server/runner/checkpoint/manager.py @@ -33,6 +33,7 @@ from xorl.models import save_model_weights from xorl.utils import helper + logger = logging.getLogger(__name__) @@ -167,8 +168,7 @@ def save_adapter_state( """ if self._adapter_manager is None: raise ValueError( - "Multi-tenancy is not enabled. save_adapter_state requires " - "LoRA with adapter_manager enabled." + "Multi-tenancy is not enabled. save_adapter_state requires LoRA with adapter_manager enabled." ) start_time = time.time() @@ -239,8 +239,7 @@ def load_adapter_state( """ if self._adapter_manager is None: raise ValueError( - "Multi-tenancy is not enabled. load_adapter_state requires " - "LoRA with adapter_manager enabled." + "Multi-tenancy is not enabled. load_adapter_state requires LoRA with adapter_manager enabled." ) result = self._adapter_manager.load_adapter_state( @@ -399,13 +398,9 @@ def save_full_weights( use_distributed = distributed_write and ps.ep_enabled and ps.world_size > 1 if use_distributed: - return self._save_full_weights_distributed( - output_path, dtype, base_model_path - ) + return self._save_full_weights_distributed(output_path, dtype, base_model_path) else: - return self._save_full_weights_single_writer( - output_path, dtype, base_model_path - ) + return self._save_full_weights_single_writer(output_path, dtype, base_model_path) @staticmethod def _copy_model_configs(output_path: str, base_model_path: str) -> None: @@ -441,9 +436,7 @@ def _save_full_weights_single_writer( # Get checkpoint handler for HF-compatible weight transforms # (e.g., splitting gate_up_proj back into gate_proj + up_proj) checkpoint_handler = ( - self.model.get_checkpoint_handler() - if hasattr(self.model, "get_checkpoint_handler") - else None + self.model.get_checkpoint_handler() if hasattr(self.model, "get_checkpoint_handler") else None ) # save_model_weights handles dtype conversion, sharding, and index.json @@ -531,9 +524,7 @@ def _save_full_weights_distributed( model_state = ModelState(self.model) state_dict_meta = model_state.state_dict() else: - state_dict_meta = { - name: param for name, param in self.model.named_parameters() - } + state_dict_meta = {name: param for name, param in self.model.named_parameters()} # Compute tensor sizes and shard assignments (all ranks compute same assignment) tensor_infos = [] # [(name, estimated_size, is_dtensor), ...] @@ -551,8 +542,10 @@ def _save_full_weights_distributed( # Estimate size after dtype conversion if tensor.dtype in (torch.float32, torch.float16, torch.bfloat16): - element_size = target_dtype.itemsize if hasattr(target_dtype, 'itemsize') else ( - 2 if target_dtype in (torch.float16, torch.bfloat16) else 4 + element_size = ( + target_dtype.itemsize + if hasattr(target_dtype, "itemsize") + else (2 if target_dtype in (torch.float16, torch.bfloat16) else 4) ) else: element_size = tensor.element_size() @@ -646,23 +639,22 @@ def _save_full_weights_distributed( torch.cuda.synchronize() torch.cuda.empty_cache() - logger.debug( - f"Rank {self.rank}: extracted {dtensor_count} DTensors, {regular_count} regular tensors" - ) + logger.debug(f"Rank {self.rank}: extracted {dtensor_count} DTensors, {regular_count} regular tensors") # Phase 3: Writers save their shards in parallel my_shard_results = [] # [(shard_idx, shard_name, weight_map, size), ...] if is_writer and my_shards: + def _save_shard(shard_idx): if num_shards == 1: shard_name = "model.safetensors" else: - shard_name = f"model-{shard_idx+1:05d}-of-{num_shards:05d}.safetensors" + shard_name = f"model-{shard_idx + 1:05d}-of-{num_shards:05d}.safetensors" shard_path = os.path.join(output_path, shard_name) shard_data = my_shard_data[shard_idx] save_file(shard_data, shard_path) - shard_weight_map = {n: shard_name for n in shard_data.keys()} + shard_weight_map = dict.fromkeys(shard_data.keys(), shard_name) shard_size = sum(t.numel() * t.element_size() for t in shard_data.values()) return shard_idx, shard_name, shard_weight_map, shard_size @@ -702,12 +694,10 @@ def _save_shard(shard_idx): # Gather serialized results max_size = max(s.item() for s in all_sizes) local_padded = torch.zeros(max_size, dtype=torch.uint8, device="cuda") - local_padded[:len(local_results_bytes)] = torch.tensor( + local_padded[: len(local_results_bytes)] = torch.tensor( list(local_results_bytes), dtype=torch.uint8, device="cuda" ) - all_results_padded = [ - torch.zeros(max_size, dtype=torch.uint8, device="cuda") for _ in range(world_size) - ] + all_results_padded = [torch.zeros(max_size, dtype=torch.uint8, device="cuda") for _ in range(world_size)] dist.all_gather(all_results_padded, local_padded) if self.rank == 0: @@ -794,10 +784,7 @@ def save_state( # Base weights never change, so we only need adapter + optimizer + metadata if self._adapter_manager is not None: target_model_id = model_id or self._adapter_manager.current_adapter_id or "default" - logger.info( - f"Multi-tenancy mode: delegating save_state to save_adapter_state " - f"(model_id={target_model_id})" - ) + logger.info(f"Multi-tenancy mode: delegating save_state to save_adapter_state (model_id={target_model_id})") return self.save_adapter_state( model_id=target_model_id, path=checkpoint_path, @@ -907,10 +894,7 @@ def load_state( # This handles weight loading + optimizer + metadata if self._adapter_manager is not None: target_model_id = model_id or self._adapter_manager.current_adapter_id or "default" - logger.info( - f"Multi-tenancy mode: delegating load_state to load_adapter_state " - f"(model_id={target_model_id})" - ) + logger.info(f"Multi-tenancy mode: delegating load_state to load_adapter_state (model_id={target_model_id})") return self.load_adapter_state( model_id=target_model_id, path=checkpoint_path, diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index b2b171bb..7e63817a 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -13,21 +13,16 @@ See xorl.server.runner.runner_dispatcher for the entry point. """ -import gc import logging import os import time -from typing import Any, Dict, List, Optional, Union - -import numpy as np +from typing import Any, Dict, List, Optional import torch import torch.distributed as dist import torch.nn.functional as F from transformers import AutoTokenizer -from torch.distributed.tensor import distribute_tensor - from xorl.checkpoint import build_checkpointer from xorl.data.constants import IGNORE_INDEX from xorl.distributed.offloading import build_activation_offloading_context @@ -35,7 +30,6 @@ from xorl.distributed.pipeline_parallel import build_pipeline_schedule, build_pp_stage from xorl.distributed.sequence_parallel.data import gather_outputs from xorl.lora import LoraLinear -from xorl.models import all_ranks_load_weights, rank0_load_and_broadcast_weights from xorl.models.layers.moe.routing_replay import set_replay_stage from xorl.ops.loss import ( causallm_loss_function, @@ -43,21 +37,22 @@ policy_loss_function, ) from xorl.optim import build_optimizer +from xorl.server.runner.adapters import LoRAAdapterManager +from xorl.server.runner.checkpoint import CheckpointManager +from xorl.server.runner.utils import MoeMetricsTracker, RoutingReplayHandler, run_self_test, validate_token_ids from xorl.trainers.model_builder import build_training_model from xorl.trainers.training_utils import ( - pp_loss_fn, clip_gradients, count_valid_tokens, forward_backward_pp, - maybe_merge_lora as _maybe_merge_lora_util, negotiate_pp_seq_len, pad_micro_batches_for_pp, + pp_loss_fn, sync_sp_gradients, ) -from xorl.server.runner.checkpoint import CheckpointManager -from xorl.server.runner.adapters import LoRAAdapterManager -from xorl.server.runner.utils import MoeMetricsTracker, RoutingReplayHandler -from xorl.server.runner.utils import run_self_test, validate_token_ids +from xorl.trainers.training_utils import ( + maybe_merge_lora as _maybe_merge_lora_util, +) from xorl.utils import helper from xorl.utils.device import get_device_id, get_device_type, get_torch_device, synchronize from xorl.utils.dist_utils import all_reduce @@ -157,7 +152,14 @@ class ModelRunner: _LOSS_EXCLUDE_KEYS = { "causallm_loss": {"labels", "_original_position_ids", "rollout_logprobs"}, "cross_entropy": {"labels", "_original_position_ids", "rollout_logprobs"}, - "importance_sampling": {"labels", "target_tokens", "logprobs", "advantages", "_original_position_ids", "rollout_logprobs"}, + "importance_sampling": { + "labels", + "target_tokens", + "logprobs", + "advantages", + "_original_position_ids", + "rollout_logprobs", + }, "policy_loss": {"labels", "target_tokens", "logprobs", "advantages", "rollout_logprobs"}, } @@ -275,8 +277,8 @@ def __init__( adapter_manager=self._adapter_manager, ) # Sync initial attributes - self._checkpoint_mgr.lora_target_modules = getattr(self, 'lora_target_modules', None) - self._checkpoint_mgr.lora_alpha_value = getattr(self, 'lora_alpha_value', None) + self._checkpoint_mgr.lora_target_modules = getattr(self, "lora_target_modules", None) + self._checkpoint_mgr.lora_alpha_value = getattr(self, "lora_alpha_value", None) # Setup MoE metrics collection if this is a Qwen3 MoE model self._moe_tracker = MoeMetricsTracker(self.model_config_obj, self.train_config, self.rank) @@ -519,7 +521,9 @@ def _initialize_model(self): trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) total_params = sum(p.numel() for p in self.model.parameters()) logger.info(f"Model loaded and parallelized on rank {self.rank}") - logger.info(f" - Trainable params: {trainable_params:,} ({100.0 * trainable_params / max(total_params, 1):.2f}%)") + logger.info( + f" - Trainable params: {trainable_params:,} ({100.0 * trainable_params / max(total_params, 1):.2f}%)" + ) logger.info(f" - Total params: {total_params:,}") def _resolve_lora_target_modules(self) -> List[str]: @@ -547,8 +551,7 @@ def _initialize_optimizer(self): if optimizer_type == "muon": optimizer_kwargs = { k: self.train_config[k] - for k in ("muon_lr", "muon_momentum", "muon_nesterov", - "muon_ns_steps", "muon_adjust_lr_fn") + for k in ("muon_lr", "muon_momentum", "muon_nesterov", "muon_ns_steps", "muon_adjust_lr_fn") if k in self.train_config } self.optimizer = build_optimizer( @@ -603,9 +606,7 @@ def register_lora_adapter(self, model_id: str, lr: float) -> Dict[str, Any]: RuntimeError: If LoRA is not enabled or adapter manager not initialized """ if self._adapter_manager is None: - raise RuntimeError( - "Cannot register adapter: LoRA is not enabled or adapter manager not initialized" - ) + raise RuntimeError("Cannot register adapter: LoRA is not enabled or adapter manager not initialized") self._adapter_manager.register_adapter( model_id=model_id, @@ -655,8 +656,12 @@ def _collect_per_token_outputs(self, per_token_tensors, micro_batch, accumulator gathered = {} for key, tensor in per_token_tensors.items(): gathered[key] = gather_outputs( - tensor, gather_dim=-1, padding_dim=-1, - unpad_dim_size=original_seq_len, scale_grad=False, group=ulysses_group, + tensor, + gather_dim=-1, + padding_dim=-1, + unpad_dim_size=original_seq_len, + scale_grad=False, + group=ulysses_group, ) if position_ids is not None: @@ -757,8 +762,11 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): if loss_fn in ["causallm_loss", "cross_entropy"]: labels = micro_batch.get("labels") _result = causallm_loss_function( - hidden_states=hidden_states, weight=effective_weight, - labels=labels, return_per_token=return_per_token, ce_mode=self.ce_mode, + hidden_states=hidden_states, + weight=effective_weight, + labels=labels, + return_per_token=return_per_token, + ce_mode=self.ce_mode, lm_head_fp32=self.lm_head_fp32, ) loss = _result.loss @@ -773,9 +781,13 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): compute_kl_stats = params.get("compute_kl_stats", False) _result = importance_sampling_loss_function( - hidden_states=hidden_states, weight=effective_weight, - labels=target_tokens, old_logprobs=old_logprobs, advantages=advantages, - ce_mode=self.ce_mode, compute_kl_stats=compute_kl_stats, + hidden_states=hidden_states, + weight=effective_weight, + labels=target_tokens, + old_logprobs=old_logprobs, + advantages=advantages, + ce_mode=self.ce_mode, + compute_kl_stats=compute_kl_stats, lm_head_fp32=self.lm_head_fp32, ) loss = _result.loss @@ -794,7 +806,7 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): diag_logits = (hs_flat @ effective_weight.t()).float() diag_log_probs = F.log_softmax(diag_logits, dim=-1) - valid = (target_tokens.reshape(-1) != IGNORE_INDEX) + valid = target_tokens.reshape(-1) != IGNORE_INDEX valid_indices = valid.nonzero(as_tuple=True)[0] valid_log_probs = diag_log_probs[valid_indices] @@ -803,7 +815,9 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): # Target token logprob and rank target_ids = target_tokens.reshape(-1)[valid_indices] - target_lps = valid_log_probs[torch.arange(len(valid_indices), device=valid_log_probs.device), target_ids] + target_lps = valid_log_probs[ + torch.arange(len(valid_indices), device=valid_log_probs.device), target_ids + ] target_ranks = (valid_log_probs > target_lps.unsqueeze(-1)).sum(dim=-1) + 1 # Entropy: -sum(p * log p) @@ -816,18 +830,23 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): cp_rank = ps.cp_rank if ps.cp_enabled else 0 diag_path_ranked = f"{diag_path}.rank{cp_rank}" os.makedirs(os.path.dirname(diag_path_ranked) or ".", exist_ok=True) - torch.save({ - "topk_logprobs": topk_vals.cpu(), - "topk_ids": topk_ids.cpu(), - "target_ids": target_ids.cpu(), - "target_logprobs": target_lps.cpu(), - "target_ranks": target_ranks.cpu(), - "entropy": entropy.cpu(), - "valid_positions": valid_indices.cpu(), - "cp_rank": cp_rank, - }, diag_path_ranked) - logger.info(f"Diagnostic top-{diagnostic_topk} saved to {diag_path_ranked} " - f"({len(valid_indices)} valid positions, cp_rank={cp_rank})") + torch.save( + { + "topk_logprobs": topk_vals.cpu(), + "topk_ids": topk_ids.cpu(), + "target_ids": target_ids.cpu(), + "target_logprobs": target_lps.cpu(), + "target_ranks": target_ranks.cpu(), + "entropy": entropy.cpu(), + "valid_positions": valid_indices.cpu(), + "cp_rank": cp_rank, + }, + diag_path_ranked, + ) + logger.info( + f"Diagnostic top-{diagnostic_topk} saved to {diag_path_ranked} " + f"({len(valid_indices)} valid positions, cp_rank={cp_rank})" + ) del diag_logits, diag_log_probs, diag_probs, valid_log_probs @@ -851,13 +870,23 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): logger.warning("use_tis=True but rollout_logprobs not provided.") _result = policy_loss_function( - hidden_states=hidden_states, weight=effective_weight, - labels=target_tokens, old_logprobs=old_logprobs, advantages=advantages, + hidden_states=hidden_states, + weight=effective_weight, + labels=target_tokens, + old_logprobs=old_logprobs, + advantages=advantages, rollout_logprobs=rollout_logprobs, - eps_clip=eps_clip, eps_clip_high=eps_clip_high, eps_clip_c=eps_clip_c, - use_tis=use_tis, tis_clip_low=tis_clip_low, tis_clip_high=tis_clip_high, - ce_mode=self.ce_mode, num_chunks=num_chunks, compute_kl_stats=compute_kl_stats, - lm_head_fp32=self.lm_head_fp32, icepop_beta=icepop_beta, + eps_clip=eps_clip, + eps_clip_high=eps_clip_high, + eps_clip_c=eps_clip_c, + use_tis=use_tis, + tis_clip_low=tis_clip_low, + tis_clip_high=tis_clip_high, + ce_mode=self.ce_mode, + num_chunks=num_chunks, + compute_kl_stats=compute_kl_stats, + lm_head_fp32=self.lm_head_fp32, + icepop_beta=icepop_beta, ) loss = _result.loss per_token_outputs["logprobs"] = _result.per_token_logprobs @@ -925,9 +954,17 @@ def _compute_per_sample_k3( # Unified forward loop # ========================================================================= - def _forward_loop(self, micro_batches, loss_fn, loss_fn_params, *, - compute_backward=True, r3_enabled=False, - model_id="default", abort_callback=None): + def _forward_loop( + self, + micro_batches, + loss_fn, + loss_fn_params, + *, + compute_backward=True, + r3_enabled=False, + model_id="default", + abort_callback=None, + ): """Core forward (+ optional backward) loop shared between forward and forward_backward.""" params = loss_fn_params or {} return_per_token = params.get("return_per_token", True) @@ -956,7 +993,9 @@ def _forward_loop(self, micro_batches, loss_fn, loss_fn_params, *, } labels = micro_batch.get("labels", micro_batch.get("target_tokens")) - local_valid_tokens = (labels != IGNORE_INDEX).sum() if labels is not None else torch.tensor(0, device=get_device_type()) + local_valid_tokens = ( + (labels != IGNORE_INDEX).sum() if labels is not None else torch.tensor(0, device=get_device_type()) + ) # R3: switch to replay_forward so MoEBlock pops pre-populated routing if r3_enabled: @@ -990,7 +1029,11 @@ def _forward_loop(self, micro_batches, loss_fn, loss_fn_params, *, if old_lp is not None: old_lp = old_lp.view(-1).to(new_lp.device) _labels = micro_batch.get("labels", micro_batch.get("target_tokens")) - _valid = (_labels.view(-1) != IGNORE_INDEX) if _labels is not None else torch.ones_like(new_lp, dtype=torch.bool) + _valid = ( + (_labels.view(-1) != IGNORE_INDEX) + if _labels is not None + else torch.ones_like(new_lp, dtype=torch.bool) + ) log_ratio = new_lp - old_lp k3_vals = (torch.exp(log_ratio) - log_ratio - 1.0).masked_fill(~_valid, 0.0) # position_ids and _original_position_ids are both kept @@ -1006,12 +1049,14 @@ def _forward_loop(self, micro_batches, loss_fn, loss_fn_params, *, ps = get_parallel_state() cp_rank = ps.ulysses_rank if ps.ulysses_enabled else 0 start = cp_rank * local_len - _pos_flat = _pos_flat[start:start + local_len] - deferred_k3.append({ - "k3_values": k3_vals.cpu(), - "valid_mask": _valid.cpu(), - "position_ids": _pos_flat.cpu(), - }) + _pos_flat = _pos_flat[start : start + local_len] + deferred_k3.append( + { + "k3_values": k3_vals.cpu(), + "valid_mask": _valid.cpu(), + "position_ids": _pos_flat.cpu(), + } + ) # Gradient accumulation — raw (unnormalized) backward. # Normalization by total accumulated valid tokens is deferred to optim_step. @@ -1086,9 +1131,7 @@ def _forward_loop(self, micro_batches, loss_fn, loss_fn_params, *, if deferred_k3: all_per_sample_k3 = [] for entry in deferred_k3: - per_sample = self._compute_per_sample_k3( - entry["k3_values"], entry["valid_mask"], entry["position_ids"] - ) + per_sample = self._compute_per_sample_k3(entry["k3_values"], entry["valid_mask"], entry["position_ids"]) all_per_sample_k3.extend(per_sample) result["per_sample_k3"] = all_per_sample_k3 @@ -1266,7 +1309,7 @@ def forward_backward( # Diagnostic: log SGLang vs Xorl logprobs comparison for first micro-batch if batch_idx == 0: orig_lp = micro_batch["logprobs"] - valid_mask = (labels.view(-1) != IGNORE_INDEX) + valid_mask = labels.view(-1) != IGNORE_INDEX if orig_lp is not None and valid_mask.any(): orig_flat = orig_lp.to(ref_logprobs.device).view(-1) ref_flat = ref_logprobs.view(-1) @@ -1314,16 +1357,16 @@ def forward_backward( "global_valid_tokens": gvt, } # Accumulate valid tokens for deferred normalization at optim_step - self._accumulated_valid_tokens[model_id] = ( - self._accumulated_valid_tokens.get(model_id, 0) + gvt - ) + self._accumulated_valid_tokens[model_id] = self._accumulated_valid_tokens.get(model_id, 0) + gvt # R3 cleanup for PP path (stage management handled by _pp_forward) if r3_enabled: self._routing_handler.cleanup() else: # Standard forward-backward via unified loop result = self._forward_loop( - micro_batches, loss_fn, loss_fn_params, + micro_batches, + loss_fn, + loss_fn_params, compute_backward=True, r3_enabled=r3_enabled, model_id=model_id, @@ -1403,7 +1446,9 @@ def forward( r3_enabled = self._routing_handler.setup(micro_batches, routed_experts, routed_expert_logits) result = self._forward_loop( - micro_batches, loss_fn, loss_fn_params, + micro_batches, + loss_fn, + loss_fn_params, compute_backward=False, r3_enabled=r3_enabled, ) @@ -1483,7 +1528,9 @@ def optim_step( # Gradients are in the adapter's params (captured by capture_gradients in forward_backward) # Pass accumulated_valid_tokens for deferred gradient normalization grad_norm = self._adapter_manager.optim_step( - model_id, effective_lr, clip_value, + model_id, + effective_lr, + clip_value, accumulated_valid_tokens=accumulated, ) current_step = self._adapter_manager.get_global_step(model_id) @@ -1503,12 +1550,13 @@ def optim_step( if lr is not None: effective_lr = lr for param_group in self.optimizer.param_groups: - param_group['lr'] = effective_lr + param_group["lr"] = effective_lr ps = get_parallel_state() grad_norm = clip_gradients( - self.model, clip_value, + self.model, + clip_value, pp_enabled=self.pp_enabled, pp_group=ps.pp_group if self.pp_enabled else None, ) @@ -1525,7 +1573,7 @@ def optim_step( self.global_step += 1 current_step = self.global_step - current_lr = self.optimizer.param_groups[0]['lr'] + current_lr = self.optimizer.param_groups[0]["lr"] # Collect mean grad_norm across data parallel group for logging grad_norm = all_reduce(grad_norm, group=ps.fsdp_group) @@ -1576,8 +1624,8 @@ def _sync_checkpoint_state(self): """Sync mutable state to checkpoint manager before save operations.""" self._checkpoint_mgr.global_step = self.global_step self._checkpoint_mgr.global_forward_backward_step = self.global_forward_backward_step - self._checkpoint_mgr.lora_target_modules = getattr(self, 'lora_target_modules', None) - self._checkpoint_mgr.lora_alpha_value = getattr(self, 'lora_alpha_value', None) + self._checkpoint_mgr.lora_target_modules = getattr(self, "lora_target_modules", None) + self._checkpoint_mgr.lora_alpha_value = getattr(self, "lora_alpha_value", None) def _sync_from_checkpoint_state(self): """Sync state back from checkpoint manager after load operations.""" diff --git a/src/xorl/server/runner/runner_dispatcher.py b/src/xorl/server/runner/runner_dispatcher.py index c55e14a0..f3c7a394 100644 --- a/src/xorl/server/runner/runner_dispatcher.py +++ b/src/xorl/server/runner/runner_dispatcher.py @@ -70,28 +70,28 @@ from xorl.data.collators import TextSequenceShardCollator from xorl.distributed.parallel_state import get_parallel_state -from xorl.server.protocol.orchestrator_runner import ( - RunnerDispatchCommand, - RunnerResponse, -) from xorl.server.protocol.operations import ( + EmptyData, + LoadStateData, ModelPassData, OptimStepData, - SaveStateData, - SaveLoraOnlyData, - LoadStateData, SaveFullWeightsData, - EmptyData, + SaveLoraOnlyData, + SaveStateData, +) +from xorl.server.protocol.orchestrator_runner import ( + RunnerDispatchCommand, + RunnerResponse, ) -from xorl.server.runner.utils import Rank0Protocol +from xorl.server.runner.adapters import AdapterCoordinator +from xorl.server.runner.model_runner import ModelRunner from xorl.server.runner.utils import ( + Rank0Protocol, apply_sequence_sharding, convert_batch_to_tensors, simple_sequence_shard, validate_batch_shapes, ) -from xorl.server.runner.adapters import AdapterCoordinator -from xorl.server.runner.model_runner import ModelRunner from xorl.server.weight_sync.handler import WeightSyncHandler @@ -187,8 +187,7 @@ def __init__( self._worker_error: Optional[str] = None logger.info( - f"RunnerDispatcher initialized (rank={rank}/{world_size}, " - f"bind_address={bind_address}, device={device})" + f"RunnerDispatcher initialized (rank={rank}/{world_size}, bind_address={bind_address}, device={device})" ) # Operations that participate in cross-rank error sync. @@ -430,7 +429,10 @@ async def _handle_request_rank0(self, request: RunnerDispatchCommand) -> RunnerR if handler: result = await handler({}) return RunnerResponse( - request_id=request.message_id, success=True, result=result, execution_time=time.time() - start_time + request_id=request.message_id, + success=True, + result=result, + execution_time=time.time() - start_time, ) # Prepare command dict (custom prepare or default pass-through) @@ -470,7 +472,8 @@ async def _handle_request_rank0(self, request: RunnerDispatchCommand) -> RunnerR if cross_rank_error: self._worker_error = None return RunnerResponse( - request_id=request.message_id, success=False, + request_id=request.message_id, + success=False, error=f"Cross-rank error: {cross_rank_error}", execution_time=time.time() - start_time, ) @@ -548,9 +551,7 @@ async def _handle_forward(self, command_dict: Dict[str, Any]) -> Dict[str, Any]: await self._handle_compute_worker_receive(command_dict, with_backward=False) return {} - async def _handle_compute_rank0_scatter( - self, command_dict: Dict[str, Any], with_backward: bool - ) -> Dict[str, Any]: + async def _handle_compute_rank0_scatter(self, command_dict: Dict[str, Any], with_backward: bool) -> Dict[str, Any]: """Rank 0: select own batches from broadcast data, run compute, gather metrics. Uses broadcast-and-select pattern: all ranks received the full payload via @@ -585,9 +586,15 @@ async def _handle_compute_rank0_scatter( cp_enabled = parallel_state.cp_enabled result = self._execute_and_gather( - my_batches, loss_fn, loss_fn_params, routed_experts, - cp_enabled, parallel_state, - with_backward=with_backward, model_id=model_id, is_rank0=True, + my_batches, + loss_fn, + loss_fn_params, + routed_experts, + cp_enabled, + parallel_state, + with_backward=with_backward, + model_id=model_id, + is_rank0=True, routed_expert_logits=routed_expert_logits, ) @@ -598,9 +605,7 @@ async def _handle_compute_rank0_scatter( return result - async def _handle_compute_worker_receive( - self, command_dict: Dict[str, Any], with_backward: bool - ) -> None: + async def _handle_compute_worker_receive(self, command_dict: Dict[str, Any], with_backward: bool) -> None: """Worker ranks: select own batches from broadcast data, run compute, participate in collective metrics. Uses broadcast-and-select pattern: workers get the full payload via @@ -634,25 +639,46 @@ async def _handle_compute_worker_receive( cp_enabled = parallel_state.cp_enabled self._execute_and_gather( - my_batches, loss_fn, loss_fn_params, routed_experts, - cp_enabled, parallel_state, - with_backward=with_backward, model_id=model_id, is_rank0=False, + my_batches, + loss_fn, + loss_fn_params, + routed_experts, + cp_enabled, + parallel_state, + with_backward=with_backward, + model_id=model_id, + is_rank0=False, routed_expert_logits=routed_expert_logits, ) # -- Helpers for compute handlers -- - def _execute_and_gather(self, my_batches, loss_fn, loss_fn_params, routed_experts, - cp_enabled, parallel_state, *, with_backward, model_id, is_rank0, - routed_expert_logits=None): + def _execute_and_gather( + self, + my_batches, + loss_fn, + loss_fn_params, + routed_experts, + cp_enabled, + parallel_state, + *, + with_backward, + model_id, + is_rank0, + routed_expert_logits=None, + ): """Shard batches, execute compute, gather IS metrics. Shared by rank-0 and workers.""" my_batches, routed_experts = self._shard_and_slice_batches( my_batches, routed_experts, cp_enabled, parallel_state ) result = self._execute_compute( - my_batches, loss_fn, loss_fn_params, routed_experts, - with_backward=with_backward, model_id=model_id, + my_batches, + loss_fn, + loss_fn_params, + routed_experts, + with_backward=with_backward, + model_id=model_id, routed_expert_logits=routed_expert_logits, ) del my_batches @@ -731,7 +757,7 @@ def _select_and_prepare_batches(self, raw_batches, routed_experts=None, routed_e remainder = num_batches % dp_size start_idx, my_real_count = self._dp_batch_range(dp_rank, base_count, remainder) - my_raw_batches = raw_batches[start_idx:start_idx + my_real_count] + my_raw_batches = raw_batches[start_idx : start_idx + my_real_count] # Convert only this rank's batches to tensors my_batches = [self._convert_batch_to_tensors(b) for b in my_raw_batches] @@ -753,14 +779,12 @@ def _select_and_prepare_batches(self, raw_batches, routed_experts=None, routed_e for bi in range(g_start, g_start + g_count): datum_offset += raw_batches[bi].get("num_samples", 1) - dp_datum_count = sum( - raw_batches[start_idx + i].get("num_samples", 1) for i in range(my_real_count) - ) + dp_datum_count = sum(raw_batches[start_idx + i].get("num_samples", 1) for i in range(my_real_count)) if routed_experts is not None: - routed_experts_slice = routed_experts[datum_offset:datum_offset + dp_datum_count] + routed_experts_slice = routed_experts[datum_offset : datum_offset + dp_datum_count] if routed_expert_logits is not None: - routed_expert_logits_slice = routed_expert_logits[datum_offset:datum_offset + dp_datum_count] + routed_expert_logits_slice = routed_expert_logits[datum_offset : datum_offset + dp_datum_count] logger.debug( f"Rank {self.rank}: _select_and_prepare_batches: dp_rank={dp_rank}/{dp_size}, " @@ -803,34 +827,25 @@ def _shard_and_slice_batches( if isinstance(original_pos_ids, torch.Tensor): batch["_original_position_ids"] = original_pos_ids.clone() else: - batch["_original_position_ids"] = torch.tensor( - original_pos_ids, dtype=torch.long - ) + batch["_original_position_ids"] = torch.tensor(original_pos_ids, dtype=torch.long) sharded_batch = self._apply_sequence_sharding(batch) sharded_batches.append(sharded_batch) except Exception as e: shapes = { - k: tuple(v.shape) if isinstance(v, torch.Tensor) else type(v).__name__ - for k, v in batch.items() + k: tuple(v.shape) if isinstance(v, torch.Tensor) else type(v).__name__ for k, v in batch.items() } - logger.error( - f"Rank {self.rank}: Sharding failed on batch {i}. " - f"Shapes: {shapes}. Error: {e}" - ) + logger.error(f"Rank {self.rank}: Sharding failed on batch {i}. Shapes: {shapes}. Error: {e}") raise my_batches = sharded_batches - logger.debug( - f"Rank {self.rank}: Applied sequence sharding locally " - f"(cp_rank={parallel_state.cp_rank})" - ) + logger.debug(f"Rank {self.rank}: Applied sequence sharding locally (cp_rank={parallel_state.cp_rank})") # Slice routed_experts for this rank's datum subset if routed_experts is not None and my_batches: r3_offset = my_batches[0].pop("_r3_datum_offset", None) r3_count = my_batches[0].pop("_r3_datum_count", None) if r3_offset is not None and r3_count is not None: - routed_experts = routed_experts[r3_offset:r3_offset + r3_count] + routed_experts = routed_experts[r3_offset : r3_offset + r3_count] return my_batches, routed_experts @@ -848,18 +863,22 @@ def _execute_compute( """Execute forward or forward+backward on the model runner.""" if with_backward: return self.trainer.forward_backward( - my_batches, loss_fn, loss_fn_params, - model_id=model_id, routed_experts=routed_experts, + my_batches, + loss_fn, + loss_fn_params, + model_id=model_id, + routed_experts=routed_experts, routed_expert_logits=routed_expert_logits, ) return self.trainer.forward( - my_batches, loss_fn, loss_fn_params, routed_experts=routed_experts, + my_batches, + loss_fn, + loss_fn_params, + routed_experts=routed_experts, routed_expert_logits=routed_expert_logits, ) - def _gather_is_metrics( - self, result: Dict[str, Any], cp_enabled: bool, *, is_rank0: bool - ) -> None: + def _gather_is_metrics(self, result: Dict[str, Any], cp_enabled: bool, *, is_rank0: bool) -> None: """Gather importance-sampling metrics across ranks via all_gather. All ranks must call this together when SP is enabled. @@ -879,9 +898,7 @@ def _gather_is_metrics( for key in list(rank_result.keys()): if key.startswith("is_") and key not in result: result[key] = rank_result[key] - logger.debug( - f"Rank {self.rank}: Copied IS metric '{key}' from rank {i}" - ) + logger.debug(f"Rank {self.rank}: Copied IS metric '{key}' from rank {i}") logger.debug(f"Rank {self.rank}: Final result keys: {list(result.keys())}") del all_results @@ -976,8 +993,7 @@ async def _handle_save_state(self, command_dict: Dict[str, Any]) -> Dict[str, An model_id = p.model_id or "default" logger.debug( - f"Rank {self.rank}: Saving state to {checkpoint_path}, " - f"model_id={model_id}, save_optimizer={save_optimizer}" + f"Rank {self.rank}: Saving state to {checkpoint_path}, model_id={model_id}, save_optimizer={save_optimizer}" ) # For LoRA models with save_optimizer=False (sampler weights), use fast PEFT-compatible save @@ -1225,7 +1241,6 @@ async def _handle_health_check(self) -> Dict[str, Any]: async def _handle_sync_inference_weights(self, command_dict: Dict[str, Any]) -> Dict[str, Any]: return await self._weight_sync_handler.handle_sync_inference_weights(command_dict) - # ======================================================================== # Adapter Operations (delegated to AdapterCoordinator) # ======================================================================== diff --git a/src/xorl/server/runner/setup.py b/src/xorl/server/runner/setup.py index 408ccf13..ede90fe4 100644 --- a/src/xorl/server/runner/setup.py +++ b/src/xorl/server/runner/setup.py @@ -18,9 +18,9 @@ import yaml from xorl.arguments import Arguments, parse_args -from xorl.server.server_arguments import ServerArguments, parse_server_args -from xorl.server.runner.runner_dispatcher import RunnerDispatcher from xorl.server.runner.model_runner import ModelRunner +from xorl.server.runner.runner_dispatcher import RunnerDispatcher +from xorl.server.server_arguments import ServerArguments, parse_server_args from xorl.utils.device import get_nccl_backend @@ -227,7 +227,7 @@ async def run_distributed_worker_from_server_args(server_args: ServerArguments): # Resolve auto bind address if server_args.worker_bind_address == "auto": with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(('', 0)) + s.bind(("", 0)) server_args.worker_bind_address = f"tcp://127.0.0.1:{s.getsockname()[1]}" config = server_args.to_config_dict() @@ -293,9 +293,7 @@ def main(): if not config_path: print("Error: Config file path required as first positional argument") - print( - "Usage: python -m xorl.server.runner.runner_dispatcher config.yaml [--worker_bind_address tcp://...]" - ) + print("Usage: python -m xorl.server.runner.runner_dispatcher config.yaml [--worker_bind_address tcp://...]") sys.exit(1) # Detect config format and parse accordingly diff --git a/src/xorl/server/runner/utils/__init__.py b/src/xorl/server/runner/utils/__init__.py index 5b9de7a7..ff120a65 100644 --- a/src/xorl/server/runner/utils/__init__.py +++ b/src/xorl/server/runner/utils/__init__.py @@ -4,10 +4,11 @@ simple_sequence_shard, validate_batch_shapes, ) -from xorl.server.runner.utils.validation import run_self_test, validate_token_ids +from xorl.server.runner.utils.moe_metrics import MoeMetricsTracker from xorl.server.runner.utils.rank0_protocol import Rank0Protocol from xorl.server.runner.utils.routing_replay_handler import RoutingReplayHandler -from xorl.server.runner.utils.moe_metrics import MoeMetricsTracker +from xorl.server.runner.utils.validation import run_self_test, validate_token_ids + __all__ = [ "apply_sequence_sharding", diff --git a/src/xorl/server/runner/utils/batch_utils.py b/src/xorl/server/runner/utils/batch_utils.py index bffe4466..bd258182 100644 --- a/src/xorl/server/runner/utils/batch_utils.py +++ b/src/xorl/server/runner/utils/batch_utils.py @@ -57,7 +57,9 @@ def convert_batch_to_tensors(batch: Dict[str, Any], rank: int = 0) -> Dict[str, if isinstance(value[0], list): # This is a list of sequences - pad them max_len = max(len(seq) for seq in value) - pad_value = -100 if key in ("labels", "target_tokens") else 0 # Use -100 for labels/target_tokens (IGNORE_INDEX) + pad_value = ( + -100 if key in ("labels", "target_tokens") else 0 + ) # Use -100 for labels/target_tokens (IGNORE_INDEX) padded = [] for seq in value: padded_seq = seq + [pad_value] * (max_len - len(seq)) @@ -71,9 +73,7 @@ def convert_batch_to_tensors(batch: Dict[str, Any], rank: int = 0) -> Dict[str, f"Rank {rank}: Padded and converted {key}: {len(value)} sequences, max_len={max_len}, dtype={dtype}" ) except Exception as e2: - logger.warning( - f"Rank {rank}: Failed to convert {key} even after padding: {e2}, keeping as-is" - ) + logger.warning(f"Rank {rank}: Failed to convert {key} even after padding: {e2}, keeping as-is") converted_batch[key] = value else: logger.warning(f"Rank {rank}: Failed to convert {key} to tensor: {e}, keeping as-is") @@ -121,8 +121,7 @@ def validate_batch_shapes(batch: Dict[str, Any], rank: int = 0, batch_idx: int = unique_lengths = set(seq_lengths.values()) if len(unique_lengths) > 1: logger.error( - f"Rank {rank}: Batch {batch_idx} has INCONSISTENT sequence lengths: {seq_lengths}. " - f"Full shapes: {shapes}" + f"Rank {rank}: Batch {batch_idx} has INCONSISTENT sequence lengths: {seq_lengths}. Full shapes: {shapes}" ) return False @@ -226,8 +225,7 @@ def pad_and_slice(tensor, pad_value=0): sharded_batch[key] = value logger.debug( - f"Rank {rank}: Simple sequence shard: {seq_len} -> {cp_chunk_size} " - f"(cp_rank={cp_rank}, cp_size={cp_size})" + f"Rank {rank}: Simple sequence shard: {seq_len} -> {cp_chunk_size} (cp_rank={cp_rank}, cp_size={cp_size})" ) return sharded_batch diff --git a/src/xorl/server/runner/utils/moe_metrics.py b/src/xorl/server/runner/utils/moe_metrics.py index 4293f081..6e1e1e11 100644 --- a/src/xorl/server/runner/utils/moe_metrics.py +++ b/src/xorl/server/runner/utils/moe_metrics.py @@ -13,11 +13,13 @@ import torch + try: from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import ( enable_expert_metrics, get_expert_metrics, ) + _HAS_MOE_METRICS = True except ImportError: _HAS_MOE_METRICS = False @@ -104,7 +106,9 @@ def collect(self, step: int, forward_backward_time: float) -> dict: metrics["expert_load"] = { "num_moe_layers": len(expert_metrics), "total_tokens": total_tokens, - "mean_imbalance_ratio": sum(all_imbalance_ratios) / len(all_imbalance_ratios) if all_imbalance_ratios else 0, + "mean_imbalance_ratio": sum(all_imbalance_ratios) / len(all_imbalance_ratios) + if all_imbalance_ratios + else 0, "max_imbalance_ratio": max(all_imbalance_ratios) if all_imbalance_ratios else 0, "mean_max_load": sum(all_max_loads) / len(all_max_loads) if all_max_loads else 0, "mean_min_load": sum(all_min_loads) / len(all_min_loads) if all_min_loads else 0, diff --git a/src/xorl/server/runner/utils/rank0_protocol.py b/src/xorl/server/runner/utils/rank0_protocol.py index 89b466b9..08dc6186 100644 --- a/src/xorl/server/runner/utils/rank0_protocol.py +++ b/src/xorl/server/runner/utils/rank0_protocol.py @@ -18,8 +18,8 @@ import zmq from xorl.server.protocol.orchestrator_runner import ( - RunnerDispatchCommand, RunnerAck, + RunnerDispatchCommand, RunnerReady, RunnerResponse, deserialize_message, @@ -28,6 +28,7 @@ from xorl.server.utils.network import get_local_ip, parse_zmq_address, write_address_file from xorl.server.utils.zmq_channels import AsyncRouterChannel + logger = logging.getLogger(__name__) diff --git a/src/xorl/server/runner/utils/routing_replay_handler.py b/src/xorl/server/runner/utils/routing_replay_handler.py index 526072cb..b66a5648 100644 --- a/src/xorl/server/runner/utils/routing_replay_handler.py +++ b/src/xorl/server/runner/utils/routing_replay_handler.py @@ -24,7 +24,7 @@ import base64 import logging import math -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Union import numpy as np import torch @@ -37,8 +37,10 @@ set_replay_stage, ) + try: from xorl.models.layers.moe import MoEBlock + _HAS_MOE_BLOCK = True except ImportError: _HAS_MOE_BLOCK = False @@ -179,10 +181,7 @@ def _infer_shape(arr: np.ndarray, num_moe_layers: int) -> Optional[np.ndarray]: if total_elements % (num_moe_layers * topk) == 0: num_tokens = total_elements // (num_moe_layers * topk) return arr.reshape(num_tokens, num_moe_layers, topk) - logger.warning( - f"R3: Cannot infer shape for {total_elements} elements " - f"with {num_moe_layers} layers" - ) + logger.warning(f"R3: Cannot infer shape for {total_elements} elements with {num_moe_layers} layers") return None def fill_routing_replay( @@ -244,9 +243,7 @@ def fill_routing_replay( ) # Build per-micro-batch routing tensors, handling packing + SP slicing - per_mb_routing = self._build_per_mb_routing( - micro_batches, decoded_routing, num_layers_in_data, topk - ) + per_mb_routing = self._build_per_mb_routing(micro_batches, decoded_routing, num_layers_in_data, topk) if not per_mb_routing: logger.warning("R3: Empty routing data after processing") @@ -271,17 +268,13 @@ def fill_routing_replay( decoded_weights.append(decoded) if decoded_weights: - per_mb_weights = self._build_per_mb_routing( - micro_batches, decoded_weights, num_layers_in_data, topk - ) + per_mb_weights = self._build_per_mb_routing(micro_batches, decoded_weights, num_layers_in_data, topk) for mb_idx, mb_weights_tensor in enumerate(per_mb_weights): num_layers_to_use_w = min(num_moe_layers, mb_weights_tensor.shape[1]) for moe_idx in range(num_layers_to_use_w): layer_weights = mb_weights_tensor[:, moe_idx, :].float() moe_blocks[moe_idx]._routing_replay.record_weights(layer_weights) - logger.debug( - f"R3: Pre-populated routing weights for {len(per_mb_weights)} micro-batches" - ) + logger.debug(f"R3: Pre-populated routing weights for {len(per_mb_weights)} micro-batches") logger.debug( f"R3: Pre-populated {len(per_mb_routing)} micro-batches x " diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index e6ea2f2d..887f2eac 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -7,8 +7,8 @@ parameters like batch size, epochs, and optimizer settings. """ -from dataclasses import dataclass, field, asdict -from typing import Optional, Dict, List, Literal, Any +from dataclasses import dataclass, field +from typing import Any, Dict, List, Literal, Optional @dataclass @@ -44,96 +44,82 @@ class ServerArguments: # ======================================================================== model_path: Optional[str] = field( - default=None, - metadata={"help": "Path to pre-trained model (HF Hub or local path)"} + default=None, metadata={"help": "Path to pre-trained model (HF Hub or local path)"} ) model_name: Optional[str] = field( default=None, - metadata={"help": "Model identifier for validation (e.g., 'Qwen/Qwen3-32B'). Defaults to model_path if not specified."} + metadata={ + "help": "Model identifier for validation (e.g., 'Qwen/Qwen3-32B'). Defaults to model_path if not specified." + }, ) - config_path: Optional[str] = field( - default=None, - metadata={"help": "Path to model config. Defaults to model_path"} - ) + config_path: Optional[str] = field(default=None, metadata={"help": "Path to model config. Defaults to model_path"}) - tokenizer_path: Optional[str] = field( - default=None, - metadata={"help": "Path to tokenizer. Defaults to config_path"} - ) + tokenizer_path: Optional[str] = field(default=None, metadata={"help": "Path to tokenizer. Defaults to config_path"}) attn_implementation: Optional[Literal["eager", "sdpa", "native", "flash_attention_3", "flash_attention_4"]] = field( default="flash_attention_3", - metadata={"help": "Attention implementation. 'native': PyTorch SDPA+cuDNN (no deps, Hopper+Blackwell). " - "'flash_attention_3': FA3 (Hopper). 'flash_attention_4': FA4 CUTE (Hopper+Blackwell)."} + metadata={ + "help": "Attention implementation. 'native': PyTorch SDPA+cuDNN (no deps, Hopper+Blackwell). " + "'flash_attention_3': FA3 (Hopper). 'flash_attention_4': FA4 CUTE (Hopper+Blackwell)." + }, ) moe_implementation: Optional[Literal[None, "eager", "triton", "native", "quack"]] = field( default=None, - metadata={"help": "MoE implementation. 'triton' uses Triton group GEMM kernels, 'native' uses torch._grouped_mm, 'quack' uses quack kernels."} + metadata={ + "help": "MoE implementation. 'triton' uses Triton group GEMM kernels, 'native' uses torch._grouped_mm, 'quack' uses quack kernels." + }, ) ep_dispatch: str = field( default="alltoall", - metadata={"help": "EP dispatch strategy: 'alltoall' (default) or 'deepep' (NVLink-optimized)."} + metadata={"help": "EP dispatch strategy: 'alltoall' (default) or 'deepep' (NVLink-optimized)."}, ) deepep_buffer_size_gb: float = field( - default=2.0, - metadata={"help": "DeepEP buffer size in GB (effective when ep_dispatch='deepep')."} + default=2.0, metadata={"help": "DeepEP buffer size in GB (effective when ep_dispatch='deepep')."} ) deepep_num_sms: int = field( - default=20, - metadata={"help": "Number of SMs for DeepEP communication kernels (must be even, default 20)."} + default=20, metadata={"help": "Number of SMs for DeepEP communication kernels (must be even, default 20)."} ) deepep_async_combine: bool = field( - default=False, - metadata={"help": "Enable async combine for DeepEP (overlap combine with next layer's compute)."} + default=False, metadata={"help": "Enable async combine for DeepEP (overlap combine with next layer's compute)."} ) # SGLang numerical alignment flags router_fp32: bool = field( - default=True, - metadata={"help": "Upcast MoE router gate computation to float32 for numerical stability."} + default=True, metadata={"help": "Upcast MoE router gate computation to float32 for numerical stability."} ) lm_head_fp32: bool = field( - default=True, - metadata={"help": "Upcast LM head logits computation to float32 for numerical stability."} + default=True, metadata={"help": "Upcast LM head logits computation to float32 for numerical stability."} ) rmsnorm_native: bool = field( - default=False, - metadata={"help": "Use native RMSNorm (no fused kernels) for SGLang alignment."} + default=False, metadata={"help": "Use native RMSNorm (no fused kernels) for SGLang alignment."} ) activation_native: bool = field( - default=False, - metadata={"help": "Use native SiLU instead of fused Triton kernel for SGLang alignment."} + default=False, metadata={"help": "Use native SiLU instead of fused Triton kernel for SGLang alignment."} ) rope_native: bool = field( - default=False, - metadata={"help": "Use naive RoPE implementation instead of flash_attn fused kernel."} + default=False, metadata={"help": "Use naive RoPE implementation instead of flash_attn fused kernel."} ) attention_cast_bf16: bool = field( - default=False, - metadata={"help": "Explicitly cast Q/K to bfloat16 after RoPE for SGLang alignment."} + default=False, metadata={"help": "Explicitly cast Q/K to bfloat16 after RoPE for SGLang alignment."} ) # Multimodal model configuration - foundation: Dict[str, str] = field( - default_factory=dict, - metadata={"help": "Foundation model extra config"} - ) + foundation: Dict[str, str] = field(default_factory=dict, metadata={"help": "Foundation model extra config"}) encoders: Dict[Literal["image", "video", "audio"], Dict[str, str]] = field( - default_factory=dict, - metadata={"help": "Multimodal encoder config"} + default_factory=dict, metadata={"help": "Multimodal encoder config"} ) # ======================================================================== @@ -142,37 +128,21 @@ class ServerArguments: data_parallel_mode: Optional[Literal["none", "ddp", "fsdp2"]] = field( default="fsdp2", - metadata={"help": "Data parallelism mode. Use 'none' for single GPU without any parallelization."} + metadata={"help": "Data parallelism mode. Use 'none' for single GPU without any parallelization."}, ) - ulysses_parallel_size: int = field( - default=1, - metadata={"help": "Ulysses sequence parallelism size"} - ) + ulysses_parallel_size: int = field(default=1, metadata={"help": "Ulysses sequence parallelism size"}) - expert_parallel_size: int = field( - default=1, - metadata={"help": "Expert parallelism size for MoE models"} - ) + expert_parallel_size: int = field(default=1, metadata={"help": "Expert parallelism size for MoE models"}) - data_parallel_replicate_size: int = field( - default=1, - metadata={"help": "Data parallel replicate size (HSDP)"} - ) + data_parallel_replicate_size: int = field(default=1, metadata={"help": "Data parallel replicate size (HSDP)"}) - data_parallel_shard_size: int = field( - default=1, - metadata={"help": "Data parallel shard size (FSDP)"} - ) + data_parallel_shard_size: int = field(default=1, metadata={"help": "Data parallel shard size (FSDP)"}) - pipeline_parallel_size: int = field( - default=1, - metadata={"help": "Pipeline parallelism size. 1 = disabled."} - ) + pipeline_parallel_size: int = field(default=1, metadata={"help": "Pipeline parallelism size. 1 = disabled."}) pipeline_parallel_schedule: str = field( - default="1F1B", - metadata={"help": "Pipeline parallelism schedule: '1F1B' or 'GPipe'."} + default="1F1B", metadata={"help": "Pipeline parallelism schedule: '1F1B' or 'GPipe'."} ) pp_variable_seq_lengths: bool = field( default=True, @@ -183,96 +153,66 @@ class ServerArguments: "static sample_packing_sequence_len. Each unique seq_len gets its own " "cached PipelineStage so P2P buffers always match the actual shape." ) - } + }, ) - tensor_parallel_size: int = field( - default=1, - metadata={"help": "Tensor parallelism size"} - ) + tensor_parallel_size: int = field(default=1, metadata={"help": "Tensor parallelism size"}) - ringattn_parallel_size: int = field( - default=1, - metadata={"help": "Ring attention parallel size"} - ) + ringattn_parallel_size: int = field(default=1, metadata={"help": "Ring attention parallel size"}) cp_fsdp_mode: str = field( - default="all", - metadata={"help": "Sequence parallel FSDP mode: 'all', 'ulysses_only', 'ring_only', 'none'"} + default="all", metadata={"help": "Sequence parallel FSDP mode: 'all', 'ulysses_only', 'ring_only', 'none'"} ) basic_modules: Optional[List[str]] = field( - default_factory=list, - metadata={"help": "Basic modules to shard in FSDP"} + default_factory=list, metadata={"help": "Basic modules to shard in FSDP"} ) merge_qkv: bool = field( - default=True, - metadata={"help": "Whether to merge QKV projections. Set False for tensor parallelism."} + default=True, metadata={"help": "Whether to merge QKV projections. Set False for tensor parallelism."} ) # ======================================================================== # Memory & Performance # ======================================================================== - seed: int = field( - default=42, - metadata={"help": "Random seed for reproducibility"} - ) + seed: int = field(default=42, metadata={"help": "Random seed for reproducibility"}) - enable_mixed_precision: bool = field( - default=True, - metadata={"help": "Enable mixed precision training"} - ) + enable_mixed_precision: bool = field(default=True, metadata={"help": "Enable mixed precision training"}) - enable_gradient_checkpointing: bool = field( - default=True, - metadata={"help": "Enable gradient checkpointing"} - ) + enable_gradient_checkpointing: bool = field(default=True, metadata={"help": "Enable gradient checkpointing"}) - enable_full_shard: bool = field( - default=True, - metadata={"help": "Enable full parameter sharding (FSDP)"} - ) + enable_full_shard: bool = field(default=True, metadata={"help": "Enable full parameter sharding (FSDP)"}) - enable_activation_offload: bool = field( - default=False, - metadata={"help": "Enable activation CPU offloading"} - ) + enable_activation_offload: bool = field(default=False, metadata={"help": "Enable activation CPU offloading"}) - enable_compile: bool = field( - default=False, - metadata={"help": "Enable torch.compile for model forward pass"} - ) + enable_compile: bool = field(default=False, metadata={"help": "Enable torch.compile for model forward pass"}) enable_reentrant: bool = field( - default=False, - metadata={"help": "Use reentrant gradient checkpointing (default: non-reentrant)"} + default=False, metadata={"help": "Use reentrant gradient checkpointing (default: non-reentrant)"} ) enable_forward_prefetch: bool = field( - default=False, - metadata={"help": "Enable FSDP forward prefetch for overlapping compute and communication"} + default=False, metadata={"help": "Enable FSDP forward prefetch for overlapping compute and communication"} ) reshard_after_forward: bool = field( - default=True, - metadata={"help": "Reshard parameters after forward pass in FSDP2"} + default=True, metadata={"help": "Reshard parameters after forward pass in FSDP2"} ) load_weights_mode: str = field( - default="auto", - metadata={"help": "Weight loading mode: 'auto', 'safetensors', 'dcp'"} + default="auto", metadata={"help": "Weight loading mode: 'auto', 'safetensors', 'dcp'"} ) init_device: Optional[Literal["cpu", "meta", "cuda"]] = field( - default="meta", - metadata={"help": "Device for model initialization"} + default="meta", metadata={"help": "Device for model initialization"} ) ce_mode: Literal["eager", "compiled"] = field( default="compiled", - metadata={"help": "Cross-entropy implementation: 'compiled' (RECOMMENDED, torch.compile) or 'eager' (baseline, may OOM at 32K)"} + metadata={ + "help": "Cross-entropy implementation: 'compiled' (RECOMMENDED, torch.compile) or 'eager' (baseline, may OOM at 32K)" + }, ) # ======================================================================== @@ -291,7 +231,9 @@ class ServerArguments: muon_lr: float = field( default=0.02, - metadata={"help": "Learning rate for Muon parameter groups (2D+ weight matrices). Only used when optimizer='muon'."}, + metadata={ + "help": "Learning rate for Muon parameter groups (2D+ weight matrices). Only used when optimizer='muon'." + }, ) muon_momentum: float = field( @@ -311,8 +253,10 @@ class ServerArguments: muon_adjust_lr_fn: Optional[str] = field( default=None, - metadata={"help": "LR adjustment for Muon. 'original': scale by sqrt(max(1,A/B)). " - "'match_rms_adamw': scale by 0.2*sqrt(max(A,B)) so Muon can reuse AdamW LR/WD."}, + metadata={ + "help": "LR adjustment for Muon. 'original': scale by sqrt(max(1,A/B)). " + "'match_rms_adamw': scale by 0.2*sqrt(max(A,B)) so Muon can reuse AdamW LR/WD." + }, ) # ======================================================================== @@ -321,57 +265,46 @@ class ServerArguments: output_dir: str = field( default="outputs", - metadata={"help": "Output directory for checkpoints, sampler weights, and logs (must be on shared filesystem for multi-node)"} + metadata={ + "help": "Output directory for checkpoints, sampler weights, and logs (must be on shared filesystem for multi-node)" + }, ) storage_limit: str = field( default="10TB", - metadata={"help": "Maximum disk usage for output_dir (e.g., '1GB', '500MB', '10GB'). Save operations will fail with StorageLimitError when limit is exceeded. Default: 10TB."} + metadata={ + "help": "Maximum disk usage for output_dir (e.g., '1GB', '500MB', '10GB'). Save operations will fail with StorageLimitError when limit is exceeded. Default: 10TB." + }, ) idle_session_timeout: float = field( default=7200.0, - metadata={"help": "Idle session timeout in seconds. Sessions inactive for this duration will be automatically cleaned up. Default: 7200 (2 hours)."} + metadata={ + "help": "Idle session timeout in seconds. Sessions inactive for this duration will be automatically cleaned up. Default: 7200 (2 hours)." + }, ) - load_checkpoint_path: str = field( - default="", - metadata={"help": "Path to checkpoint to load"} - ) + load_checkpoint_path: str = field(default="", metadata={"help": "Path to checkpoint to load"}) - ckpt_manager: Optional[Literal["torch", "dcp"]] = field( - default="dcp", - metadata={"help": "Checkpoint manager type"} - ) + ckpt_manager: Optional[Literal["torch", "dcp"]] = field(default="dcp", metadata={"help": "Checkpoint manager type"}) # ======================================================================== # Logging # ======================================================================== - log_level: str = field( - default="INFO", - metadata={"help": "Logging level (DEBUG, INFO, WARNING, ERROR)"} - ) + log_level: str = field(default="INFO", metadata={"help": "Logging level (DEBUG, INFO, WARNING, ERROR)"}) - enable_self_test: bool = field( - default=False, - metadata={"help": "Enable self-test after model initialization"} - ) + enable_self_test: bool = field(default=False, metadata={"help": "Enable self-test after model initialization"}) skip_initial_checkpoint: bool = field( - default=False, - metadata={"help": "Skip saving initial checkpoint (000000) on startup"} + default=False, metadata={"help": "Skip saving initial checkpoint (000000) on startup"} ) log_gradient_norms: bool = field( - default=True, - metadata={"help": "Log gradient norms by layer type after backward pass"} + default=True, metadata={"help": "Log gradient norms by layer type after backward pass"} ) - log_router_stats: bool = field( - default=True, - metadata={"help": "Log MoE router token distribution statistics"} - ) + log_router_stats: bool = field(default=True, metadata={"help": "Log MoE router token distribution statistics"}) # ======================================================================== # Worker Configuration @@ -379,80 +312,72 @@ class ServerArguments: worker_bind_host: str = field( default="0.0.0.0", - metadata={"help": "Host for worker ZMQ ROUTER socket to bind. Use '0.0.0.0' for multi-node (accepts connections from any interface)."} + metadata={ + "help": "Host for worker ZMQ ROUTER socket to bind. Use '0.0.0.0' for multi-node (accepts connections from any interface)." + }, ) worker_bind_port: int = field( - default=5556, - metadata={"help": "Port for worker ZMQ ROUTER socket to bind (rank 0 worker)"} + default=5556, metadata={"help": "Port for worker ZMQ ROUTER socket to bind (rank 0 worker)"} ) engine_connect_host: Optional[str] = field( default=None, - metadata={"help": "Host for Engine to connect to rank 0 worker. If None, auto-discovered (localhost for single-node, file-based for multi-node)."} + metadata={ + "help": "Host for Engine to connect to rank 0 worker. If None, auto-discovered (localhost for single-node, file-based for multi-node)." + }, ) worker_bind_address: str = field( default="auto", - metadata={"help": "ZMQ ROUTER socket address to bind (rank 0 worker). 'auto' picks a free port."} + metadata={"help": "ZMQ ROUTER socket address to bind (rank 0 worker). 'auto' picks a free port."}, ) worker_connection_timeout: float = field( default=120.0, - metadata={"help": "Timeout in seconds for worker-executor connection. Increased for multi-node scenarios."} + metadata={"help": "Timeout in seconds for worker-executor connection. Increased for multi-node scenarios."}, ) - worker_max_retries: int = field( - default=3, - metadata={"help": "Maximum number of retries for failed operations"} - ) + worker_max_retries: int = field(default=3, metadata={"help": "Maximum number of retries for failed operations"}) # ======================================================================== # Data Processing Configuration # ======================================================================== sample_packing_sequence_len: int = field( - default=32000, - metadata={"help": "Maximum sequence length for sample packing (default: 32000)"} + default=32000, metadata={"help": "Maximum sequence length for sample packing (default: 32000)"} ) enable_packing: bool = field( - default=True, - metadata={"help": "Enable sample packing to combine multiple samples into one sequence"} + default=True, metadata={"help": "Enable sample packing to combine multiple samples into one sequence"} ) # ======================================================================== # LoRA Configuration # ======================================================================== - enable_lora: bool = field( - default=False, - metadata={"help": "Enable LoRA adapters for training"} - ) + enable_lora: bool = field(default=False, metadata={"help": "Enable LoRA adapters for training"}) - lora_rank: int = field( - default=32, - metadata={"help": "LoRA rank (r parameter)"} - ) + lora_rank: int = field(default=32, metadata={"help": "LoRA rank (r parameter)"}) - lora_alpha: int = field( - default=16, - metadata={"help": "LoRA alpha scaling parameter"} - ) + lora_alpha: int = field(default=16, metadata={"help": "LoRA alpha scaling parameter"}) lora_target_modules: Optional[List[str]] = field( default=None, - metadata={"help": "List of module names to apply LoRA to (e.g., ['q_proj', 'k_proj', 'v_proj', 'o_proj']). If None, uses default based on model architecture."} + metadata={ + "help": "List of module names to apply LoRA to (e.g., ['q_proj', 'k_proj', 'v_proj', 'o_proj']). If None, uses default based on model architecture." + }, ) moe_shared_lora: bool = field( - default=False, - metadata={"help": "Enable shared LoRA for MoE: share LoRA weights across experts"} + default=False, metadata={"help": "Enable shared LoRA for MoE: share LoRA weights across experts"} ) moe_hybrid_shared_lora: bool = field( default=False, - metadata={"help": "Enable hybrid shared LoRA for MoE: share lora_A for gate/up_proj, lora_B for down_proj across experts"} + metadata={ + "help": "Enable hybrid shared LoRA for MoE: share lora_A for gate/up_proj, lora_B for down_proj across experts" + }, ) # ======================================================================== @@ -460,42 +385,29 @@ class ServerArguments: # ======================================================================== enable_qlora: bool = field( - default=False, - metadata={"help": "Enable QLoRA (quantized LoRA) for memory-efficient training"} + default=False, metadata={"help": "Enable QLoRA (quantized LoRA) for memory-efficient training"} ) quant_format: str = field( - default="nvfp4", - metadata={"help": "Quantization format for QLoRA: 'nvfp4', 'block_fp8', or 'nf4'"} + default="nvfp4", metadata={"help": "Quantization format for QLoRA: 'nvfp4', 'block_fp8', or 'nf4'"} ) - quant_group_size: int = field( - default=16, - metadata={"help": "Quantization group size for QLoRA"} - ) + quant_group_size: int = field(default=16, metadata={"help": "Quantization group size for QLoRA"}) qlora_exclude_modules: Optional[List[str]] = field( - default=None, - metadata={"help": "Modules to exclude from QLoRA quantization (e.g., ['lm_head'])"} + default=None, metadata={"help": "Modules to exclude from QLoRA quantization (e.g., ['lm_head'])"} ) - merge_lora_interval: int = field( - default=0, - metadata={"help": "Merge LoRA weights every N steps (0 = never merge)"} - ) + merge_lora_interval: int = field(default=0, metadata={"help": "Merge LoRA weights every N steps (0 = never merge)"}) reset_optimizer_on_merge: bool = field( - default=False, - metadata={"help": "ReLoRA-style partial optimizer reset after each LoRA merge"} + default=False, metadata={"help": "ReLoRA-style partial optimizer reset after each LoRA merge"} ) # ======================================================================== # MoE Training Configuration # ======================================================================== - freeze_router: bool = field( - default=True, - metadata={"help": "Freeze MoE router weights during training"} - ) + freeze_router: bool = field(default=True, metadata={"help": "Freeze MoE router weights during training"}) # ======================================================================== # Inference Weight Sync Configuration @@ -503,8 +415,10 @@ class ServerArguments: sync_inference_method: Literal["nccl_broadcast"] = field( default="nccl_broadcast", - metadata={"help": "Method for syncing weights to inference endpoints: " - "'nccl_broadcast' (rank-0 broadcast via SGLang update_weights_from_distributed)"} + metadata={ + "help": "Method for syncing weights to inference endpoints: " + "'nccl_broadcast' (rank-0 broadcast via SGLang update_weights_from_distributed)" + }, ) def __post_init__(self): @@ -621,7 +535,7 @@ def to_config_dict(self) -> Dict[str, Any]: def get_world_size(self) -> int: """ Calculate world size from parallelism configuration. - + Note: EP (Expert Parallel) is NOT included in world_size calculation. EP creates a separate 2D mesh: world_size = ep_size * ep_fsdp_size where ep_fsdp_size contains all other parallelism dimensions. @@ -630,14 +544,14 @@ def get_world_size(self) -> int: Required world size (number of GPUs) """ return ( - self.pipeline_parallel_size * - self.tensor_parallel_size * - self.ringattn_parallel_size * - self.ulysses_parallel_size * - self.data_parallel_replicate_size * - self.data_parallel_shard_size + self.pipeline_parallel_size + * self.tensor_parallel_size + * self.ringattn_parallel_size + * self.ulysses_parallel_size + * self.data_parallel_replicate_size + * self.data_parallel_shard_size ) - + def get_total_gpus(self) -> int: """ Calculate total number of GPUs required for the parallelism configuration. @@ -672,7 +586,7 @@ def get_total_gpus(self) -> int: def get_ep_fsdp_size(self) -> int: """ Calculate ep_fsdp_size (the size of each expert parallel group). - + For non-EP models (ep_size=1), this equals world_size. For EP models, ep_fsdp_size = world_size / ep_size. @@ -681,18 +595,19 @@ def get_ep_fsdp_size(self) -> int: """ world_size = self.get_world_size() if self.expert_parallel_size > 1: - assert world_size % self.expert_parallel_size == 0, \ + assert world_size % self.expert_parallel_size == 0, ( f"world_size ({world_size}) must be divisible by expert_parallel_size ({self.expert_parallel_size})" + ) return world_size // self.expert_parallel_size return world_size - + def get_dp_size(self) -> int: """ Calculate data parallel size (auto-calculated from other dimensions). - + IMPORTANT: EP (Expert Parallel) is NOT included in this calculation! EP creates a separate 2D mesh and doesn't participate in the main mesh. - + dp_size is the remaining parallelism after accounting for ulysses. It's then split into dp_replicate_size and dp_shard_size. @@ -702,8 +617,7 @@ def get_dp_size(self) -> int: world_size = self.get_world_size() if world_size % self.ulysses_parallel_size != 0: raise ValueError( - f"world_size ({world_size}) must be divisible by " - f"ulysses_parallel_size ({self.ulysses_parallel_size})" + f"world_size ({world_size}) must be divisible by ulysses_parallel_size ({self.ulysses_parallel_size})" ) return world_size // self.ulysses_parallel_size @@ -719,14 +633,14 @@ def parse_server_args() -> ServerArguments: Returns: ServerArguments with all fields populated from YAML and CLI """ - from xorl.arguments import parse_args - import yaml import sys + import yaml + # Read YAML directly to get flat structure config_path = None for i, arg in enumerate(sys.argv): - if not arg.startswith('--') and i > 0 and not sys.argv[i-1].startswith('--'): + if not arg.startswith("--") and i > 0 and not sys.argv[i - 1].startswith("--"): config_path = arg break @@ -734,7 +648,7 @@ def parse_server_args() -> ServerArguments: raise ValueError("Config file path required as first positional argument") # Load YAML - with open(config_path, 'r') as f: + with open(config_path, "r") as f: config_data = yaml.safe_load(f) if not config_data: @@ -746,31 +660,31 @@ def parse_server_args() -> ServerArguments: i = 1 while i < len(sys.argv): arg = sys.argv[i] - if arg.startswith('--') and not arg.startswith('---'): + if arg.startswith("--") and not arg.startswith("---"): key_part = arg[2:] # Remove '--' # Check for --key=value format - if '=' in key_part: - key, value = key_part.split('=', 1) + if "=" in key_part: + key, value = key_part.split("=", 1) # Convert dotted notation (worker.bind_address) to flat (worker_bind_address) - key = key.replace('.', '_') + key = key.replace(".", "_") # Try to parse as number or boolean - if value.lower() in ('true', 'false'): - value = value.lower() == 'true' - elif value.replace('.', '', 1).replace('-', '', 1).isdigit(): - value = float(value) if '.' in value else int(value) + if value.lower() in ("true", "false"): + value = value.lower() == "true" + elif value.replace(".", "", 1).replace("-", "", 1).isdigit(): + value = float(value) if "." in value else int(value) cli_overrides[key] = value i += 1 - elif i + 1 < len(sys.argv) and not sys.argv[i + 1].startswith('--'): + elif i + 1 < len(sys.argv) and not sys.argv[i + 1].startswith("--"): # --key value format # Convert dotted notation (worker.bind_address) to flat (worker_bind_address) - key = key_part.replace('.', '_') + key = key_part.replace(".", "_") value = sys.argv[i + 1] # Try to parse as number or boolean - if value.lower() in ('true', 'false'): - value = value.lower() == 'true' - elif value.replace('.', '', 1).replace('-', '', 1).isdigit(): - value = float(value) if '.' in value else int(value) + if value.lower() in ("true", "false"): + value = value.lower() == "true" + elif value.replace(".", "", 1).replace("-", "", 1).isdigit(): + value = float(value) if "." in value else int(value) cli_overrides[key] = value i += 2 else: @@ -782,12 +696,11 @@ def parse_server_args() -> ServerArguments: config_data.update(cli_overrides) # Validate config keys against ServerArguments fields - valid_fields = {f.name for f in __import__('dataclasses').fields(ServerArguments)} + valid_fields = {f.name for f in __import__("dataclasses").fields(ServerArguments)} unknown_fields = set(config_data.keys()) - valid_fields if unknown_fields: raise ValueError( - f"Unrecognized config fields: {sorted(unknown_fields)}. " - f"Check your config file for typos or removed fields." + f"Unrecognized config fields: {sorted(unknown_fields)}. Check your config file for typos or removed fields." ) # Create ServerArguments diff --git a/src/xorl/server/utils/__init__.py b/src/xorl/server/utils/__init__.py index 0aeae5a7..9212452a 100644 --- a/src/xorl/server/utils/__init__.py +++ b/src/xorl/server/utils/__init__.py @@ -1,14 +1,15 @@ """Server utilities for Xorl.""" from xorl.server.utils.network import ( - get_local_ip, build_worker_bind_address, build_worker_connect_address, + get_local_ip, parse_zmq_address, - write_address_file, read_address_file, + write_address_file, ) + __all__ = [ "get_local_ip", "build_worker_bind_address", diff --git a/src/xorl/server/utils/network.py b/src/xorl/server/utils/network.py index 314251c4..927562d2 100644 --- a/src/xorl/server/utils/network.py +++ b/src/xorl/server/utils/network.py @@ -10,6 +10,7 @@ import socket from typing import Optional + logger = logging.getLogger(__name__) @@ -148,7 +149,6 @@ def write_address_file(address: str, output_dir: str, filename: str = ".rank0_ad >>> write_address_file("tcp://10.0.0.5:5556", "/shared/outputs") '/shared/outputs/.rank0_address' """ - import os from pathlib import Path # Ensure output directory exists diff --git a/src/xorl/server/utils/storage.py b/src/xorl/server/utils/storage.py index 05bc63d9..b301bbc0 100644 --- a/src/xorl/server/utils/storage.py +++ b/src/xorl/server/utils/storage.py @@ -12,6 +12,7 @@ import re from typing import Optional + logger = logging.getLogger(__name__) @@ -78,39 +79,36 @@ def parse_size_string(size_str: str) -> int: size_str = size_str.strip().upper() # Pattern: number followed by optional unit - pattern = r'^(\d+(?:\.\d+)?)\s*(GB?|MB?|KB?|TB?|B)?$' + pattern = r"^(\d+(?:\.\d+)?)\s*(GB?|MB?|KB?|TB?|B)?$" match = re.match(pattern, size_str) if not match: - raise ValueError( - f"Invalid size format: '{size_str}'. " - "Expected format like '1GB', '500MB', '100KB', or '1024B'" - ) + raise ValueError(f"Invalid size format: '{size_str}'. Expected format like '1GB', '500MB', '100KB', or '1024B'") value = float(match.group(1)) - unit = match.group(2) or 'B' + unit = match.group(2) or "B" # Normalize single-letter units unit_map = { - 'T': 'TB', - 'G': 'GB', - 'M': 'MB', - 'K': 'KB', - 'B': 'B', - 'TB': 'TB', - 'GB': 'GB', - 'MB': 'MB', - 'KB': 'KB', + "T": "TB", + "G": "GB", + "M": "MB", + "K": "KB", + "B": "B", + "TB": "TB", + "GB": "GB", + "MB": "MB", + "KB": "KB", } unit = unit_map.get(unit, unit) multipliers = { - 'B': 1, - 'KB': 1024, - 'MB': 1024 ** 2, - 'GB': 1024 ** 3, - 'TB': 1024 ** 4, + "B": 1, + "KB": 1024, + "MB": 1024**2, + "GB": 1024**3, + "TB": 1024**4, } if unit not in multipliers: @@ -138,7 +136,7 @@ def bytes_to_human(size_bytes: int) -> str: if size_bytes < 0: return "0B" - units = ['B', 'KB', 'MB', 'GB', 'TB'] + units = ["B", "KB", "MB", "GB", "TB"] unit_idx = 0 size = float(size_bytes) diff --git a/src/xorl/server/utils/zmq_channels.py b/src/xorl/server/utils/zmq_channels.py index 379ae89a..349d92f6 100644 --- a/src/xorl/server/utils/zmq_channels.py +++ b/src/xorl/server/utils/zmq_channels.py @@ -18,6 +18,7 @@ import zmq import zmq.asyncio + logger = logging.getLogger(__name__) diff --git a/src/xorl/server/weight_sync/__init__.py b/src/xorl/server/weight_sync/__init__.py index 4f4607bb..2419adc0 100644 --- a/src/xorl/server/weight_sync/__init__.py +++ b/src/xorl/server/weight_sync/__init__.py @@ -2,4 +2,5 @@ from .backends.nccl_broadcast import NCCLBroadcastBackend + __all__ = ["NCCLBroadcastBackend"] diff --git a/src/xorl/server/weight_sync/backends/__init__.py b/src/xorl/server/weight_sync/backends/__init__.py index 4d4fc89d..a94addf1 100644 --- a/src/xorl/server/weight_sync/backends/__init__.py +++ b/src/xorl/server/weight_sync/backends/__init__.py @@ -2,6 +2,7 @@ from .base import TransportConfig, WeightTransportBackend + __all__ = [ "TransportConfig", "WeightTransportBackend", @@ -26,8 +27,6 @@ def create_backend( """ if method == "nccl_broadcast": from .nccl_broadcast import NCCLBroadcastBackend + return NCCLBroadcastBackend(config, **kwargs) - raise ValueError( - f"Unknown weight sync backend: {method!r}. " - f"Supported: 'nccl_broadcast'." - ) + raise ValueError(f"Unknown weight sync backend: {method!r}. Supported: 'nccl_broadcast'.") diff --git a/src/xorl/server/weight_sync/backends/base.py b/src/xorl/server/weight_sync/backends/base.py index 7845797b..acddba7e 100644 --- a/src/xorl/server/weight_sync/backends/base.py +++ b/src/xorl/server/weight_sync/backends/base.py @@ -40,7 +40,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import Any, Dict, FrozenSet, List, Optional, Tuple +from typing import Any, Dict, FrozenSet, List, Tuple import torch @@ -48,6 +48,7 @@ @dataclass class EndpointConfig: """Description of a single inference endpoint.""" + host: str port: int world_size: int = 1 # TP size on the inference side @@ -60,6 +61,7 @@ class TransportConfig: The handler populates this from the ``SyncWeightsData`` payload, and the backend reads whichever fields it needs. """ + endpoints: List[EndpointConfig] = field(default_factory=list) master_address: str = "localhost" master_port: int = 29600 diff --git a/src/xorl/server/weight_sync/backends/nccl_broadcast.py b/src/xorl/server/weight_sync/backends/nccl_broadcast.py index 6d90c828..613315e3 100644 --- a/src/xorl/server/weight_sync/backends/nccl_broadcast.py +++ b/src/xorl/server/weight_sync/backends/nccl_broadcast.py @@ -23,6 +23,7 @@ import torch import torch.distributed as dist + logger = logging.getLogger(__name__) # Reusable session for HTTP connection pooling @@ -48,6 +49,7 @@ def _get_http_session() -> requests.Session: @dataclass class EndpointInfo: """Information about an inference endpoint.""" + host: str port: int world_size: int # tensor_parallel_size for this endpoint @@ -56,6 +58,7 @@ class EndpointInfo: @dataclass class SyncResult: """Result of weight synchronization.""" + success: bool message: str transfer_time: float = 0.0 @@ -132,7 +135,7 @@ def _init_training_process_group(self) -> dist.ProcessGroup: Returns: The initialized process group """ - from torch.distributed import TCPStore, PrefixStore + from torch.distributed import PrefixStore, TCPStore from torch.distributed.distributed_c10d import ( Backend, _new_process_group_helper, @@ -146,7 +149,7 @@ def _init_training_process_group(self) -> dist.ProcessGroup: # torchrun sets TORCHELASTIC_USE_AGENT_STORE=True which forces all ranks # to be TCPStore clients. We need to unset this for our separate weight # sync process group where rank 0 must be the store master. - old_agent_store = os.environ.pop('TORCHELASTIC_USE_AGENT_STORE', None) + old_agent_store = os.environ.pop("TORCHELASTIC_USE_AGENT_STORE", None) rank = 0 # Training is always rank 0 @@ -165,7 +168,7 @@ def _init_training_process_group(self) -> dist.ProcessGroup: timeout = default_pg_timeout # Create TCPStore directly - rank 0 is the master - is_master = (rank == 0) + is_master = rank == 0 logger.info(f"[Training] Creating TCPStore (is_master={is_master})...") store = TCPStore( host_name=self.master_address, @@ -177,10 +180,11 @@ def _init_training_process_group(self) -> dist.ProcessGroup: # Use PrefixStore with group_name to namespace keys store = PrefixStore(self.group_name, store) - logger.info(f"[Training] TCPStore created, creating process group...") + logger.info("[Training] TCPStore created, creating process group...") # Handle different PyTorch versions by inspecting the actual signature import inspect + _pg_params = inspect.signature(_new_process_group_helper).parameters if "backend_options" in _pg_params: pg_options_param_name = "backend_options" @@ -206,7 +210,7 @@ def _init_training_process_group(self) -> dist.ProcessGroup: finally: # Restore the environment variable if old_agent_store is not None: - os.environ['TORCHELASTIC_USE_AGENT_STORE'] = old_agent_store + os.environ["TORCHELASTIC_USE_AGENT_STORE"] = old_agent_store _world.pg_group_ranks[pg] = {i: i for i in range(self.world_size)} @@ -453,10 +457,7 @@ def _endpoint_request_with_retry( } if success: - logger.info( - f"[{endpoint_label}] {operation} succeeded " - f"(attempt {attempt + 1}/{max_retries})" - ) + logger.info(f"[{endpoint_label}] {operation} succeeded (attempt {attempt + 1}/{max_retries})") return endpoint_result else: logger.warning( @@ -471,10 +472,7 @@ def _endpoint_request_with_retry( "message": str(e), "attempts": attempt + 1, } - logger.warning( - f"[{endpoint_label}] {operation} error " - f"(attempt {attempt + 1}/{max_retries}): {e}" - ) + logger.warning(f"[{endpoint_label}] {operation} error (attempt {attempt + 1}/{max_retries}): {e}") if attempt < max_retries - 1: time.sleep(retry_delay_seconds) @@ -497,9 +495,13 @@ def pause_inference_endpoints( futures = { executor.submit( self._endpoint_request_with_retry, - ep, "/pause_generation", "Pause", - {"mode": pause_mode}, timeout=60, - max_retries=max_retries, retry_delay_seconds=retry_delay_seconds, + ep, + "/pause_generation", + "Pause", + {"mode": pause_mode}, + timeout=60, + max_retries=max_retries, + retry_delay_seconds=retry_delay_seconds, ): ep for ep in self.endpoints } @@ -525,9 +527,13 @@ def resume_inference_endpoints( futures = { executor.submit( self._endpoint_request_with_retry, - ep, "/continue_generation", "Resume", - {}, timeout=30, - max_retries=max_retries, retry_delay_seconds=retry_delay_seconds, + ep, + "/continue_generation", + "Resume", + {}, + timeout=30, + max_retries=max_retries, + retry_delay_seconds=retry_delay_seconds, ): ep for ep in self.endpoints } @@ -595,15 +601,15 @@ def call_single_endpoint(endpoint: EndpointInfo, endpoint_idx: int): timeout=600, ) result = response.json() - update_results.append({ - "endpoint": f"{endpoint.host}:{endpoint.port}", - "success": result.get("success", False), - "message": result.get("message", ""), - }) + update_results.append( + { + "endpoint": f"{endpoint.host}:{endpoint.port}", + "success": result.get("success", False), + "message": result.get("message", ""), + } + ) if not result.get("success"): - update_errors.append( - f"API failed on {endpoint.host}:{endpoint.port}: {result}" - ) + update_errors.append(f"API failed on {endpoint.host}:{endpoint.port}: {result}") except Exception as e: update_errors.append(f"Exception calling {endpoint.host}:{endpoint.port}: {e}") @@ -696,10 +702,7 @@ def sync_weights( logger.info("[Training] Skipping lm_head.weight (tied with embed_tokens)") total_params = sum(state_dict[n].numel() for n in param_names) - total_bytes = sum( - state_dict[n].numel() * state_dict[n].element_size() - for n in param_names - ) + total_bytes = sum(state_dict[n].numel() * state_dict[n].element_size() for n in param_names) logger.info(f"[Training] Transferring {len(param_names)} parameters") logger.info(f"[Training] Total parameters: {total_params:,}") @@ -738,11 +741,11 @@ def sync_weights( for i, bucket in enumerate(buckets): bucket_size = sum(p.numel() * p.element_size() for _, p in bucket) - is_last_bucket = (i == len(buckets) - 1) + is_last_bucket = i == len(buckets) - 1 logger.info( - f"[Training] Bucket {i+1}/{len(buckets)} " - f"({len(bucket)} params, {bucket_size/1e6:.1f} MB)" + f"[Training] Bucket {i + 1}/{len(buckets)} " + f"({len(bucket)} params, {bucket_size / 1e6:.1f} MB)" f"{' [final, flush_cache=True]' if is_last_bucket else ''}" ) @@ -775,10 +778,7 @@ def sync_weights( total_bytes=total_bytes, num_parameters=len(param_names), num_buckets=len(buckets), - endpoint_results=[ - {"host": ep.host, "port": ep.port, "success": True} - for ep in self.endpoints - ], + endpoint_results=[{"host": ep.host, "port": ep.port, "success": True} for ep in self.endpoints], ) except Exception as e: @@ -823,10 +823,7 @@ def __init__(self, config: TransportConfig, **kwargs) -> None: def initialize(self) -> bool: cfg = self.config - ep_infos = [ - EndpointInfo(host=e.host, port=e.port, world_size=e.world_size) - for e in cfg.endpoints - ] + ep_infos = [EndpointInfo(host=e.host, port=e.port, world_size=e.world_size) for e in cfg.endpoints] self._synchronizer = NCCLWeightSynchronizer( endpoints=ep_infos, master_address=cfg.master_address, @@ -835,10 +832,7 @@ def initialize(self) -> bool: buffer_size_mb=cfg.buffer_size_mb, device=cfg.device, ) - logger.info( - f"[NCCLBroadcast] Initializing NCCL sync group " - f"({len(ep_infos)} endpoints, device={cfg.device})" - ) + logger.info(f"[NCCLBroadcast] Initializing NCCL sync group ({len(ep_infos)} endpoints, device={cfg.device})") ok = self._synchronizer.init_nccl_group() if ok: self._process_group = self._synchronizer.process_group @@ -861,9 +855,7 @@ def transfer_bucket( flush_cache: bool = False, ) -> None: if src_rank != 0: - raise ValueError( - f"NCCLBroadcastBackend only supports src_rank=0, got {src_rank}" - ) + raise ValueError(f"NCCLBroadcastBackend only supports src_rank=0, got {src_rank}") if self._synchronizer is None: raise RuntimeError("Backend not initialized — call initialize() first") self._synchronizer._transfer_single_bucket(bucket, flush_cache=flush_cache) diff --git a/src/xorl/server/weight_sync/endpoint_manager.py b/src/xorl/server/weight_sync/endpoint_manager.py index 3918b67d..03b31f6b 100644 --- a/src/xorl/server/weight_sync/endpoint_manager.py +++ b/src/xorl/server/weight_sync/endpoint_manager.py @@ -11,6 +11,7 @@ import requests + logger = logging.getLogger(__name__) # Reusable session for HTTP connection pooling @@ -52,10 +53,7 @@ def health_check(self) -> None: resp.raise_for_status() logger.info(f"[EndpointMgr] {ep['host']}:{ep['port']} healthy") except Exception as e: - raise RuntimeError( - f"Inference endpoint {ep['host']}:{ep['port']} " - f"health check failed: {e}" - ) + raise RuntimeError(f"Inference endpoint {ep['host']}:{ep['port']} health check failed: {e}") def pause( self, @@ -113,8 +111,13 @@ def _parallel_request( futures = { pool.submit( self._request_with_retry, - ep, url_path, operation, payload, - timeout, max_retries, retry_delay, + ep, + url_path, + operation, + payload, + timeout, + max_retries, + retry_delay, ): ep for ep in self.endpoints } @@ -149,10 +152,7 @@ def _request_with_retry( "attempts": attempt + 1, } if success: - logger.info( - f"[EndpointMgr] {label} {operation} ok " - f"(attempt {attempt + 1}/{max_retries})" - ) + logger.info(f"[EndpointMgr] {label} {operation} ok (attempt {attempt + 1}/{max_retries})") return result logger.warning( f"[EndpointMgr] {label} {operation} failed " @@ -165,10 +165,7 @@ def _request_with_retry( "message": str(e), "attempts": attempt + 1, } - logger.warning( - f"[EndpointMgr] {label} {operation} error " - f"(attempt {attempt + 1}/{max_retries}): {e}" - ) + logger.warning(f"[EndpointMgr] {label} {operation} error (attempt {attempt + 1}/{max_retries}): {e}") if attempt < max_retries - 1: time.sleep(retry_delay) diff --git a/src/xorl/server/weight_sync/handler.py b/src/xorl/server/weight_sync/handler.py index 732fb45f..6990a9e8 100644 --- a/src/xorl/server/weight_sync/handler.py +++ b/src/xorl/server/weight_sync/handler.py @@ -33,6 +33,7 @@ import torch import torch.distributed as dist + logger = logging.getLogger(__name__) # Default bucket size for MoE expert broadcasting (256 MB) @@ -71,6 +72,7 @@ async def handle_sync_inference_weights(self, command_dict: Dict[str, Any]) -> D logger.info(f"Rank {self.rank}: [WeightSync] Starting sync_inference_weights") from xorl.server.protocol.operations import SyncWeightsData + p: SyncWeightsData = command_dict.get("payload", SyncWeightsData()) endpoints = p.endpoints @@ -167,14 +169,15 @@ def _sync_weights( # ------------------------------------------------------------------ # Step 3: Create backend + endpoint manager, initialize on sender ranks # ------------------------------------------------------------------ - from xorl.server.weight_sync.backends import create_backend, TransportConfig + from xorl.server.weight_sync.backends import TransportConfig, create_backend from xorl.server.weight_sync.backends.base import EndpointConfig from xorl.server.weight_sync.endpoint_manager import EndpointManager transport_cfg = TransportConfig( endpoints=[ EndpointConfig( - host=ep["host"], port=ep["port"], + host=ep["host"], + port=ep["port"], world_size=ep.get("world_size", 1), ) for ep in endpoints @@ -239,6 +242,7 @@ def _sync_weights( # Detect EP mode from xorl.distributed.parallel_state import get_parallel_state + _ps = get_parallel_state() _ep_enabled = _ps.ep_enabled and _ps.ep_size > 1 @@ -254,7 +258,7 @@ def _sync_weights( try: for pp_stage in range(_pp_size): - _is_my_stage = (_pp_rank == pp_stage) + _is_my_stage = _pp_rank == pp_stage _is_remote = _pp_enabled and pp_stage > 0 # Stage leader: the dp_shard_rank==0 rank within this stage _stage_leader = _ps.dp_shard_rank == 0 if _pp_enabled else (self.rank == 0) @@ -279,10 +283,7 @@ def _sync_weights( ) for mod_idx in range(num_stage_modules): - is_last_overall = ( - mod_idx == num_stage_modules - 1 - and pp_stage == _pp_size - 1 - ) + is_last_overall = mod_idx == num_stage_modules - 1 and pp_stage == _pp_size - 1 # ── FSDP ops (only ranks owning this stage) ────────── current_buffer = None @@ -295,12 +296,16 @@ def _sync_weights( fsdp_mod.unshard() qlora_linear_buffer, moe_contexts = self._qlora_collective_ops( - fsdp_mod, mod_name, collect_results=_stage_leader, + fsdp_mod, + mod_name, + collect_results=_stage_leader, ) if _ep_enabled: ep_moe_contexts = self._collect_ep_moe_data( - fsdp_mod, mod_name, _ps, + fsdp_mod, + mod_name, + _ps, ) # EP MoE prefixes to skip in extraction @@ -311,19 +316,21 @@ def _sync_weights( if p == mod_name: ep_moe_prefixes.add("") elif p.startswith(mod_name + "."): - ep_moe_prefixes.add(p[len(mod_name) + 1:]) + ep_moe_prefixes.add(p[len(mod_name) + 1 :]) else: ep_moe_prefixes.add(p) if _stage_leader: if ep_moe_prefixes: logger.info( - f"Rank {self.rank}: [WeightSync] ep_moe_prefixes=" - f"{ep_moe_prefixes} for {mod_name}" + f"Rank {self.rank}: [WeightSync] ep_moe_prefixes={ep_moe_prefixes} for {mod_name}" ) from torch.distributed._tensor import DTensor + current_buffer = self._extract_params_for_sync( - fsdp_mod, mod_name, DTensor, + fsdp_mod, + mod_name, + DTensor, skip_moe_prefixes=ep_moe_prefixes, ) current_buffer.extend(qlora_linear_buffer) @@ -336,23 +343,19 @@ def _sync_weights( # Stage 0: sender rank(s) broadcast directly to SGLang if _is_sender and current_buffer: current_buffer = self._unfuse_for_inference( - current_buffer, model, + current_buffer, + model, ) if quantization and quantization.get("quant_method") == "fp8": current_buffer = self._quantize_buffer_for_fp8( current_buffer, quantization_config=quantization, ) - logger.info( - f"Rank 0: [WeightSync] Module {mod_name}: " - f"{len(current_buffer)} params" - ) + logger.info(f"Rank 0: [WeightSync] Module {mod_name}: {len(current_buffer)} params") b, p = self._broadcast_buffer( - backend, current_buffer, - flush_cache=( - flush_cache and is_last_overall - and not moe_contexts - ), + backend, + current_buffer, + flush_cache=(flush_cache and is_last_overall and not moe_contexts), ) total_bytes += b total_params += p @@ -362,9 +365,10 @@ def _sync_weights( # Stage 0 MoE handling (unchanged) if moe_contexts or ep_moe_contexts: if _ep_enabled: - for ctx in (moe_contexts + ep_moe_contexts): + for ctx in moe_contexts + ep_moe_contexts: b, p, n = self._gather_and_broadcast_ep_moe_experts( - backend, ctx, + backend, + ctx, flush_cache=(flush_cache and is_last_overall), quantization=quantization, ps=_ps, @@ -375,7 +379,8 @@ def _sync_weights( elif _is_sender: for ctx in moe_contexts: b, p, n = self._broadcast_moe_experts_bucketed( - backend, ctx, + backend, + ctx, flush_cache=(flush_cache and is_last_overall), quantization=quantization, ) @@ -398,7 +403,9 @@ def _sync_weights( # NCCL transfer: metadata + flat bf16 tensor received = self._pp_nccl_transfer_buffer( send_buf if (_is_my_stage and _stage_leader) else None, - pp_grp, stage_src, device, + pp_grp, + stage_src, + device, ) # Rank 0: quantize + broadcast to SGLang @@ -413,7 +420,8 @@ def _sync_weights( f"{mod_idx}: {len(received)} params via NCCL" ) b, p = self._broadcast_buffer( - backend, received, + backend, + received, flush_cache=(flush_cache and is_last_overall), ) total_bytes += b @@ -449,13 +457,10 @@ def _sync_weights( "total_bytes": total_bytes, "num_parameters": total_params, "num_buckets": num_buckets, - "endpoint_results": [ - {"host": ep["host"], "port": ep["port"], "success": True} - for ep in endpoints - ], + "endpoint_results": [{"host": ep["host"], "port": ep["port"], "success": True} for ep in endpoints], } - except Exception as e: + except Exception: if endpoint_mgr is not None: try: endpoint_mgr.resume() @@ -465,7 +470,9 @@ def _sync_weights( try: backend.destroy() except Exception as destroy_err: - logger.warning(f"Rank {self.rank}: [WeightSync] Failed to destroy backend during cleanup: {destroy_err}") + logger.warning( + f"Rank {self.rank}: [WeightSync] Failed to destroy backend during cleanup: {destroy_err}" + ) raise # ======================================================================== @@ -504,13 +511,14 @@ def _qlora_collective_ops( moe_contexts: list of dicts with info for per-expert processing """ if collect_results is None: - collect_results = (self.rank == 0) - from xorl.qlora.modules.linear import QLoRALinear - from xorl.qlora.modules.moe_experts import QLoRAMoeExperts - + collect_results = self.rank == 0 # Discover QLoRA modules (only those directly owned by this FSDP module, # not by child FSDP modules — child modules will be processed separately) from torch.distributed.fsdp._fully_shard import FSDPModule + + from xorl.qlora.modules.linear import QLoRALinear + from xorl.qlora.modules.moe_experts import QLoRAMoeExperts + child_fsdp_prefixes = set() for mname, mod in fsdp_mod.named_modules(): if isinstance(mod, FSDPModule) and mname != "": @@ -550,6 +558,7 @@ def _qlora_collective_ops( # --- Phase 2: QLoRAMoeExperts (gather lora params, defer heavy work) --- from xorl.distributed.parallel_state import get_parallel_state + ps = get_parallel_state() ep_enabled = ps.ep_enabled and ps.ep_size > 1 @@ -567,13 +576,13 @@ def _qlora_collective_ops( param = getattr(mod, key) if ep_enabled: # EP: use local shard directly (no collective needed) - if hasattr(param, 'to_local'): + if hasattr(param, "to_local"): lora_params[key] = param.to_local().clone() else: lora_params[key] = param.data.clone() else: # Non-EP: collective full_tensor to get global params on stage leader - if hasattr(param, 'full_tensor'): + if hasattr(param, "full_tensor"): full = param.full_tensor() # collective: all ranks else: full = param.data @@ -598,7 +607,10 @@ def _qlora_collective_ops( # ======================================================================== def _collect_ep_moe_data( - self, fsdp_mod, mod_name: str, ps, + self, + fsdp_mod, + mod_name: str, + ps, ) -> List[Dict[str, Any]]: """Collect local EP-sharded MoE expert data during unshard phase. @@ -611,13 +623,15 @@ def _collect_ep_moe_data( Must be called between unshard() and reshard(). Returns list of context dicts for _gather_and_broadcast_ep_moe_experts. """ - from xorl.models.layers.moe.experts import MoEExperts - from xorl.models.layers.moe.lora import MoEExpertsLoRA - from xorl.qlora.modules.moe_experts import QLoRAMoeExperts from torch.distributed._tensor import DTensor # Skip modules owned by child FSDP modules (processed separately) from torch.distributed.fsdp._fully_shard import FSDPModule + + from xorl.models.layers.moe.experts import MoEExperts + from xorl.models.layers.moe.lora import MoEExpertsLoRA + from xorl.qlora.modules.moe_experts import QLoRAMoeExperts + child_fsdp_prefixes = set() for mname, mod in fsdp_mod.named_modules(): if isinstance(mod, FSDPModule) and mname != "": @@ -634,7 +648,7 @@ def _collect_ep_moe_data( continue # Get expert params — after unshard they may be plain tensors or DTensors - gate = getattr(mod, 'gate_proj', None) + gate = getattr(mod, "gate_proj", None) if gate is None or not isinstance(gate, torch.nn.Parameter): continue @@ -662,7 +676,7 @@ def _collect_ep_moe_data( if isinstance(delta, DTensor): delta = delta.to_local() elif delta.shape[0] > E_local: - delta = delta[start:start + E_local] + delta = delta[start : start + E_local] local = local.to(torch.bfloat16) + delta.to(torch.bfloat16) else: local = local.to(torch.bfloat16) @@ -671,12 +685,14 @@ def _collect_ep_moe_data( local_experts[proj_name] = local.clone() - contexts.append({ - "type": "full_weight", - "prefix": full_prefix, - "local_experts": local_experts, # {proj: [E_local, K, N]} - "num_local_experts": E_local, - }) + contexts.append( + { + "type": "full_weight", + "prefix": full_prefix, + "local_experts": local_experts, # {proj: [E_local, K, N]} + "num_local_experts": E_local, + } + ) return contexts @@ -798,10 +814,13 @@ def _gather_and_broadcast_ep_moe_experts( if bucket_bytes >= bucket_size_bytes: if quantization and quantization.get("quant_method") == "fp8": bucket = self._quantize_buffer_for_fp8( - bucket, quantization_config=quantization, + bucket, + quantization_config=quantization, ) b, p = self._broadcast_buffer( - backend, bucket, flush_cache=False, + backend, + bucket, + flush_cache=False, ) total_bytes += b total_params += p @@ -815,10 +834,13 @@ def _gather_and_broadcast_ep_moe_experts( if self.rank == 0 and bucket: if quantization and quantization.get("quant_method") == "fp8": bucket = self._quantize_buffer_for_fp8( - bucket, quantization_config=quantization, + bucket, + quantization_config=quantization, ) b, p = self._broadcast_buffer( - backend, bucket, flush_cache=flush_cache, + backend, + bucket, + flush_cache=flush_cache, ) total_bytes += b total_params += p @@ -910,10 +932,13 @@ def _broadcast_moe_experts_bucketed( if bucket_bytes >= bucket_size_bytes: if quantization and quantization.get("quant_method") == "fp8": bucket = self._quantize_buffer_for_fp8( - bucket, quantization_config=quantization, + bucket, + quantization_config=quantization, ) b, p = self._broadcast_buffer( - backend, bucket, flush_cache=False, + backend, + bucket, + flush_cache=False, ) total_bytes += b total_params += p @@ -925,10 +950,13 @@ def _broadcast_moe_experts_bucketed( if bucket: if quantization and quantization.get("quant_method") == "fp8": bucket = self._quantize_buffer_for_fp8( - bucket, quantization_config=quantization, + bucket, + quantization_config=quantization, ) b, p = self._broadcast_buffer( - backend, bucket, flush_cache=flush_cache, + backend, + bucket, + flush_cache=flush_cache, ) total_bytes += b total_params += p @@ -966,7 +994,7 @@ def _pp_nccl_transfer_buffer( For non-source ranks: list of (name, tensor) on device For source rank: None """ - is_src = (self.rank == src_global) + is_src = self.rank == src_global # Step 1: broadcast metadata (names + shapes) if is_src: @@ -997,7 +1025,7 @@ def _pp_nccl_transfer_buffer( offset = 0 for name, shape in meta: numel = _prod(shape) - result.append((name, flat[offset:offset + numel].view(shape))) + result.append((name, flat[offset : offset + numel].view(shape))) offset += numel return result @@ -1042,9 +1070,7 @@ def _compute_moe_experts_buffer( hf_name = f"{full_prefix}.{expert_idx}.{hf_proj}.weight" items.append((hf_name, merged)) - logger.info( - f"Rank {self.rank}: [WeightSync] MoE {full_prefix}: {len(items)} experts" - ) + logger.info(f"Rank {self.rank}: [WeightSync] MoE {full_prefix}: {len(items)} experts") # Free cloned lora params del lora_params @@ -1102,8 +1128,7 @@ def _quantize_buffer_for_fp8( # "model.layers.0.mlp.gate" matches "model.layers.0.mlp.gate.weight" if modules_to_not_convert: skip = any( - name == prefix + ".weight" or name.startswith(prefix + ".") - for prefix in modules_to_not_convert + name == prefix + ".weight" or name.startswith(prefix + ".") for prefix in modules_to_not_convert ) if skip: result.append((name, tensor)) @@ -1121,8 +1146,10 @@ def _quantize_buffer_for_fp8( if pad_rows > 0 or pad_cols > 0: padded = torch.zeros( - rows + pad_rows, cols + pad_cols, - dtype=tensor.dtype, device=tensor.device, + rows + pad_rows, + cols + pad_cols, + dtype=tensor.dtype, + device=tensor.device, ) padded[:rows, :cols] = tensor else: @@ -1146,9 +1173,7 @@ def _quantize_buffer_for_fp8( quantized_blocks = quantized_blocks.to(fp8_dtype) # Reshape back: [nr, nc, block_size, block_size] → [padded_rows, padded_cols] - quantized = quantized_blocks.permute(0, 2, 1, 3).reshape( - padded.shape[0], padded.shape[1] - ) + quantized = quantized_blocks.permute(0, 2, 1, 3).reshape(padded.shape[0], padded.shape[1]) # Remove padding if pad_rows > 0 or pad_cols > 0: @@ -1185,8 +1210,8 @@ def _extract_params_for_sync( """ from xorl.lora.modules.base import LoraModule from xorl.lora.modules.linear import LoraLinear - from xorl.models.layers.moe.lora import MoEExpertsLoRA from xorl.models.layers.moe.experts import MoEExperts + from xorl.models.layers.moe.lora import MoEExpertsLoRA from xorl.qlora.modules.linear import QLoRALinear from xorl.qlora.modules.moe_experts import QLoRAMoeExperts @@ -1196,6 +1221,7 @@ def _extract_params_for_sync( # when that child module is unsharded separately. Parent unshard may # expose child params as plain tensors, so we can't rely on DTensor check. from torch.distributed.fsdp._fully_shard import FSDPModule + child_fsdp_prefixes = set() for mname, mod in fsdp_mod.named_modules(): if isinstance(mod, FSDPModule) and mname != "": @@ -1319,8 +1345,7 @@ def _extract_params_for_sync( for buf_name, buf_tensor in buffer: if buf_name == full_source: logger.info( - f"Rank 0: [WeightSync] Tied weight: emitting " - f"{full_tied} (clone of {full_source})" + f"Rank 0: [WeightSync] Tied weight: emitting {full_tied} (clone of {full_source})" ) buffer.append((full_tied, buf_tensor.clone())) break @@ -1354,8 +1379,8 @@ def _unfuse_for_inference( # Split [q_size + 2*kv_size, hidden] → q, k, v prefix, suffix = name.rsplit(".qkv_proj.", 1) q = tensor[:q_size].clone() - k = tensor[q_size:q_size + kv_size].clone() - v = tensor[q_size + kv_size:].clone() + k = tensor[q_size : q_size + kv_size].clone() + v = tensor[q_size + kv_size :].clone() result.append((f"{prefix}.q_proj.{suffix}", q)) result.append((f"{prefix}.k_proj.{suffix}", k)) result.append((f"{prefix}.v_proj.{suffix}", v)) @@ -1390,10 +1415,7 @@ def _broadcast_buffer( return 0, 0 bucket_bytes = sum(t.numel() * t.element_size() for _, t in buffer) - logger.info( - f"Rank {self.rank}: [WeightSync] Broadcasting {len(buffer)} params, " - f"{bucket_bytes / 1e6:.1f} MB" - ) + logger.info(f"Rank {self.rank}: [WeightSync] Broadcasting {len(buffer)} params, {bucket_bytes / 1e6:.1f} MB") backend.transfer_bucket(buffer, flush_cache=flush_cache) return bucket_bytes, len(buffer) @@ -1423,4 +1445,3 @@ def _get_fsdp_modules(model) -> Tuple[Optional[Any], List[Tuple[str, Any]]]: layer_modules.append((fqn, mod)) return root_module, layer_modules - diff --git a/src/xorl/trainers/__init__.py b/src/xorl/trainers/__init__.py index bc0ddfaf..10639c3f 100644 --- a/src/xorl/trainers/__init__.py +++ b/src/xorl/trainers/__init__.py @@ -1,3 +1,4 @@ from xorl.trainers.trainer import Trainer, TrainState + __all__ = ["Trainer", "TrainState"] diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index 36a994fa..e5395c25 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -13,6 +13,7 @@ from xorl.utils import helper + logger = helper.create_logger(__name__) @@ -126,6 +127,7 @@ def build_training_model( # Set module-level flags for rope and activation if rope_native: from xorl.models.layers.rope import set_rope_native + set_rope_native(True) logger.info_rank0("Using native RoPE (flash_attn fused kernel disabled)") if activation_native: @@ -237,6 +239,7 @@ def build_training_model( helper.print_device_mem_info("VRAM usage after QLoRA quantization") elif enable_lora: from xorl.lora import freeze_base_parameters + freeze_base_parameters(model) logger.info_rank0("Base model parameters frozen, only LoRA parameters trainable") else: @@ -279,6 +282,7 @@ def build_training_model( # Internal helpers # ====================================================================== + def _inject_qlora( model: nn.Module, *, @@ -296,9 +300,9 @@ def _inject_qlora( Returns (is_prequantized, checkpoint_quant_format, exclude_modules). """ from xorl.qlora import ( - inject_qlora_into_model, - detect_prequantized_nvfp4, detect_prequantized_block_fp8, + detect_prequantized_nvfp4, + inject_qlora_into_model, ) is_prequantized = False @@ -319,11 +323,11 @@ def _inject_qlora( logger.info_rank0(f"Using user-specified exclude_modules: {exclude_modules}") elif is_prequantized: from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules + exclude_modules = get_prequantized_exclude_modules(weights_path) if exclude_modules: logger.info_rank0( - f"Auto-detected {len(exclude_modules)} excluded modules " - f"from checkpoint config: {exclude_modules}" + f"Auto-detected {len(exclude_modules)} excluded modules from checkpoint config: {exclude_modules}" ) # NF4 quantizes bf16 weights on-the-fly — no pre-quantized checkpoint needed. @@ -383,9 +387,9 @@ def _inject_lora( if is_moe_model and (moe_shared_lora or moe_hybrid_shared_lora): from xorl.lora.utils import inject_lora_into_model_with_moe + logger.info_rank0( - f"MoE-aware LoRA injection " - f"(shared={moe_shared_lora}, hybrid_shared={moe_hybrid_shared_lora})" + f"MoE-aware LoRA injection (shared={moe_shared_lora}, hybrid_shared={moe_hybrid_shared_lora})" ) inject_lora_into_model_with_moe( model, @@ -397,6 +401,7 @@ def _inject_lora( ) else: from xorl.lora.utils import inject_lora_into_model + inject_lora_into_model( model, r=lora_rank, @@ -425,47 +430,47 @@ def _deferred_qlora_quantize( # 1. Pre-quantized linear/MoE loading (nvfp4/block_fp8) needs_prequant_linear = any( - isinstance(m, QLoRALinear) and m._is_prequantized and not m._inline_loaded - for m in model.modules() + isinstance(m, QLoRALinear) and m._is_prequantized and not m._inline_loaded for m in model.modules() ) needs_prequant_moe = any( - isinstance(m, QLoRAMoeExperts) and not m._weights_loaded - and m._source_quant_format is not None + isinstance(m, QLoRAMoeExperts) and not m._weights_loaded and m._source_quant_format is not None for m in model.modules() ) if needs_prequant_linear or needs_prequant_moe: from xorl.qlora import maybe_load_prequantized_qlora + logger.info(f"Starting pre-quantized weight loading (mode={load_weights_mode}) …") helper.print_device_mem_info("VRAM before pre-quantized loading") maybe_load_prequantized_qlora(model, weights_path, load_mode=load_weights_mode) logger.info("Done pre-quantized weight loading") # 2. NF4 linear: FSDP loaded bf16 into weight param → quantize to NF4 - needs_bf16_quantize = any( - isinstance(m, QLoRALinear) and m.weight is not None - for m in model.modules() - ) + needs_bf16_quantize = any(isinstance(m, QLoRALinear) and m.weight is not None for m in model.modules()) if needs_bf16_quantize: from xorl.qlora import maybe_quantize_qlora + logger.info("Quantizing bf16 linear weights to NF4 …") maybe_quantize_qlora(model) # 3. NF4 MoE: load bf16 experts from checkpoint → quantize needs_bf16_moe = any( - isinstance(m, QLoRAMoeExperts) and not m._weights_loaded + isinstance(m, QLoRAMoeExperts) + and not m._weights_loaded and m._source_quant_format is None # NF4 (no source quant format) for m in model.modules() ) if needs_bf16_moe: from xorl.qlora import maybe_load_and_quantize_moe_qlora + logger.info("Loading and quantizing bf16 MoE expert weights …") maybe_load_and_quantize_moe_qlora( - model, weights_path, load_mode=load_weights_mode, + model, + weights_path, + load_mode=load_weights_mode, ) - if not (needs_prequant_linear or needs_prequant_moe - or needs_bf16_quantize or needs_bf16_moe): + if not (needs_prequant_linear or needs_prequant_moe or needs_bf16_quantize or needs_bf16_moe): logger.info("All QLoRA modules loaded inline, skipping deferred disk I/O") # Always deregister packed_weight_f32 from FSDP2 (prevent mixed-precision corruption) diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index 7beded2c..526dd777 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -147,6 +147,7 @@ def _bootstrap(self) -> None: save_args(args, args.train.output_dir) if args.train.use_wandb: import wandb + wandb.init( project=args.train.wandb_project, name=args.train.wandb_name, @@ -211,6 +212,7 @@ def _maybe_log_startup_metrics(self, metrics: Dict[str, Any], commit: bool = Fal if self.args.train.global_rank != 0 or not self.args.train.use_wandb or not self._wandb_initialized: return import wandb + wandb.log(metrics, step=0, commit=commit) def _write_startup_metrics_file(self) -> None: @@ -248,7 +250,8 @@ def _log_host_inventory(self) -> None: unique_hostnames = sorted({item["hostname"] for item in inventory}) rank_to_hostname = {str(item["global_rank"]): item["hostname"] for item in inventory} logger.info_rank0( - "Host inventory:\n" + json.dumps( + "Host inventory:\n" + + json.dumps( { "master_addr": os.environ.get("MASTER_ADDR"), "master_port": os.environ.get("MASTER_PORT"), @@ -271,6 +274,7 @@ def _log_host_inventory(self) -> None: self._maybe_log_startup_metrics({"startup/node_count": len(unique_hostnames)}, commit=False) if self.args.train.use_wandb and self._wandb_initialized: import wandb + wandb.config.update( { "master_addr": os.environ.get("MASTER_ADDR"), @@ -988,9 +992,7 @@ def _maybe_log( tokens_per_sec = train_metrics.get("efficiency/tokens_per_second(K)", 0) * 1e3 if use_tqdm and tqdm_bar is not None: - tqdm_bar.set_postfix_str( - f"loss={total_loss:.2f} gn={grad_norm:.2f} lr={lr:.1e} tok/s={tokens_per_sec:.0f}" - ) + tqdm_bar.set_postfix_str(f"loss={total_loss:.2f} gn={grad_norm:.2f} lr={lr:.1e} tok/s={tokens_per_sec:.0f}") tqdm_bar.update() else: max_steps_str = args.train.max_steps or "?" diff --git a/src/xorl/trainers/training_utils.py b/src/xorl/trainers/training_utils.py index 1fa41ba8..0453c050 100644 --- a/src/xorl/trainers/training_utils.py +++ b/src/xorl/trainers/training_utils.py @@ -143,18 +143,19 @@ def maybe_merge_lora( return if enable_qlora: from xorl.qlora.utils import maybe_requant_qlora + maybe_requant_qlora(model) elif enable_lora: from xorl.lora.utils import maybe_merge_lora as _merge + _merge(model) if reset_optimizer and optimizer is not None: count = reset_lora_optimizer_states(model, optimizer) if count > 0: import logging - logging.getLogger(__name__).info( - f"ReLoRA optimizer reset: pruned states for {count} LoRA parameters" - ) + + logging.getLogger(__name__).info(f"ReLoRA optimizer reset: pruned states for {count} LoRA parameters") def negotiate_pp_seq_len(micro_batches: List[Dict[str, Any]], pp_group) -> int: @@ -262,12 +263,8 @@ def forward_backward_pp( """ device = get_device_type() - input_ids = torch.cat( - [mb["input_ids"].to(device, non_blocking=True) for mb in micro_batches], dim=0 - ) - labels = torch.cat( - [mb["labels"].to(device, non_blocking=True) for mb in micro_batches], dim=0 - ) + input_ids = torch.cat([mb["input_ids"].to(device, non_blocking=True) for mb in micro_batches], dim=0) + labels = torch.cat([mb["labels"].to(device, non_blocking=True) for mb in micro_batches], dim=0) # Per-microbatch metadata for PP forward (position_ids, flash-attn kwargs) _PP_FA_KEYS = ("cu_seq_lens_q", "cu_seq_lens_k", "max_length_q", "max_length_k") diff --git a/src/xorl/utils/checkpoint_utils.py b/src/xorl/utils/checkpoint_utils.py index 5fc9b541..4c498fde 100644 --- a/src/xorl/utils/checkpoint_utils.py +++ b/src/xorl/utils/checkpoint_utils.py @@ -88,6 +88,6 @@ def get_checkpoint_path(output_dir, is_local_rank0: bool, ckpt_manager: str): return None checkpoint_path = os.path.join(output_dir, "checkpoints", f"global_step_{iteration}") - logger.info_rank0(f"Sucessfully get the latest checkpoint path: {checkpoint_path}") + logger.info_rank0(f"Successfully get the latest checkpoint path: {checkpoint_path}") return checkpoint_path diff --git a/src/xorl/utils/count_flops.py b/src/xorl/utils/count_flops.py index 7a621fca..71601390 100644 --- a/src/xorl/utils/count_flops.py +++ b/src/xorl/utils/count_flops.py @@ -49,7 +49,7 @@ def unit_convert(number, level): return number device_name = get_device_name() - flops = float("inf") # INF flops for unkown gpu type + flops = float("inf") # INF flops for unknown gpu type if "H100" in device_name or "H800" in device_name: flops = 989e12 elif "A100" in device_name or "A800" in device_name: @@ -143,24 +143,21 @@ def _estimate_deepseek_v3_flops(self, tokens_sum, batch_seqlens, delta_time): ) attn_linear_N += num_query_heads * self.config.v_head_dim * hidden_size - router_N = hidden_size * moe_num_expert - gate_up_N = hidden_size * moe_intermediate_size * (moe_topk + share_expert_num) * 2 - down_N = hidden_size * moe_intermediate_size * (moe_topk + share_expert_num) + router_N = hidden_size * moe_num_expert + gate_up_N = hidden_size * moe_intermediate_size * (moe_topk + share_expert_num) * 2 + down_N = hidden_size * moe_intermediate_size * (moe_topk + share_expert_num) dense_mlp_N = hidden_size * self.config.intermediate_size * 3 - embed_lm_N = vocab_size * hidden_size # lm_head only; embedding is a lookup (0 FLOPs) + embed_lm_N = vocab_size * hidden_size # lm_head only; embedding is a lookup (0 FLOPs) moe_layer_flops = ( - m["router"] * router_N * tokens_sum - + m["gate"] * gate_up_N * tokens_sum - + m["down"] * down_N * tokens_sum - + m["attn_linear"] * attn_linear_N * tokens_sum - ) - dense_layer_flops = ( - m["dense_mlp"] * dense_mlp_N * tokens_sum + m["router"] * router_N * tokens_sum + + m["gate"] * gate_up_N * tokens_sum + + m["down"] * down_N * tokens_sum + m["attn_linear"] * attn_linear_N * tokens_sum ) + dense_layer_flops = m["dense_mlp"] * dense_mlp_N * tokens_sum + m["attn_linear"] * attn_linear_N * tokens_sum dense_N_flops = ( - moe_layer_flops * (num_hidden_layers - first_k_dense_replace) + moe_layer_flops * (num_hidden_layers - first_k_dense_replace) + dense_layer_flops * first_k_dense_replace + 6 * embed_lm_N * tokens_sum ) @@ -188,17 +185,17 @@ def _estimate_qwen3_moe_flops(self, tokens_sum, batch_seqlens, delta_time): v_size = num_key_value_heads * head_dim # Per-parameter counts (number of multiplications in one forward pass) - router_N = hidden_size * moe_num_expert - gate_up_N = hidden_size * moe_intermediate_size * moe_topk * 2 # gate_proj + up_proj - down_N = hidden_size * moe_intermediate_size * moe_topk # down_proj + router_N = hidden_size * moe_num_expert + gate_up_N = hidden_size * moe_intermediate_size * moe_topk * 2 # gate_proj + up_proj + down_N = hidden_size * moe_intermediate_size * moe_topk # down_proj attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim) - embed_lm_N = vocab_size * hidden_size # lm_head only; embedding is a lookup (0 FLOPs) + embed_lm_N = vocab_size * hidden_size # lm_head only; embedding is a lookup (0 FLOPs) # Per-layer FLOPs with GC-corrected multipliers per_layer_flops = ( - m["router"] * router_N * tokens_sum - + m["gate"] * gate_up_N * tokens_sum - + m["down"] * down_N * tokens_sum + m["router"] * router_N * tokens_sum + + m["gate"] * gate_up_N * tokens_sum + + m["down"] * down_N * tokens_sum + m["attn_linear"] * attn_linear_N * tokens_sum ) dense_N_flops = per_layer_flops * num_hidden_layers + 6 * embed_lm_N * tokens_sum @@ -239,29 +236,27 @@ def _estimate_qwen3_5_flops(self, tokens_sum, batch_seqlens, delta_time): mlp_N = hidden_size * intermediate_size * 3 embed_lm_N = vocab_size * hidden_size # lm_head only; embedding is a lookup (0 FLOPs) - full_attn_layer_flops = ( - m["attn_linear"] * full_attn_linear_N * tokens_sum - + m["dense_mlp"] * mlp_N * tokens_sum - ) + full_attn_layer_flops = m["attn_linear"] * full_attn_linear_N * tokens_sum + m["dense_mlp"] * mlp_N * tokens_sum linear_layer_flops = 0 if linear_attn_layers > 0: - lin_key_dim = getattr(self.config, "linear_num_key_heads", num_attention_heads) * getattr(self.config, "linear_key_head_dim", 128) - lin_value_dim = getattr(self.config, "linear_num_value_heads", num_attention_heads) * getattr(self.config, "linear_value_head_dim", 128) + lin_key_dim = getattr(self.config, "linear_num_key_heads", num_attention_heads) * getattr( + self.config, "linear_key_head_dim", 128 + ) + lin_value_dim = getattr(self.config, "linear_num_value_heads", num_attention_heads) * getattr( + self.config, "linear_value_head_dim", 128 + ) lin_num_v_heads = getattr(self.config, "linear_num_value_heads", num_attention_heads) linear_proj_N = hidden_size * ( - lin_key_dim # q_proj - + lin_key_dim # k_proj - + lin_value_dim # v_proj - + lin_num_v_heads # a_proj - + lin_num_v_heads # b_proj - + lin_value_dim # g_proj (gate) - + lin_value_dim # o_proj - ) - linear_layer_flops = ( - m["attn_linear"] * linear_proj_N * tokens_sum - + m["dense_mlp"] * mlp_N * tokens_sum + lin_key_dim # q_proj + + lin_key_dim # k_proj + + lin_value_dim # v_proj + + lin_num_v_heads # a_proj + + lin_num_v_heads # b_proj + + lin_value_dim # g_proj (gate) + + lin_value_dim # o_proj ) + linear_layer_flops = m["attn_linear"] * linear_proj_N * tokens_sum + m["dense_mlp"] * mlp_N * tokens_sum dense_N_flops = ( full_attn_layer_flops * full_attn_layers @@ -314,13 +309,21 @@ def _estimate_qwen3_5_moe_flops(self, tokens_sum, batch_seqlens, delta_time): linear_layer_flops = 0 if linear_attn_layers > 0: - lin_key_dim = getattr(self.config, "linear_num_key_heads", num_attention_heads) * getattr(self.config, "linear_key_head_dim", 128) - lin_value_dim = getattr(self.config, "linear_num_value_heads", num_attention_heads) * getattr(self.config, "linear_value_head_dim", 128) + lin_key_dim = getattr(self.config, "linear_num_key_heads", num_attention_heads) * getattr( + self.config, "linear_key_head_dim", 128 + ) + lin_value_dim = getattr(self.config, "linear_num_value_heads", num_attention_heads) * getattr( + self.config, "linear_value_head_dim", 128 + ) lin_num_v_heads = getattr(self.config, "linear_num_value_heads", num_attention_heads) linear_proj_N = hidden_size * ( - lin_key_dim + lin_key_dim + lin_value_dim - + lin_num_v_heads + lin_num_v_heads - + lin_value_dim + lin_value_dim + lin_key_dim + + lin_key_dim + + lin_value_dim + + lin_num_v_heads + + lin_num_v_heads + + lin_value_dim + + lin_value_dim ) linear_layer_flops = m["attn_linear"] * linear_proj_N * tokens_sum + moe_mlp_flops @@ -351,14 +354,11 @@ def _estimate_qwen2_flops(self, tokens_sum, batch_seqlens, delta_time): k_size = num_key_value_heads * head_dim v_size = num_key_value_heads * head_dim - mlp_N = hidden_size * intermediate_size * 3 + mlp_N = hidden_size * intermediate_size * 3 attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim) - embed_lm_N = vocab_size * hidden_size # lm_head only; embedding is a lookup (0 FLOPs) + embed_lm_N = vocab_size * hidden_size # lm_head only; embedding is a lookup (0 FLOPs) - per_layer_flops = ( - m["dense_mlp"] * mlp_N * tokens_sum - + m["attn_linear"] * attn_linear_N * tokens_sum - ) + per_layer_flops = m["dense_mlp"] * mlp_N * tokens_sum + m["attn_linear"] * attn_linear_N * tokens_sum dense_N_flops = per_layer_flops * num_hidden_layers + 6 * embed_lm_N * tokens_sum seqlen_square_sum = sum(s * s for s in batch_seqlens) @@ -444,9 +444,7 @@ def _estimate_qwen_vit_flop(self, image_seqlens, config): # Qwen 2.5 VL uses SiLU, thus 3. mlp_N = dim * mlp_hidden_dim * (2 if is_qwen2_vl else 3) attn_linear_N = dim * (4 * dim) # qkv and output proj - patch_embed_and_merger_N = (out_hidden_size + (dim * (spatial_merge_size**2))) * ( - dim * (spatial_merge_size**2) - ) + patch_embed_and_merger_N = (out_hidden_size + (dim * (spatial_merge_size**2))) * (dim * (spatial_merge_size**2)) # non-attn all_layer parm dense_N = (mlp_N + attn_linear_N) * depth + patch_embed_and_merger_N diff --git a/src/xorl/utils/distillation_utils.py b/src/xorl/utils/distillation_utils.py index 7d1fc5f8..9a701c11 100644 --- a/src/xorl/utils/distillation_utils.py +++ b/src/xorl/utils/distillation_utils.py @@ -2,11 +2,10 @@ import json import logging -import os from huggingface_hub import hf_hub_download from safetensors import safe_open -import torch + logger = logging.getLogger(__name__) @@ -14,44 +13,35 @@ def load_lm_head_from_hf_model(model_id: str, token: str = None): """ Load only the lm_head weights from a HuggingFace model. - + Args: model_id: HuggingFace model ID (e.g., "Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8") token: HuggingFace token if model is private - + Returns: lm_head weight tensor """ # 1. Download the index file - index_path = hf_hub_download( - repo_id=model_id, - filename="model.safetensors.index.json", - token=token - ) - + index_path = hf_hub_download(repo_id=model_id, filename="model.safetensors.index.json", token=token) + # 2. Parse the index to find which shard contains lm_head - with open(index_path, 'r') as f: + with open(index_path, "r") as f: index = json.load(f) - + # 3. Find the shard file for lm_head.weight weight_map = index.get("weight_map", {}) lm_head_file = weight_map.get("lm_head.weight") - + if not lm_head_file: raise ValueError(f"lm_head.weight not found in {model_id}") - + logger.info(f"lm_head.weight is in: {lm_head_file}") - + # 4. Download only that specific shard - shard_path = hf_hub_download( - repo_id=model_id, - filename=lm_head_file, - token=token - ) - + shard_path = hf_hub_download(repo_id=model_id, filename=lm_head_file, token=token) + # 5. Load only the lm_head.weight from the shard with safe_open(shard_path, framework="pt", device="cpu") as f: lm_head_weight = f.get_tensor("lm_head.weight") - - return lm_head_weight + return lm_head_weight diff --git a/src/xorl/utils/helper.py b/src/xorl/utils/helper.py index ce313ded..115db65a 100644 --- a/src/xorl/utils/helper.py +++ b/src/xorl/utils/helper.py @@ -9,7 +9,7 @@ from collections import defaultdict from dataclasses import dataclass from functools import lru_cache -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional import numpy as np import psutil @@ -33,7 +33,6 @@ from .multisource_utils import parse_multisource_config - if TYPE_CHECKING: from torch.utils.data import DataLoader from transformers import PretrainedConfig @@ -44,9 +43,7 @@ CACHE_DIR = os.path.expanduser(os.getenv("CACHE_DIR", os.path.join("~/.cache", "xorl"))) -def _get_global_token_count( - micro_batch: Dict[str, "torch.Tensor"] -) -> List[int]: +def _get_global_token_count(micro_batch: Dict[str, "torch.Tensor"]) -> List[int]: """Return per-document token counts for FLOPs accounting. Uses ``_original_position_ids`` (pre-zigzag, pre-SP-padding) so document @@ -66,9 +63,7 @@ def _get_global_token_count( starts = (pos == 0).nonzero(as_tuple=True)[0] if len(starts) <= 1: return [int(pos.shape[0])] - seqlens = torch.diff( - torch.cat([starts, torch.tensor([pos.shape[0]], device=starts.device)]) - ).tolist() + seqlens = torch.diff(torch.cat([starts, torch.tensor([pos.shape[0]], device=starts.device)])).tolist() return [int(s) for s in seqlens] @@ -545,12 +540,12 @@ def handler_fn(p): warmup = 0 if start_step == 1 else 1 wait = start_step - warmup - 1 active = end_step - start_step - + # Ensure active >= 1 (PyTorch profiler requirement) if active <= 0: logger.warning(f"Profiler active steps is {active}, adjusting to 1 for valid profiling.") active = 1 - + logger.info(f"build profiler schedule - wait: {wait}, warmup: {warmup}, active: {active}.") schedule = profiler_module.schedule( diff --git a/src/xorl/utils/import_utils.py b/src/xorl/utils/import_utils.py index ee37928a..e37e3bc8 100644 --- a/src/xorl/utils/import_utils.py +++ b/src/xorl/utils/import_utils.py @@ -72,5 +72,3 @@ def is_torch_version_greater_than(value: str) -> bool: @lru_cache def is_transformers_version_greater_or_equal_to(value: str) -> bool: return _get_package_version("transformers") > version.parse(value) - - diff --git a/tests/conftest.py b/tests/conftest.py index 411af585..d2e2e6bc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,7 +1,8 @@ +from typing import Any, Dict, Sequence + import pytest import torch from torch.utils.data import Dataset -from typing import Dict, List, Sequence, Any class SimpleCollator: @@ -57,7 +58,7 @@ def __init__( min_seq_len: int = 64, max_seq_len: int = 256, vocab_size: int = 1000, - num_sequences_per_sample: int = 3 + num_sequences_per_sample: int = 3, ): self.num_samples = num_samples self.min_seq_len = min_seq_len @@ -109,11 +110,7 @@ def fake_text_dataset(): def fake_packed_dataset(): """Provides a fake packed dataset with position IDs.""" return FakePackedDataset( - num_samples=100, - min_seq_len=64, - max_seq_len=256, - vocab_size=1000, - num_sequences_per_sample=3 + num_samples=100, min_seq_len=64, max_seq_len=256, vocab_size=1000, num_sequences_per_sample=3 ) diff --git a/tests/data/collators/test_collate_pipeline.py b/tests/data/collators/test_collate_pipeline.py index b6073aa2..919b1fd1 100644 --- a/tests/data/collators/test_collate_pipeline.py +++ b/tests/data/collators/test_collate_pipeline.py @@ -1,7 +1,10 @@ +from typing import Any, Dict, Sequence + import pytest import torch + from xorl.data.collators import CollatePipeline, DataCollator -from typing import Dict, Sequence, Any + pytestmark = [pytest.mark.cpu, pytest.mark.collator] diff --git a/tests/data/collators/test_packing_concat_collator.py b/tests/data/collators/test_packing_concat_collator.py index 97a6b216..ae26657f 100644 --- a/tests/data/collators/test_packing_concat_collator.py +++ b/tests/data/collators/test_packing_concat_collator.py @@ -1,8 +1,11 @@ +from unittest.mock import Mock, patch + import pytest import torch -from unittest.mock import Mock, patch + from xorl.data.collators import PackingConcatCollator, add_flash_attention_kwargs_from_position_ids + pytestmark = [pytest.mark.cpu, pytest.mark.collator] @@ -12,7 +15,10 @@ class TestAddFlashAttentionKwargs: def test_kwargs_and_correctness(self): """Covers all required kwargs added, cu_seq_lens correctness, max_length correctness, and single sequence.""" # Multiple sequences: [0,1,2] and [0,1,2,3] - batch = {"input_ids": torch.tensor([[1, 2, 3, 4, 5, 6, 7]]), "position_ids": torch.tensor([[0, 1, 2, 0, 1, 2, 3]])} + batch = { + "input_ids": torch.tensor([[1, 2, 3, 4, 5, 6, 7]]), + "position_ids": torch.tensor([[0, 1, 2, 0, 1, 2, 3]]), + } cu_q, cu_k, max_q, max_k = add_flash_attention_kwargs_from_position_ids(batch) assert all(k in batch for k in ["cu_seq_lens_q", "cu_seq_lens_k", "max_length_q", "max_length_k"]) @@ -29,7 +35,7 @@ def test_kwargs_and_correctness(self): class TestPackingConcatCollator: """Tests for PackingConcatCollator.""" - @patch('xorl.data.collators.packing_concat_collator.get_parallel_state') + @patch("xorl.data.collators.packing_concat_collator.get_parallel_state") def test_concatenation_and_flash_attn(self, mock_parallel_state, sample_packed_features, sample_features): """Covers basic concatenation shapes, value correctness, position_ids generation, flash attn kwargs with SP disabled/enabled, and single feature handling.""" @@ -63,11 +69,18 @@ def test_concatenation_and_flash_attn(self, mock_parallel_state, sample_packed_f # Single feature mock_ps.cp_enabled = False - single = [{"input_ids": torch.tensor([1, 2, 3], dtype=torch.long), "attention_mask": torch.tensor([1, 1, 1], dtype=torch.long), "labels": torch.tensor([1, 2, 3], dtype=torch.long), "position_ids": torch.tensor([0, 1, 2], dtype=torch.long)}] + single = [ + { + "input_ids": torch.tensor([1, 2, 3], dtype=torch.long), + "attention_mask": torch.tensor([1, 1, 1], dtype=torch.long), + "labels": torch.tensor([1, 2, 3], dtype=torch.long), + "position_ids": torch.tensor([0, 1, 2], dtype=torch.long), + } + ] batch_single = collator(single) assert batch_single["input_ids"].shape == (1, 3) - @patch('xorl.data.collators.packing_concat_collator.get_parallel_state') + @patch("xorl.data.collators.packing_concat_collator.get_parallel_state") def test_extra_fields_and_multiple_seqs(self, mock_parallel_state): """Covers extra field handling and multiple packed sequences per sample.""" mock_ps = Mock() @@ -78,16 +91,38 @@ def test_extra_fields_and_multiple_seqs(self, mock_parallel_state): # Extra fields collated as batch features = [ - {"input_ids": torch.tensor([1, 2, 3], dtype=torch.long), "attention_mask": torch.tensor([1, 1, 1], dtype=torch.long), "labels": torch.tensor([1, 2, 3], dtype=torch.long), "position_ids": torch.tensor([0, 1, 2], dtype=torch.long), "extra_field": torch.tensor([100], dtype=torch.long)}, - {"input_ids": torch.tensor([4, 5, 6], dtype=torch.long), "attention_mask": torch.tensor([1, 1, 1], dtype=torch.long), "labels": torch.tensor([4, 5, 6], dtype=torch.long), "position_ids": torch.tensor([0, 1, 2], dtype=torch.long), "extra_field": torch.tensor([200], dtype=torch.long)}, + { + "input_ids": torch.tensor([1, 2, 3], dtype=torch.long), + "attention_mask": torch.tensor([1, 1, 1], dtype=torch.long), + "labels": torch.tensor([1, 2, 3], dtype=torch.long), + "position_ids": torch.tensor([0, 1, 2], dtype=torch.long), + "extra_field": torch.tensor([100], dtype=torch.long), + }, + { + "input_ids": torch.tensor([4, 5, 6], dtype=torch.long), + "attention_mask": torch.tensor([1, 1, 1], dtype=torch.long), + "labels": torch.tensor([4, 5, 6], dtype=torch.long), + "position_ids": torch.tensor([0, 1, 2], dtype=torch.long), + "extra_field": torch.tensor([200], dtype=torch.long), + }, ] batch = collator(features) assert "extra_field" in batch and batch["extra_field"].shape == (2, 1) # Multiple sequences per sample features2 = [ - {"input_ids": torch.tensor([1, 2, 3, 10, 11], dtype=torch.long), "attention_mask": torch.tensor([1, 1, 1, 1, 1], dtype=torch.long), "labels": torch.tensor([1, 2, 3, 10, 11], dtype=torch.long), "position_ids": torch.tensor([0, 1, 2, 0, 1], dtype=torch.long)}, - {"input_ids": torch.tensor([4, 5, 6, 7], dtype=torch.long), "attention_mask": torch.tensor([1, 1, 1, 1], dtype=torch.long), "labels": torch.tensor([4, 5, 6, 7], dtype=torch.long), "position_ids": torch.tensor([0, 1, 2, 3], dtype=torch.long)}, + { + "input_ids": torch.tensor([1, 2, 3, 10, 11], dtype=torch.long), + "attention_mask": torch.tensor([1, 1, 1, 1, 1], dtype=torch.long), + "labels": torch.tensor([1, 2, 3, 10, 11], dtype=torch.long), + "position_ids": torch.tensor([0, 1, 2, 0, 1], dtype=torch.long), + }, + { + "input_ids": torch.tensor([4, 5, 6, 7], dtype=torch.long), + "attention_mask": torch.tensor([1, 1, 1, 1], dtype=torch.long), + "labels": torch.tensor([4, 5, 6, 7], dtype=torch.long), + "position_ids": torch.tensor([0, 1, 2, 3], dtype=torch.long), + }, ] batch2 = collator(features2) assert batch2["input_ids"].shape == (1, 9) diff --git a/tests/data/collators/test_sequence_shard_collator.py b/tests/data/collators/test_sequence_shard_collator.py index 38d2e120..216ae5c9 100644 --- a/tests/data/collators/test_sequence_shard_collator.py +++ b/tests/data/collators/test_sequence_shard_collator.py @@ -1,9 +1,12 @@ +from unittest.mock import Mock, patch + import pytest import torch -from unittest.mock import Mock, patch + from xorl.data.collators import TextSequenceShardCollator from xorl.data.constants import IGNORE_INDEX + pytestmark = [pytest.mark.cpu, pytest.mark.collator] @@ -18,7 +21,7 @@ def _make_mock_ps(cp_size=2, cp_rank=0, ringattn_size=1): class TestSPSliceAndPadding: """Tests for sp_slice and sp_padding utility methods.""" - @patch('xorl.data.collators.sequence_shard_collator.get_parallel_state') + @patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") def test_sp_slice_across_ranks_and_uneven(self, mock_parallel_state): """Covers initialization, basic slicing rank 0/1, and uneven split.""" mock_parallel_state.return_value = _make_mock_ps(cp_size=2, cp_rank=0) @@ -38,22 +41,27 @@ def test_sp_slice_across_ranks_and_uneven(self, mock_parallel_state): collator0 = TextSequenceShardCollator() assert torch.equal(collator0.sp_slice(torch.tensor([[1, 2, 3, 4, 5]]), dim=-1), torch.tensor([[1, 2, 3]])) - @patch('xorl.data.collators.sequence_shard_collator.get_parallel_state') + @patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") def test_sp_padding_basic_sequential_zero(self, mock_parallel_state): """Covers basic padding, sequential padding, and zero-length padding.""" mock_parallel_state.return_value = _make_mock_ps(cp_size=2, cp_rank=0) collator = TextSequenceShardCollator() tensor = torch.tensor([[1, 2, 3]]) - assert torch.equal(collator.sp_padding(tensor, dim=-1, pad_value=0, pad_length=2), torch.tensor([[1, 2, 3, 0, 0]])) - assert torch.equal(collator.sp_padding(tensor, dim=-1, pad_value=0, pad_length=2, sequential=True), torch.tensor([[1, 2, 3, 0, 1]])) + assert torch.equal( + collator.sp_padding(tensor, dim=-1, pad_value=0, pad_length=2), torch.tensor([[1, 2, 3, 0, 0]]) + ) + assert torch.equal( + collator.sp_padding(tensor, dim=-1, pad_value=0, pad_length=2, sequential=True), + torch.tensor([[1, 2, 3, 0, 1]]), + ) assert torch.equal(collator.sp_padding(tensor, dim=-1, pad_value=0, pad_length=0), tensor) class TestCollatorCall: """Tests for the full collator __call__ method.""" - @patch('xorl.data.collators.sequence_shard_collator.get_parallel_state') + @patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") def test_preshifted_labels_and_packed_sequences(self, mock_parallel_state): """Covers pre-shifted labels pass-through and packed sequence boundary masking with cp_size=1.""" mock_parallel_state.return_value = _make_mock_ps(cp_size=1, cp_rank=0) @@ -80,7 +88,7 @@ def test_preshifted_labels_and_packed_sequences(self, mock_parallel_state): assert packed_result["labels"][0, 2] == IGNORE_INDEX assert packed_result["labels"][0, 4] == IGNORE_INDEX - @patch('xorl.data.collators.sequence_shard_collator.get_parallel_state') + @patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") def test_sp_splitting_padding_and_flash_attn_kwargs(self, mock_parallel_state): """Covers SP padding to multiple, splitting across ranks, flash attention kwargs, attention_mask/position_ids preservation, and padding values.""" diff --git a/tests/data/collators/test_tensor_collator.py b/tests/data/collators/test_tensor_collator.py index b086f5f5..974e6d7d 100644 --- a/tests/data/collators/test_tensor_collator.py +++ b/tests/data/collators/test_tensor_collator.py @@ -6,6 +6,7 @@ from xorl.data.collators import ToTensorCollator + pytestmark = [pytest.mark.cpu, pytest.mark.collator] @@ -17,7 +18,9 @@ def test_type_conversion_and_passthrough(self): collator = ToTensorCollator() # Lists to tensors - result = collator([{"input_ids": [1, 2, 3], "labels": [4, 5, 6]}, {"input_ids": [7, 8, 9], "labels": [10, 11, 12]}]) + result = collator( + [{"input_ids": [1, 2, 3], "labels": [4, 5, 6]}, {"input_ids": [7, 8, 9], "labels": [10, 11, 12]}] + ) assert isinstance(result, list) and len(result) == 2 assert isinstance(result[0]["input_ids"], torch.Tensor) and result[0]["input_ids"].shape == (3,) assert torch.equal(result[0]["input_ids"], torch.tensor([1, 2, 3])) @@ -55,7 +58,17 @@ def test_scalars_empty_and_dtype_inference(self): assert collator([]) == {} # Dtype inference - result_dtype = collator([{"input_ids": [1, 2, 3], "labels": [4, 5, 6], "position_ids": [0, 1, 2], "attention_mask": [1, 1, 1], "other_field": [1.0, 2.0, 3.0]}]) + result_dtype = collator( + [ + { + "input_ids": [1, 2, 3], + "labels": [4, 5, 6], + "position_ids": [0, 1, 2], + "attention_mask": [1, 1, 1], + "other_field": [1.0, 2.0, 3.0], + } + ] + ) assert result_dtype[0]["input_ids"].dtype == torch.long assert result_dtype[0]["labels"].dtype == torch.long assert result_dtype[0]["position_ids"].dtype == torch.long @@ -63,7 +76,14 @@ def test_scalars_empty_and_dtype_inference(self): assert result_dtype[0]["other_field"].dtype in [torch.float32, torch.float64] # Numpy dtype preservation - result_np = collator([{"input_ids": np.array([1, 2, 3], dtype=np.int32), "embeddings": np.array([1.0, 2.0, 3.0], dtype=np.float32)}]) + result_np = collator( + [ + { + "input_ids": np.array([1, 2, 3], dtype=np.int32), + "embeddings": np.array([1.0, 2.0, 3.0], dtype=np.float32), + } + ] + ) assert result_np[0]["embeddings"].dtype == torch.float32 # Batch size one @@ -96,15 +116,30 @@ def test_different_lengths_and_packed_sequences(self): collator = ToTensorCollator() # Different lengths fallback - result = collator([{"input_ids": [1, 2, 3], "labels": [4, 5, 6]}, {"input_ids": [7, 8, 9, 10, 11], "labels": [12, 13, 14, 15, 16]}]) + result = collator( + [ + {"input_ids": [1, 2, 3], "labels": [4, 5, 6]}, + {"input_ids": [7, 8, 9, 10, 11], "labels": [12, 13, 14, 15, 16]}, + ] + ) assert isinstance(result, list) assert result[0]["input_ids"].shape == (3,) assert result[1]["input_ids"].shape == (5,) # Packed sequences (same length, should be stacked) packed = [ - {"input_ids": [1, 2, 3, 101, 4, 5, 6, 102], "labels": [-100, -100, 7, -100, -100, 8, 9, -100], "position_ids": [0, 1, 2, 3, 0, 1, 2, 3], "length": 8}, - {"input_ids": [10, 11, 12, 101, 13, 14, 15, 16], "labels": [-100, 15, 16, -100, 17, 18, 19, 20], "position_ids": [0, 1, 2, 0, 1, 2, 3, 4], "length": 8}, + { + "input_ids": [1, 2, 3, 101, 4, 5, 6, 102], + "labels": [-100, -100, 7, -100, -100, 8, 9, -100], + "position_ids": [0, 1, 2, 3, 0, 1, 2, 3], + "length": 8, + }, + { + "input_ids": [10, 11, 12, 101, 13, 14, 15, 16], + "labels": [-100, 15, 16, -100, 17, 18, 19, 20], + "position_ids": [0, 1, 2, 0, 1, 2, 3, 4], + "length": 8, + }, ] result_packed = collator(packed) assert result_packed[0]["input_ids"].shape == (8,) diff --git a/tests/data/prepare/test_file_lock_loader.py b/tests/data/prepare/test_file_lock_loader.py index 7bd4ce5d..fcd11193 100644 --- a/tests/data/prepare/test_file_lock_loader.py +++ b/tests/data/prepare/test_file_lock_loader.py @@ -1,15 +1,15 @@ """Tests for xorl.data.prepare.file_lock_loader module.""" -import tempfile from pathlib import Path from unittest.mock import Mock, patch + import pytest from xorl.data.prepare.file_lock_loader import ( - FileLockLoader, LOCK_FILE_NAME, - READY_FILE_NAME, PROCESS_COUNTER_FILE_NAME, + READY_FILE_NAME, + FileLockLoader, ) @@ -122,7 +122,7 @@ def test_corrupted_counter_io_error_and_exceptions(self, mock_args, temp_dataset loader3 = FileLockLoader(mock_args) loader3.cleanup() Path(temp_dataset_path).mkdir(parents=True, exist_ok=True) - with patch.object(Path, 'read_text', side_effect=OSError("IO error")): + with patch.object(Path, "read_text", side_effect=OSError("IO error")): loader3.load(lambda: "data") assert loader3.counter_path.read_text().strip() == "1" diff --git a/tests/data/prepare/test_hash.py b/tests/data/prepare/test_hash.py index b47d86d1..3cb5ef8a 100644 --- a/tests/data/prepare/test_hash.py +++ b/tests/data/prepare/test_hash.py @@ -1,15 +1,16 @@ """Tests for xorl.data.prepare.hash module.""" -from unittest.mock import Mock, patch +from unittest.mock import Mock + import pytest from datasets import Dataset +from xorl.arguments import DatasetConfig from xorl.data.prepare.hash import ( - generate_split_fingerprints, - generate_packing_hash, generate_dataset_hash_from_config, + generate_packing_hash, + generate_split_fingerprints, ) -from xorl.arguments import DatasetConfig pytestmark = pytest.mark.cpu @@ -17,9 +18,16 @@ def _make_config(path="dataset1"): return DatasetConfig( - path=path, type="tokenized", shards=None, shards_idx=None, - preprocess_shards=None, name=None, split="train", revision=None, - trust_remote_code=False, max_seq_len=None, + path=path, + type="tokenized", + shards=None, + shards_idx=None, + preprocess_shards=None, + name=None, + split="train", + revision=None, + trust_remote_code=False, + max_seq_len=None, ) diff --git a/tests/data/prepare/test_packing.py b/tests/data/prepare/test_packing.py index 154db436..456fa0b5 100644 --- a/tests/data/prepare/test_packing.py +++ b/tests/data/prepare/test_packing.py @@ -1,20 +1,21 @@ """Tests for xorl.data.prepare.packing module.""" +from unittest.mock import Mock, patch + import numpy as np import pytest -from unittest.mock import Mock, patch from datasets import Dataset as HFDataset from xorl.data.prepare.packing import ( - ffd_check, - pack_group, - allocate_sequentially, + PackingDataset, add_position_ids, + allocate_sequentially, drop_no_trainable_tokens, + ffd_check, filter_dataset_with_logging, - process_datasets_for_packing, - PackingDataset, + pack_group, pack_parallel, + process_datasets_for_packing, ) @@ -128,10 +129,12 @@ def test_drop_no_trainable_tokens(): def test_filter_dataset_with_logging(): """Filter dataset and verify correct samples are kept.""" - dataset = HFDataset.from_dict({ - "input_ids": [[1, 2, 3], [4, 5, 6], [7, 8, 9]], - "labels": [[1, 2, 3], [-100, -100, -100], [7, 8, 9]], - }) + dataset = HFDataset.from_dict( + { + "input_ids": [[1, 2, 3], [4, 5, 6], [7, 8, 9]], + "labels": [[1, 2, 3], [-100, -100, -100], [7, 8, 9]], + } + ) filtered = filter_dataset_with_logging(dataset, lambda x: x["labels"][0] != -100, "test", num_proc=1) assert len(filtered) == 2 assert filtered[0]["labels"] == [1, 2, 3] @@ -171,11 +174,13 @@ def test_packing_dataset(): args.data.test_datasets = [] args.data.dataset_num_proc = 1 - dataset = HFDataset.from_dict({ - "input_ids": [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], - "labels": [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], - "length": [3, 3, 3, 3], - }) + dataset = HFDataset.from_dict( + { + "input_ids": [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], + "labels": [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], + "length": [3, 3, 3, 3], + } + ) tokenizer = Mock() tokenizer.name_or_path = "test-tokenizer" diff --git a/tests/data/prepare/test_shared.py b/tests/data/prepare/test_shared.py index 7bdb1720..b9676908 100644 --- a/tests/data/prepare/test_shared.py +++ b/tests/data/prepare/test_shared.py @@ -1,19 +1,20 @@ """Tests for xorl.data.prepare.shared module.""" from unittest.mock import Mock, patch + import pytest from datasets import Dataset as HFDataset +from xorl.arguments import DatasetConfig from xorl.data.prepare.shared import ( - get_dataset_type, - datasets_with_name_generator, - load_dataset_with_config, _check_if_hub_dataset, _get_remote_filesystem, create_train_validation_split, + datasets_with_name_generator, + get_dataset_type, + load_dataset_with_config, merge_datasets, ) -from xorl.arguments import DatasetConfig pytestmark = pytest.mark.cpu @@ -90,6 +91,7 @@ def test_hub_detection_and_remote_filesystem(self, mock_snapshot_download): # Invalid hub dataset from huggingface_hub.errors import RepositoryNotFoundError + mock_response = Mock() mock_response.status_code = 404 mock_response.headers = {} @@ -116,10 +118,12 @@ class TestSplitAndMerge: def test_split_and_merge_operations(self): """Covers absolute/fractional split, merge with shuffle variants, and empty merge error.""" - dataset = HFDataset.from_dict({ - "input_ids": [[i] for i in range(10)], - "labels": [[i] for i in range(10)], - }) + dataset = HFDataset.from_dict( + { + "input_ids": [[i] for i in range(10)], + "labels": [[i] for i in range(10)], + } + ) args = Mock() args.train.seed = 42 @@ -178,17 +182,25 @@ def test_local_and_hub_loading(self, tmp_path): assert ds is not None and len(ds) == 1 # Hub dataset - with patch("xorl.data.prepare.shared._check_if_hub_dataset", return_value=True), \ - patch("xorl.data.prepare.shared._load_from_hub", - return_value=HFDataset.from_dict({"input_ids": [[1]], "labels": [[1]]})) as mock_hub: + with ( + patch("xorl.data.prepare.shared._check_if_hub_dataset", return_value=True), + patch( + "xorl.data.prepare.shared._load_from_hub", + return_value=HFDataset.from_dict({"input_ids": [[1]], "labels": [[1]]}), + ) as mock_hub, + ): config = _make_config(path="username/dataset", split=None) assert load_dataset_with_config(config, use_auth_token=False, streaming=False) is not None mock_hub.assert_called_once() # HTTPS URL - with patch("xorl.data.prepare.shared._check_if_hub_dataset", return_value=False), \ - patch("xorl.data.prepare.shared._load_from_url", - return_value=HFDataset.from_dict({"input_ids": [[1]], "labels": [[1]]})) as mock_url: + with ( + patch("xorl.data.prepare.shared._check_if_hub_dataset", return_value=False), + patch( + "xorl.data.prepare.shared._load_from_url", + return_value=HFDataset.from_dict({"input_ids": [[1]], "labels": [[1]]}), + ) as mock_url, + ): config = _make_config(path="https://example.com/dataset.json", split=None) assert load_dataset_with_config(config, use_auth_token=False, streaming=False) is not None mock_url.assert_called_once() @@ -226,7 +238,7 @@ class TestSaveAndLoadPreprocessedDataset: def test_save_load_and_missing(self, tmp_path): """Covers save+load round-trip and load returning None when not found.""" - from xorl.data.prepare.shared import save_preprocessed_dataset, load_preprocessed_dataset + from xorl.data.prepare.shared import load_preprocessed_dataset, save_preprocessed_dataset args = Mock() args.data.dataset_prepared_path = str(tmp_path) @@ -236,10 +248,12 @@ def test_save_load_and_missing(self, tmp_path): args.data.skip_prepare_dataset = False args.data.is_preprocess = False - dataset = HFDataset.from_dict({ - "input_ids": [[1, 2, 3], [4, 5, 6]], - "labels": [[1, 2, 3], [4, 5, 6]], - }) + dataset = HFDataset.from_dict( + { + "input_ids": [[1, 2, 3], [4, 5, 6]], + "labels": [[1, 2, 3], [4, 5, 6]], + } + ) save_preprocessed_dataset(args, dataset, "test_hash_123", split="train") loaded = load_preprocessed_dataset(args, "test_hash_123") diff --git a/tests/data/prepare/test_utils.py b/tests/data/prepare/test_utils.py index 34b6522d..75407c91 100644 --- a/tests/data/prepare/test_utils.py +++ b/tests/data/prepare/test_utils.py @@ -3,14 +3,14 @@ import time from unittest.mock import Mock +import huggingface_hub import pytest import requests -import huggingface_hub from xorl.data.prepare.utils import ( RetryStrategy, - retry_on_request_exceptions, md5, + retry_on_request_exceptions, sha256, ) @@ -24,17 +24,23 @@ class TestRetryOnRequestExceptions: def test_success_retries_and_max_retries(self): """Covers immediate success, retry on ReadTimeout/HfHubHTTPError, max retries exhaustion, and non-request exception passthrough.""" + # Immediate success @retry_on_request_exceptions(max_retries=3, delay=0.01) def success_func(): return "success" + assert success_func() == "success" # Retry on ReadTimeout - mock_func = Mock(side_effect=[requests.exceptions.ReadTimeout("t"), requests.exceptions.ReadTimeout("t"), "success"]) + mock_func = Mock( + side_effect=[requests.exceptions.ReadTimeout("t"), requests.exceptions.ReadTimeout("t"), "success"] + ) + @retry_on_request_exceptions(max_retries=3, delay=0.01) def retry_func(): return mock_func() + assert retry_func() == "success" assert mock_func.call_count == 3 @@ -42,16 +48,21 @@ def retry_func(): mock_response = Mock() mock_response.status_code = 500 mock_response.headers = {} - mock_func2 = Mock(side_effect=[huggingface_hub.errors.HfHubHTTPError("HF error", response=mock_response), "success"]) + mock_func2 = Mock( + side_effect=[huggingface_hub.errors.HfHubHTTPError("HF error", response=mock_response), "success"] + ) + @retry_on_request_exceptions(max_retries=3, delay=0.01) def hf_func(): return mock_func2() + assert hf_func() == "success" # Max retries exhausted @retry_on_request_exceptions(max_retries=2, delay=0.01) def failing_func(): raise requests.exceptions.ReadTimeout("persistent timeout") + with pytest.raises(requests.exceptions.ReadTimeout): failing_func() @@ -59,6 +70,7 @@ def failing_func(): @retry_on_request_exceptions(max_retries=3, delay=0.01) def value_error_func(): raise ValueError("not a request exception") + with pytest.raises(ValueError): value_error_func() @@ -70,7 +82,9 @@ def test_backoff_strategies(self): (RetryStrategy.CONSTANT, lambda d1, d2: abs(d2 - d1) < d1 * 0.5), ]: call_times = [] - mock_func = Mock(side_effect=[requests.exceptions.ReadTimeout("t"), requests.exceptions.ReadTimeout("t"), "success"]) + mock_func = Mock( + side_effect=[requests.exceptions.ReadTimeout("t"), requests.exceptions.ReadTimeout("t"), "success"] + ) @retry_on_request_exceptions(max_retries=3, delay=0.1, retry_strategy=strategy) def func(): diff --git a/tests/data/test_data_loader.py b/tests/data/test_data_loader.py index de788a8c..52af64b6 100644 --- a/tests/data/test_data_loader.py +++ b/tests/data/test_data_loader.py @@ -1,14 +1,17 @@ +from typing import Any, Dict, Sequence +from unittest.mock import Mock, patch + import pytest import torch -from unittest.mock import Mock, patch from torch.utils.data import Dataset + +from xorl.data.collators import DataCollator, PackingConcatCollator from xorl.data.data_loader import ( - MicroBatchCollator, - DistributedDataloader, DataLoaderBuilder, + DistributedDataloader, + MicroBatchCollator, ) -from xorl.data.collators import DataCollator, PackingConcatCollator -from typing import Dict, Sequence, Any + pytestmark = [pytest.mark.cpu, pytest.mark.dataloader] @@ -31,7 +34,9 @@ def test_micro_batch_splitting_set_epoch_and_builder(self): internal_collator = SimpleCollator() # Multiple micro-batches with correct data - collator = MicroBatchCollator(micro_batch_size=2, gradient_accumulation_steps=2, internal_collator=internal_collator) + collator = MicroBatchCollator( + micro_batch_size=2, gradient_accumulation_steps=2, internal_collator=internal_collator + ) features = [{"input_ids": torch.tensor([i]), "labels": torch.tensor([i * 10])} for i in range(4)] result = collator(features) assert len(result) == 2 @@ -41,14 +46,18 @@ def test_micro_batch_splitting_set_epoch_and_builder(self): assert torch.equal(result[1]["labels"], torch.tensor([[20], [30]])) # Single micro-batch - collator_single = MicroBatchCollator(micro_batch_size=4, gradient_accumulation_steps=1, internal_collator=internal_collator) + collator_single = MicroBatchCollator( + micro_batch_size=4, gradient_accumulation_steps=1, internal_collator=internal_collator + ) features = [{"input_ids": torch.tensor([i]), "labels": torch.tensor([i])} for i in range(4)] result = collator_single(features) assert len(result) == 1 assert result[0]["input_ids"].shape[0] == 4 # Uneven division raises ValueError - collator_uneven = MicroBatchCollator(micro_batch_size=3, gradient_accumulation_steps=2, internal_collator=internal_collator) + collator_uneven = MicroBatchCollator( + micro_batch_size=3, gradient_accumulation_steps=2, internal_collator=internal_collator + ) features = [{"input_ids": torch.tensor([i]), "labels": torch.tensor([i])} for i in range(5)] with pytest.raises(ValueError, match="Expected 6 samples"): collator_uneven(features) @@ -76,10 +85,12 @@ def test_micro_batch_splitting_set_epoch_and_builder(self): dl3.dataset = Mock(spec=[]) dl3.set_epoch(5) - @patch('xorl.data.data_loader.get_parallel_state') - @patch('xorl.data.data_loader.StatefulDistributedSampler') - @patch('xorl.data.data_loader.DistributedDataloader') - def test_builder_batch_size_sampler_sp_and_custom_collators(self, mock_dataloader_cls, mock_sampler_cls, mock_parallel_state, fake_text_dataset): + @patch("xorl.data.data_loader.get_parallel_state") + @patch("xorl.data.data_loader.StatefulDistributedSampler") + @patch("xorl.data.data_loader.DistributedDataloader") + def test_builder_batch_size_sampler_sp_and_custom_collators( + self, mock_dataloader_cls, mock_sampler_cls, mock_parallel_state, fake_text_dataset + ): """Covers batch size, sampler params, SP collator, single/multiple custom collators.""" mock_ps = Mock() mock_ps.dp_size = 2 @@ -88,14 +99,18 @@ def test_builder_batch_size_sampler_sp_and_custom_collators(self, mock_dataloade mock_ps.cp_enabled = False mock_parallel_state.return_value = mock_ps - DataLoaderBuilder(dataset=fake_text_dataset, micro_batch_size=4, gradient_accumulation_steps=2, num_workers=4, seed=42).build(verbose=False) + DataLoaderBuilder( + dataset=fake_text_dataset, micro_batch_size=4, gradient_accumulation_steps=2, num_workers=4, seed=42 + ).build(verbose=False) assert mock_dataloader_cls.call_args[1]["batch_size"] == 8 # Sampler params mock_ps.dp_size = 4 mock_ps.dp_rank = 1 mock_parallel_state.return_value = mock_ps - DataLoaderBuilder(dataset=fake_text_dataset, micro_batch_size=2, gradient_accumulation_steps=2, seed=123).build(verbose=False) + DataLoaderBuilder(dataset=fake_text_dataset, micro_batch_size=2, gradient_accumulation_steps=2, seed=123).build( + verbose=False + ) call_kwargs = mock_sampler_cls.call_args[1] assert call_kwargs["num_replicas"] == 4 assert call_kwargs["rank"] == 1 @@ -108,10 +123,13 @@ def test_builder_batch_size_sampler_sp_and_custom_collators(self, mock_dataloade mock_ps.cp_size = 2 mock_ps.cp_enabled = True mock_parallel_state.return_value = mock_ps - DataLoaderBuilder(dataset=fake_text_dataset, micro_batch_size=2, gradient_accumulation_steps=2).build(verbose=False) + DataLoaderBuilder(dataset=fake_text_dataset, micro_batch_size=2, gradient_accumulation_steps=2).build( + verbose=False + ) collate_fn = mock_dataloader_cls.call_args[1]["collate_fn"] assert isinstance(collate_fn, MicroBatchCollator) from xorl.data.collators import CollatePipeline + assert isinstance(collate_fn.internal_collator, CollatePipeline) assert len(collate_fn.internal_collator.data_collators) == 5 @@ -120,7 +138,9 @@ def test_builder_batch_size_sampler_sp_and_custom_collators(self, mock_dataloade mock_ps.cp_enabled = False mock_parallel_state.return_value = mock_ps custom = SimpleCollator() - builder = DataLoaderBuilder(dataset=fake_text_dataset, micro_batch_size=2, gradient_accumulation_steps=2, use_default_collators=False) + builder = DataLoaderBuilder( + dataset=fake_text_dataset, micro_batch_size=2, gradient_accumulation_steps=2, use_default_collators=False + ) builder.add_collator(custom) builder.build(verbose=False) micro_batch_collator = mock_dataloader_cls.call_args[1]["collate_fn"] @@ -128,7 +148,9 @@ def test_builder_batch_size_sampler_sp_and_custom_collators(self, mock_dataloade assert micro_batch_collator.internal_collator is custom # Multiple custom collators - builder2 = DataLoaderBuilder(dataset=fake_text_dataset, micro_batch_size=2, gradient_accumulation_steps=2, use_default_collators=False) + builder2 = DataLoaderBuilder( + dataset=fake_text_dataset, micro_batch_size=2, gradient_accumulation_steps=2, use_default_collators=False + ) builder2.add_collator(SimpleCollator()) builder2.add_collator(SimpleCollator()) builder2.build(verbose=False) @@ -139,8 +161,8 @@ def test_builder_batch_size_sampler_sp_and_custom_collators(self, mock_dataloade class TestDataLoaderBuilderPipelineAndIntegration: """Tests for pipeline manipulation and integration with packed sequences.""" - @patch('xorl.data.data_loader.get_parallel_state') - @patch('xorl.data.collators.packing_concat_collator.get_parallel_state') + @patch("xorl.data.data_loader.get_parallel_state") + @patch("xorl.data.collators.packing_concat_collator.get_parallel_state") def test_pipeline_manipulation_and_packed_integration(self, mock_ps_collator, mock_ps_loader, fake_text_dataset): """Covers add/insert/remove, defaults flag, invalid position, empty build, training loop, nested lists, packed dataset, flatten+packing, and variable seqs.""" @@ -183,9 +205,13 @@ def test_pipeline_manipulation_and_packed_integration(self, mock_ps_collator, mo assert len(builder4.get_collator_pipeline()) == 4 # use_default_collators flag - with_defaults = DataLoaderBuilder(dataset=fake_text_dataset, micro_batch_size=2, gradient_accumulation_steps=2, use_default_collators=True) + with_defaults = DataLoaderBuilder( + dataset=fake_text_dataset, micro_batch_size=2, gradient_accumulation_steps=2, use_default_collators=True + ) assert len(with_defaults.get_collator_pipeline()) > 0 - without_defaults = DataLoaderBuilder(dataset=fake_text_dataset, micro_batch_size=2, gradient_accumulation_steps=2, use_default_collators=False) + without_defaults = DataLoaderBuilder( + dataset=fake_text_dataset, micro_batch_size=2, gradient_accumulation_steps=2, use_default_collators=False + ) assert len(without_defaults.get_collator_pipeline()) == 0 # Invalid position raises ValueError @@ -193,7 +219,9 @@ def test_pipeline_manipulation_and_packed_integration(self, mock_ps_collator, mo builder.add_collator(SimpleCollator(), position="middle") # Empty pipeline build raises ValueError - empty_builder = DataLoaderBuilder(dataset=fake_text_dataset, micro_batch_size=2, gradient_accumulation_steps=2, use_default_collators=False) + empty_builder = DataLoaderBuilder( + dataset=fake_text_dataset, micro_batch_size=2, gradient_accumulation_steps=2, use_default_collators=False + ) with pytest.raises(ValueError, match="Collator pipeline is empty"): empty_builder.build(verbose=False) @@ -202,8 +230,10 @@ class SimpleTokenizedDataset(Dataset): def __init__(self, num_samples=32, seq_len=10): self.num_samples = num_samples self.seq_len = seq_len + def __len__(self): return self.num_samples + def __getitem__(self, idx): return { "input_ids": torch.arange(idx * 10, idx * 10 + self.seq_len, dtype=torch.long), @@ -214,8 +244,12 @@ def __getitem__(self, idx): dataset = SimpleTokenizedDataset(num_samples=16, seq_len=5) int_builder = DataLoaderBuilder( - dataset=dataset, micro_batch_size=2, gradient_accumulation_steps=2, - num_workers=0, prefetch_factor=None, pad_to_multiple_of=1, + dataset=dataset, + micro_batch_size=2, + gradient_accumulation_steps=2, + num_workers=0, + prefetch_factor=None, + pad_to_multiple_of=1, ) dataloader = int_builder.build(verbose=False) @@ -233,12 +267,34 @@ def __getitem__(self, idx): # Nested list features handled correctly internal_collator = PackingConcatCollator() - micro_batch_collator = MicroBatchCollator(micro_batch_size=2, gradient_accumulation_steps=2, internal_collator=internal_collator) + micro_batch_collator = MicroBatchCollator( + micro_batch_size=2, gradient_accumulation_steps=2, internal_collator=internal_collator + ) correct_features = [ - {"input_ids": torch.tensor([1, 2]), "attention_mask": torch.ones(2), "labels": torch.tensor([1, 2]), "position_ids": torch.tensor([0, 1])}, - {"input_ids": torch.tensor([3, 4]), "attention_mask": torch.ones(2), "labels": torch.tensor([3, 4]), "position_ids": torch.tensor([0, 1])}, - {"input_ids": torch.tensor([5, 6]), "attention_mask": torch.ones(2), "labels": torch.tensor([5, 6]), "position_ids": torch.tensor([0, 1])}, - {"input_ids": torch.tensor([7, 8]), "attention_mask": torch.ones(2), "labels": torch.tensor([7, 8]), "position_ids": torch.tensor([0, 1])}, + { + "input_ids": torch.tensor([1, 2]), + "attention_mask": torch.ones(2), + "labels": torch.tensor([1, 2]), + "position_ids": torch.tensor([0, 1]), + }, + { + "input_ids": torch.tensor([3, 4]), + "attention_mask": torch.ones(2), + "labels": torch.tensor([3, 4]), + "position_ids": torch.tensor([0, 1]), + }, + { + "input_ids": torch.tensor([5, 6]), + "attention_mask": torch.ones(2), + "labels": torch.tensor([5, 6]), + "position_ids": torch.tensor([0, 1]), + }, + { + "input_ids": torch.tensor([7, 8]), + "attention_mask": torch.ones(2), + "labels": torch.tensor([7, 8]), + "position_ids": torch.tensor([0, 1]), + }, ] result = micro_batch_collator(correct_features) assert len(result) == 2 @@ -249,18 +305,22 @@ def __init__(self, num_samples=32, num_seqs_per_sample=2, seq_len=10): self.num_samples = num_samples self.num_seqs_per_sample = num_seqs_per_sample self.seq_len = seq_len + def __len__(self): return self.num_samples + def __getitem__(self, idx): sequences = [] for seq_idx in range(self.num_seqs_per_sample): offset = (idx * self.num_seqs_per_sample + seq_idx) * self.seq_len - sequences.append({ - "input_ids": torch.arange(offset, offset + self.seq_len, dtype=torch.long), - "attention_mask": torch.ones(self.seq_len, dtype=torch.long), - "labels": torch.arange(offset, offset + self.seq_len, dtype=torch.long), - "position_ids": torch.arange(self.seq_len, dtype=torch.long), - }) + sequences.append( + { + "input_ids": torch.arange(offset, offset + self.seq_len, dtype=torch.long), + "attention_mask": torch.ones(self.seq_len, dtype=torch.long), + "labels": torch.arange(offset, offset + self.seq_len, dtype=torch.long), + "position_ids": torch.arange(self.seq_len, dtype=torch.long), + } + ) return sequences packed_dataset = PackedSequencesDataset(num_samples=16, num_seqs_per_sample=2, seq_len=5) @@ -269,8 +329,12 @@ def __getitem__(self, idx): assert isinstance(sample[0], dict) packed_builder = DataLoaderBuilder( - dataset=packed_dataset, micro_batch_size=2, gradient_accumulation_steps=2, - num_workers=0, prefetch_factor=None, pad_to_multiple_of=1, + dataset=packed_dataset, + micro_batch_size=2, + gradient_accumulation_steps=2, + num_workers=0, + prefetch_factor=None, + pad_to_multiple_of=1, ) packed_dl = packed_builder.build(verbose=False) batch = next(iter(packed_dl)) @@ -281,16 +345,37 @@ def __getitem__(self, idx): # Flatten + Packing collator from xorl.data.collators import FlattenCollator + flatten_collator = FlattenCollator() packing_collator = PackingConcatCollator(pad_to_multiple_of=1) features = [ [ - {"input_ids": torch.tensor([1, 2, 3]), "attention_mask": torch.ones(3), "labels": torch.tensor([1, 2, 3]), "position_ids": torch.tensor([0, 1, 2])}, - {"input_ids": torch.tensor([4, 5]), "attention_mask": torch.ones(2), "labels": torch.tensor([4, 5]), "position_ids": torch.tensor([0, 1])}, + { + "input_ids": torch.tensor([1, 2, 3]), + "attention_mask": torch.ones(3), + "labels": torch.tensor([1, 2, 3]), + "position_ids": torch.tensor([0, 1, 2]), + }, + { + "input_ids": torch.tensor([4, 5]), + "attention_mask": torch.ones(2), + "labels": torch.tensor([4, 5]), + "position_ids": torch.tensor([0, 1]), + }, ], [ - {"input_ids": torch.tensor([6, 7]), "attention_mask": torch.ones(2), "labels": torch.tensor([6, 7]), "position_ids": torch.tensor([0, 1])}, - {"input_ids": torch.tensor([8, 9, 10]), "attention_mask": torch.ones(3), "labels": torch.tensor([8, 9, 10]), "position_ids": torch.tensor([0, 1, 2])}, + { + "input_ids": torch.tensor([6, 7]), + "attention_mask": torch.ones(2), + "labels": torch.tensor([6, 7]), + "position_ids": torch.tensor([0, 1]), + }, + { + "input_ids": torch.tensor([8, 9, 10]), + "attention_mask": torch.ones(3), + "labels": torch.tensor([8, 9, 10]), + "position_ids": torch.tensor([0, 1, 2]), + }, ], ] result = packing_collator(flatten_collator(features)) @@ -301,26 +386,70 @@ def __getitem__(self, idx): class VariablePackedDataset(Dataset): def __len__(self): return 4 + def __getitem__(self, idx): if idx == 0: - return [{"input_ids": torch.tensor([1, 2]), "attention_mask": torch.ones(2), "labels": torch.tensor([1, 2]), "position_ids": torch.tensor([0, 1])}] + return [ + { + "input_ids": torch.tensor([1, 2]), + "attention_mask": torch.ones(2), + "labels": torch.tensor([1, 2]), + "position_ids": torch.tensor([0, 1]), + } + ] elif idx == 1: return [ - {"input_ids": torch.tensor([3, 4, 5]), "attention_mask": torch.ones(3), "labels": torch.tensor([3, 4, 5]), "position_ids": torch.tensor([0, 1, 2])}, - {"input_ids": torch.tensor([6]), "attention_mask": torch.ones(1), "labels": torch.tensor([6]), "position_ids": torch.tensor([0])}, + { + "input_ids": torch.tensor([3, 4, 5]), + "attention_mask": torch.ones(3), + "labels": torch.tensor([3, 4, 5]), + "position_ids": torch.tensor([0, 1, 2]), + }, + { + "input_ids": torch.tensor([6]), + "attention_mask": torch.ones(1), + "labels": torch.tensor([6]), + "position_ids": torch.tensor([0]), + }, ] elif idx == 2: return [ - {"input_ids": torch.tensor([7, 8]), "attention_mask": torch.ones(2), "labels": torch.tensor([7, 8]), "position_ids": torch.tensor([0, 1])}, - {"input_ids": torch.tensor([9]), "attention_mask": torch.ones(1), "labels": torch.tensor([9]), "position_ids": torch.tensor([0])}, - {"input_ids": torch.tensor([10, 11]), "attention_mask": torch.ones(2), "labels": torch.tensor([10, 11]), "position_ids": torch.tensor([0, 1])}, + { + "input_ids": torch.tensor([7, 8]), + "attention_mask": torch.ones(2), + "labels": torch.tensor([7, 8]), + "position_ids": torch.tensor([0, 1]), + }, + { + "input_ids": torch.tensor([9]), + "attention_mask": torch.ones(1), + "labels": torch.tensor([9]), + "position_ids": torch.tensor([0]), + }, + { + "input_ids": torch.tensor([10, 11]), + "attention_mask": torch.ones(2), + "labels": torch.tensor([10, 11]), + "position_ids": torch.tensor([0, 1]), + }, ] else: - return [{"input_ids": torch.tensor([12, 13, 14]), "attention_mask": torch.ones(3), "labels": torch.tensor([12, 13, 14]), "position_ids": torch.tensor([0, 1, 2])}] + return [ + { + "input_ids": torch.tensor([12, 13, 14]), + "attention_mask": torch.ones(3), + "labels": torch.tensor([12, 13, 14]), + "position_ids": torch.tensor([0, 1, 2]), + } + ] var_builder = DataLoaderBuilder( - dataset=VariablePackedDataset(), micro_batch_size=2, gradient_accumulation_steps=1, - num_workers=0, prefetch_factor=None, pad_to_multiple_of=1, + dataset=VariablePackedDataset(), + micro_batch_size=2, + gradient_accumulation_steps=1, + num_workers=0, + prefetch_factor=None, + pad_to_multiple_of=1, ) var_dl = var_builder.build(verbose=False) var_batch = next(iter(var_dl)) diff --git a/tests/data/test_data_loader_distributed.py b/tests/data/test_data_loader_distributed.py index 156a0cb7..515483d2 100644 --- a/tests/data/test_data_loader_distributed.py +++ b/tests/data/test_data_loader_distributed.py @@ -4,17 +4,17 @@ These tests verify that the data loader correctly distributes data across multiple processes/ranks in a distributed training setting, and that the data alignment is correct. """ + +from unittest.mock import Mock, patch + import pytest import torch -import os -from unittest.mock import Mock, patch + +from tests.conftest import FakeTextDataset, SimpleCollator from xorl.data.data_loader import ( - MicroBatchCollator, - DistributedDataloader, DataLoaderBuilder, + MicroBatchCollator, ) -from xorl.data.collators import PackingConcatCollator -from tests.conftest import FakeTextDataset, SimpleCollator pytestmark = [pytest.mark.gpu, pytest.mark.dataloader, pytest.mark.distributed] @@ -23,9 +23,11 @@ class TestDistributedDataAlignment: """Tests for data partitioning, micro-batch splitting, SP sharding, batch size, and drop_last.""" - @patch('xorl.data.data_loader.get_parallel_state') - @patch('xorl.data.data_loader.StatefulDistributedSampler') - def test_partitioning_micro_batches_sp_batch_size_and_drop_last(self, mock_sampler_cls, mock_parallel_state, fake_text_dataset): + @patch("xorl.data.data_loader.get_parallel_state") + @patch("xorl.data.data_loader.StatefulDistributedSampler") + def test_partitioning_micro_batches_sp_batch_size_and_drop_last( + self, mock_sampler_cls, mock_parallel_state, fake_text_dataset + ): """Covers data partitioning across ranks, micro-batch splitting, SP pipeline setup, batch size calculation, and drop_last behavior.""" # --- Data partitioning across 4 DP ranks --- @@ -85,6 +87,7 @@ def test_partitioning_micro_batches_sp_batch_size_and_drop_last(self, mock_sampl dataloader_sp = builder_sp.build(verbose=False) from xorl.data.collators import CollatePipeline, TextSequenceShardCollator + internal = builder_sp.collate_fn.internal_collator assert isinstance(internal, CollatePipeline) assert any(isinstance(c, TextSequenceShardCollator) for c in internal.data_collators) @@ -111,7 +114,9 @@ def test_partitioning_micro_batches_sp_batch_size_and_drop_last(self, mock_sampl mock_sampler.__len__ = Mock(return_value=len(small_dataset)) mock_sampler_cls.return_value = mock_sampler - builder_dl = DataLoaderBuilder(dataset=small_dataset, micro_batch_size=2, gradient_accumulation_steps=2, drop_last=True) + builder_dl = DataLoaderBuilder( + dataset=small_dataset, micro_batch_size=2, gradient_accumulation_steps=2, drop_last=True + ) dl = builder_dl.build(verbose=False) batches = list(dl) assert len(batches) == 2 @@ -124,26 +129,32 @@ class TestMicroBatchCollatorAndSequenceSharding: def test_order_preservation_ga_alignment_and_error(self): """Covers data order within micro-batches, gradient accumulation alignment, and incorrect batch size error.""" - collator = MicroBatchCollator(micro_batch_size=2, gradient_accumulation_steps=3, internal_collator=SimpleCollator()) + collator = MicroBatchCollator( + micro_batch_size=2, gradient_accumulation_steps=3, internal_collator=SimpleCollator() + ) features = [{"input_ids": torch.tensor([i]), "value": torch.tensor([i * 100])} for i in range(6)] micro_batches = collator(features) assert torch.equal(micro_batches[0]["value"], torch.tensor([[0], [100]])) assert torch.equal(micro_batches[1]["value"], torch.tensor([[200], [300]])) assert torch.equal(micro_batches[2]["value"], torch.tensor([[400], [500]])) - collator2 = MicroBatchCollator(micro_batch_size=3, gradient_accumulation_steps=2, internal_collator=SimpleCollator()) + collator2 = MicroBatchCollator( + micro_batch_size=3, gradient_accumulation_steps=2, internal_collator=SimpleCollator() + ) features2 = [{"input_ids": torch.tensor([i]), "labels": torch.tensor([i])} for i in range(6)] mbs = collator2(features2) assert len(mbs) == 2 for mb in mbs: assert mb["input_ids"].shape[0] == 3 - collator3 = MicroBatchCollator(micro_batch_size=4, gradient_accumulation_steps=2, internal_collator=SimpleCollator()) + collator3 = MicroBatchCollator( + micro_batch_size=4, gradient_accumulation_steps=2, internal_collator=SimpleCollator() + ) features3 = [{"input_ids": torch.tensor([i]), "labels": torch.tensor([i])} for i in range(7)] with pytest.raises(ValueError, match="Expected 8 samples"): collator3(features3) - @patch('xorl.data.collators.sequence_shard_collator.get_parallel_state') + @patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") def test_sequence_sharding_and_padding(self, mock_parallel_state): """Covers correct chunk sizes across SP ranks and padding for non-divisible lengths.""" from xorl.data.collators import TextSequenceShardCollator @@ -198,9 +209,11 @@ def test_sequence_sharding_and_padding(self, mock_parallel_state): class TestEndToEndAndPackedSequences: """End-to-end integration tests for distributed data loading and packed sequences.""" - @patch('xorl.data.data_loader.get_parallel_state') - @patch('xorl.data.collators.packing_concat_collator.get_parallel_state') - def test_pipeline_epoch_consistency_and_packed_sequences(self, mock_ps_collator, mock_ps_loader, fake_packed_dataset, fake_text_dataset): + @patch("xorl.data.data_loader.get_parallel_state") + @patch("xorl.data.collators.packing_concat_collator.get_parallel_state") + def test_pipeline_epoch_consistency_and_packed_sequences( + self, mock_ps_collator, mock_ps_loader, fake_packed_dataset, fake_text_dataset + ): """Covers complete pipeline output, epoch consistency, packed sequence flattening, multi-DP, SP with packing, and variable lengths.""" mock_ps = Mock() @@ -213,8 +226,12 @@ def test_pipeline_epoch_consistency_and_packed_sequences(self, mock_ps_collator, # Complete pipeline output builder = DataLoaderBuilder( - dataset=fake_packed_dataset, micro_batch_size=2, gradient_accumulation_steps=2, - num_workers=0, prefetch_factor=None, seed=42, + dataset=fake_packed_dataset, + micro_batch_size=2, + gradient_accumulation_steps=2, + num_workers=0, + prefetch_factor=None, + seed=42, ) dataloader = builder.build(verbose=False) @@ -239,8 +256,12 @@ def test_pipeline_epoch_consistency_and_packed_sequences(self, mock_ps_collator, mock_ps_loader.return_value = mock_ps mock_ps_collator.return_value = mock_ps builder2 = DataLoaderBuilder( - dataset=fake_text_dataset, micro_batch_size=4, gradient_accumulation_steps=1, - num_workers=0, prefetch_factor=None, seed=42, + dataset=fake_text_dataset, + micro_batch_size=4, + gradient_accumulation_steps=1, + num_workers=0, + prefetch_factor=None, + seed=42, ) dl2 = builder2.build(verbose=False) @@ -255,10 +276,21 @@ def test_pipeline_epoch_consistency_and_packed_sequences(self, mock_ps_collator, class PackedDataset(Dataset): def __len__(self): return 8 + def __getitem__(self, idx): return [ - {"input_ids": torch.tensor([idx * 10, idx * 10 + 1, idx * 10 + 2]), "labels": torch.tensor([idx * 10, idx * 10 + 1, idx * 10 + 2]), "position_ids": torch.tensor([0, 1, 2]), "attention_mask": torch.ones(3, dtype=torch.long)}, - {"input_ids": torch.tensor([idx * 10 + 3, idx * 10 + 4]), "labels": torch.tensor([idx * 10 + 3, idx * 10 + 4]), "position_ids": torch.tensor([0, 1]), "attention_mask": torch.ones(2, dtype=torch.long)}, + { + "input_ids": torch.tensor([idx * 10, idx * 10 + 1, idx * 10 + 2]), + "labels": torch.tensor([idx * 10, idx * 10 + 1, idx * 10 + 2]), + "position_ids": torch.tensor([0, 1, 2]), + "attention_mask": torch.ones(3, dtype=torch.long), + }, + { + "input_ids": torch.tensor([idx * 10 + 3, idx * 10 + 4]), + "labels": torch.tensor([idx * 10 + 3, idx * 10 + 4]), + "position_ids": torch.tensor([0, 1]), + "attention_mask": torch.ones(2, dtype=torch.long), + }, ] mock_ps.dp_size = 1 @@ -269,8 +301,12 @@ def __getitem__(self, idx): mock_ps_loader.return_value = mock_ps builder3 = DataLoaderBuilder( - dataset=PackedDataset(), micro_batch_size=2, gradient_accumulation_steps=1, - num_workers=0, prefetch_factor=None, pad_to_multiple_of=1, + dataset=PackedDataset(), + micro_batch_size=2, + gradient_accumulation_steps=1, + num_workers=0, + prefetch_factor=None, + pad_to_multiple_of=1, ) dl3 = builder3.build(verbose=False) mbs = next(iter(dl3)) @@ -283,15 +319,30 @@ def __getitem__(self, idx): class SimplePacked(Dataset): def __len__(self): return 16 + def __getitem__(self, idx): - return [{"input_ids": torch.tensor([idx, idx + 1]), "labels": torch.tensor([idx, idx + 1]), "position_ids": torch.tensor([0, 1]), "attention_mask": torch.ones(2, dtype=torch.long)}] + return [ + { + "input_ids": torch.tensor([idx, idx + 1]), + "labels": torch.tensor([idx, idx + 1]), + "position_ids": torch.tensor([0, 1]), + "attention_mask": torch.ones(2, dtype=torch.long), + } + ] for dp_rank in [0, 1]: mock_ps.dp_size = 2 mock_ps.dp_rank = dp_rank mock_ps_collator.return_value = mock_ps mock_ps_loader.return_value = mock_ps - builder4 = DataLoaderBuilder(dataset=SimplePacked(), micro_batch_size=2, gradient_accumulation_steps=1, num_workers=0, prefetch_factor=None, pad_to_multiple_of=1) + builder4 = DataLoaderBuilder( + dataset=SimplePacked(), + micro_batch_size=2, + gradient_accumulation_steps=1, + num_workers=0, + prefetch_factor=None, + pad_to_multiple_of=1, + ) dl4 = builder4.build(verbose=False) mbs = next(iter(dl4)) assert len(mbs) == 1 @@ -301,14 +352,37 @@ def __getitem__(self, idx): class VariablePackedDataset(Dataset): def __len__(self): return 4 + def __getitem__(self, idx): if idx % 2 == 0: - return [{"input_ids": torch.tensor([idx, idx + 1]), "labels": torch.tensor([idx, idx + 1]), "position_ids": torch.tensor([0, 1]), "attention_mask": torch.ones(2, dtype=torch.long)}] + return [ + { + "input_ids": torch.tensor([idx, idx + 1]), + "labels": torch.tensor([idx, idx + 1]), + "position_ids": torch.tensor([0, 1]), + "attention_mask": torch.ones(2, dtype=torch.long), + } + ] else: return [ - {"input_ids": torch.tensor([idx]), "labels": torch.tensor([idx]), "position_ids": torch.tensor([0]), "attention_mask": torch.ones(1, dtype=torch.long)}, - {"input_ids": torch.tensor([idx + 1, idx + 2]), "labels": torch.tensor([idx + 1, idx + 2]), "position_ids": torch.tensor([0, 1]), "attention_mask": torch.ones(2, dtype=torch.long)}, - {"input_ids": torch.tensor([idx + 3]), "labels": torch.tensor([idx + 3]), "position_ids": torch.tensor([0]), "attention_mask": torch.ones(1, dtype=torch.long)}, + { + "input_ids": torch.tensor([idx]), + "labels": torch.tensor([idx]), + "position_ids": torch.tensor([0]), + "attention_mask": torch.ones(1, dtype=torch.long), + }, + { + "input_ids": torch.tensor([idx + 1, idx + 2]), + "labels": torch.tensor([idx + 1, idx + 2]), + "position_ids": torch.tensor([0, 1]), + "attention_mask": torch.ones(2, dtype=torch.long), + }, + { + "input_ids": torch.tensor([idx + 3]), + "labels": torch.tensor([idx + 3]), + "position_ids": torch.tensor([0]), + "attention_mask": torch.ones(1, dtype=torch.long), + }, ] mock_ps.dp_size = 1 @@ -317,7 +391,14 @@ def __getitem__(self, idx): mock_ps.cp_enabled = False mock_ps_collator.return_value = mock_ps mock_ps_loader.return_value = mock_ps - builder5 = DataLoaderBuilder(dataset=VariablePackedDataset(), micro_batch_size=2, gradient_accumulation_steps=1, num_workers=0, prefetch_factor=None, pad_to_multiple_of=1) + builder5 = DataLoaderBuilder( + dataset=VariablePackedDataset(), + micro_batch_size=2, + gradient_accumulation_steps=1, + num_workers=0, + prefetch_factor=None, + pad_to_multiple_of=1, + ) dl5 = builder5.build(verbose=False) mbs = next(iter(dl5)) assert len(mbs) == 1 diff --git a/tests/distributed/distributed_utils.py b/tests/distributed/distributed_utils.py index ced35bd7..c9b5a6bd 100644 --- a/tests/distributed/distributed_utils.py +++ b/tests/distributed/distributed_utils.py @@ -19,6 +19,7 @@ # GPU helpers # --------------------------------------------------------------------------- + def gpu_count() -> int: if not torch.cuda.is_available(): return 0 @@ -36,6 +37,7 @@ def skip_if_gpu_count_less_than(n: int): # Distributed test result # --------------------------------------------------------------------------- + @dataclass class DistributedTestResult: """Captures the outcome of a distributed test subprocess.""" @@ -62,6 +64,7 @@ def assert_success(self, msg: str = ""): # Torchrun launcher # --------------------------------------------------------------------------- + def _get_free_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind(("", 0)) @@ -80,9 +83,13 @@ def run_distributed_script( env.update(extra_env) cmd = [ - sys.executable, "-m", "torch.distributed.run", - "--nproc_per_node", str(num_gpus), - "--master_port", str(_get_free_port()), + sys.executable, + "-m", + "torch.distributed.run", + "--nproc_per_node", + str(num_gpus), + "--master_port", + str(_get_free_port()), script_path, ] diff --git a/tests/distributed/test_deepep_correctness.py b/tests/distributed/test_deepep_correctness.py index 7ba9e2e6..9686ca36 100644 --- a/tests/distributed/test_deepep_correctness.py +++ b/tests/distributed/test_deepep_correctness.py @@ -27,6 +27,7 @@ # ── Helpers ─────────────────────────────────────────────────────────────────── + def setup(): dist.init_process_group(backend="nccl") rank = dist.get_rank() @@ -66,6 +67,7 @@ def allclose_distributed(a: torch.Tensor, b: torch.Tensor, atol: float, rtol: fl # ── Core dispatch/combine runner ────────────────────────────────────────────── + def run_one_pass( ep_dispatch: str, hidden_states: torch.Tensor, @@ -81,11 +83,11 @@ def run_one_pass( buffer=None, ): """One forward + backward pass with a given dispatch backend.""" - from xorl.models.layers.moe.backend import EP_DISPATCH, EP_COMBINE, EP_EXPERT_COMPUTE + from xorl.models.layers.moe.backend import EP_COMBINE, EP_DISPATCH, EP_EXPERT_COMPUTE dispatch_fn = EP_DISPATCH[ep_dispatch] - combine_fn = EP_COMBINE[ep_dispatch] - compute_fn = EP_EXPERT_COMPUTE["triton"] + combine_fn = EP_COMBINE[ep_dispatch] + compute_fn = EP_EXPERT_COMPUTE["triton"] x = hidden_states.detach().requires_grad_(True) @@ -117,6 +119,7 @@ def run_one_pass( # ── Single test case ────────────────────────────────────────────────────────── + def run_test( num_tokens: int, hidden_dim: int, @@ -131,7 +134,9 @@ def run_test( num_experts = world_size num_local_experts = 1 - log(f" [{num_tokens}tok, h={hidden_dim}, ffn={intermediate_size}, top{topk}] ", ) + log( + f" [{num_tokens}tok, h={hidden_dim}, ffn={intermediate_size}, top{topk}] ", + ) ep_group = dist.new_group(list(range(world_size))) @@ -139,7 +144,7 @@ def run_test( torch.manual_seed(0) scale = (2.0 / (hidden_dim + intermediate_size)) ** 0.5 gate = nn.Parameter(torch.randn(1, hidden_dim, intermediate_size, device=device, dtype=torch.bfloat16) * scale) - up = nn.Parameter(torch.randn(1, hidden_dim, intermediate_size, device=device, dtype=torch.bfloat16) * scale) + up = nn.Parameter(torch.randn(1, hidden_dim, intermediate_size, device=device, dtype=torch.bfloat16) * scale) down = nn.Parameter(torch.randn(1, intermediate_size, hidden_dim, device=device, dtype=torch.bfloat16) * scale) for p in [gate, up, down]: dist.broadcast(p.data, src=0) @@ -158,21 +163,44 @@ def run_test( # ── AllToAll ────────────────────────────────────────────────────────────── out_a2a, grad_a2a = run_one_pass( - "alltoall", hidden, topk_w, topk_idx, - gate, up, down, ep_group, num_experts, num_local_experts, grad_out, + "alltoall", + hidden, + topk_w, + topk_idx, + gate, + up, + down, + ep_group, + num_experts, + num_local_experts, + grad_out, ) - gate.grad = None; up.grad = None; down.grad = None + gate.grad = None + up.grad = None + down.grad = None # ── DeepEP ─────────────────────────────────────────────────────────────── from xorl.distributed.moe.deepep import DeepEPBuffer + buffer = DeepEPBuffer(ep_group=ep_group, buffer_size_gb=1.0) out_dep, grad_dep = run_one_pass( - "deepep", hidden, topk_w, topk_idx, - gate, up, down, ep_group, num_experts, num_local_experts, grad_out, + "deepep", + hidden, + topk_w, + topk_idx, + gate, + up, + down, + ep_group, + num_experts, + num_local_experts, + grad_out, buffer=buffer, ) buffer.destroy_buffer() - gate.grad = None; up.grad = None; down.grad = None + gate.grad = None + up.grad = None + down.grad = None # ── Check ───────────────────────────────────────────────────────────────── fwd_abs = max_abs_diff(out_dep, out_a2a) @@ -182,8 +210,7 @@ def run_test( bwd_ok = allclose_distributed(grad_dep, grad_a2a, atol=atol, rtol=rtol) log( - f" fwd max_abs={fwd_abs:.2e} {'✓' if fwd_ok else '✗'} | " - f"bwd max_abs={bwd_abs:.2e} {'✓' if bwd_ok else '✗'}" + f" fwd max_abs={fwd_abs:.2e} {'✓' if fwd_ok else '✗'} | bwd max_abs={bwd_abs:.2e} {'✓' if bwd_ok else '✗'}" ) return fwd_ok and bwd_ok @@ -191,20 +218,20 @@ def run_test( # ── Test suite ──────────────────────────────────────────────────────────────── + def run_all_tests() -> bool: world_size = dist.get_world_size() cross_node = world_size > 8 - log(f"\n{'='*65}") - log(f" DeepEP vs AllToAll Correctness — world_size={world_size} " - f"({'cross-node' if cross_node else 'single-node'})") - log(f"{'='*65}") + log(f"\n{'=' * 65}") + log(f" DeepEP vs AllToAll Correctness — world_size={world_size} ({'cross-node' if cross_node else 'single-node'})") + log(f"{'=' * 65}") configs = [ - dict(num_tokens=16, hidden_dim=256, intermediate_size=512, topk=2), - dict(num_tokens=32, hidden_dim=512, intermediate_size=1024, topk=4), - dict(num_tokens=64, hidden_dim=1024, intermediate_size=2048, topk=4), - dict(num_tokens=64, hidden_dim=2048, intermediate_size=4096, topk=8), + dict(num_tokens=16, hidden_dim=256, intermediate_size=512, topk=2), + dict(num_tokens=32, hidden_dim=512, intermediate_size=1024, topk=4), + dict(num_tokens=64, hidden_dim=1024, intermediate_size=2048, topk=4), + dict(num_tokens=64, hidden_dim=2048, intermediate_size=4096, topk=8), ] all_ok = True @@ -214,7 +241,7 @@ def run_all_tests() -> bool: all_ok = False log(f"\n {'[PASS] All tests passed ✓' if all_ok else '[FAIL] Some tests failed ✗'}") - log(f"{'='*65}\n") + log(f"{'=' * 65}\n") return all_ok @@ -224,6 +251,7 @@ def run_all_tests() -> bool: # Add NVSHMEM lib to LD_LIBRARY_PATH for DeepEP internode transport try: import nvidia.nvshmem + nvshmem_lib = os.path.join(list(nvidia.nvshmem.__path__)[0], "lib") existing = os.environ.get("LD_LIBRARY_PATH", "") if nvshmem_lib not in existing: diff --git a/tests/distributed/test_ep_lora_weight_slicing.py b/tests/distributed/test_ep_lora_weight_slicing.py index d6161bdd..a70696e8 100644 --- a/tests/distributed/test_ep_lora_weight_slicing.py +++ b/tests/distributed/test_ep_lora_weight_slicing.py @@ -8,14 +8,15 @@ 4. Gradients flow correctly through the EP path """ +from unittest.mock import MagicMock, patch + import pytest import torch -import torch.nn as nn -from unittest.mock import MagicMock, patch from torch.distributed._tensor import Shard from xorl.models.layers.moe import MoEExpertsLoRA, MoELoRAConfig + pytestmark = [pytest.mark.distributed] @@ -144,9 +145,17 @@ def test_ep_and_non_ep_forward_with_gradients(self): scaling = 8.0 / 4.0 output = TritonEPGroupGemmWithLoRA.apply( - permute_tokens, cumsum, - gate_proj, up_proj, down_proj, - gate_A, gate_B, up_A, up_B, down_A, down_B, + permute_tokens, + cumsum, + gate_proj, + up_proj, + down_proj, + gate_A, + gate_B, + up_A, + up_B, + down_A, + down_B, scaling, ) assert output.shape == (num_tokens, hidden_dim) @@ -170,9 +179,19 @@ def test_ep_and_non_ep_forward_with_gradients(self): expert_index = torch.randint(0, num_experts, (num_tokens2, top_k), device=device) output2 = TritonMoeExpertsLoRAFunction.apply( - num_experts, gate_weights, expert_index, hidden_states, - gate_proj2, up_proj2, down_proj2, - gate_A2, gate_B2, up_A2, up_B2, down_A2, down_B2, + num_experts, + gate_weights, + expert_index, + hidden_states, + gate_proj2, + up_proj2, + down_proj2, + gate_A2, + gate_B2, + up_A2, + up_B2, + down_A2, + down_B2, scaling, ) assert output2.shape == (num_tokens2, hidden_dim) diff --git a/tests/distributed/test_linear_attention_cp_equivalence.py b/tests/distributed/test_linear_attention_cp_equivalence.py index d7834a9c..310fe39d 100644 --- a/tests/distributed/test_linear_attention_cp_equivalence.py +++ b/tests/distributed/test_linear_attention_cp_equivalence.py @@ -15,12 +15,14 @@ from xorl.ops.linear_attention.ops.cp import build_linear_attention_cp_context from xorl.utils.device import get_nccl_backend + THIS_DIR = Path(__file__).resolve().parent if str(THIS_DIR) not in sys.path: sys.path.insert(0, str(THIS_DIR)) from distributed_utils import run_distributed_script, skip_if_gpu_count_less_than + pytestmark = [pytest.mark.distributed] @@ -160,6 +162,7 @@ def _main() -> None: if __name__ != "__main__": + @skip_if_gpu_count_less_than(2) def test_linear_attention_cp_matches_single_gpu_reference(): result = run_distributed_script( diff --git a/tests/distributed/test_parallel_state.py b/tests/distributed/test_parallel_state.py index 5473e140..1442df0d 100644 --- a/tests/distributed/test_parallel_state.py +++ b/tests/distributed/test_parallel_state.py @@ -1,17 +1,19 @@ """Tests for xorl.distributed.parallel_state module.""" +from unittest.mock import Mock, patch + import pytest import torch -from unittest.mock import Mock, patch from xorl.distributed.parallel_state import ( ParallelState, get_parallel_state, - init_parallel_state, init_ep_mesh_matrix, + init_parallel_state, requires_mesh, ) + pytestmark = [pytest.mark.cpu, pytest.mark.distributed] @@ -37,6 +39,7 @@ def test_ep_mesh_matrix_and_requires_mesh(self): class MC: def __init__(self, m): self.device_mesh = m + @requires_mesh def go(self): return "ok" @@ -49,7 +52,7 @@ def go(self): class TestParallelStateConstruction: """Test ParallelState construction, validation, properties, and enabled flags.""" - @patch('xorl.distributed.parallel_state.dist.is_initialized', return_value=False) + @patch("xorl.distributed.parallel_state.dist.is_initialized", return_value=False) def test_defaults_and_uninitialized_properties(self, mock_is_init): """Default ParallelState: all sizes 1, not initialized, rank/world_size defaults; invalid cp_fsdp_mode raises.""" state = ParallelState() @@ -63,10 +66,10 @@ def test_defaults_and_uninitialized_properties(self, mock_is_init): with pytest.raises(ValueError, match="Invalid cp_fsdp_mode"): ParallelState(cp_fsdp_mode="invalid") - @patch('xorl.distributed.parallel_state.dist.is_initialized', return_value=True) - @patch('xorl.distributed.parallel_state.dist.get_rank', return_value=5) - @patch('xorl.distributed.parallel_state.dist.get_world_size', return_value=8) - @patch('xorl.distributed.sequence_parallel.init_sequence_parallel') + @patch("xorl.distributed.parallel_state.dist.is_initialized", return_value=True) + @patch("xorl.distributed.parallel_state.dist.get_rank", return_value=5) + @patch("xorl.distributed.parallel_state.dist.get_world_size", return_value=8) + @patch("xorl.distributed.sequence_parallel.init_sequence_parallel") def test_custom_init_validation_and_enabled_flags(self, mock_sp, mock_ws, mock_rank, mock_init): """Custom init; validation errors; initialized properties; sp/fsdp enabled flags.""" state = ParallelState(dp_size=4, dp_replicate_size=2, dp_shard_size=2, tp_size=2) @@ -90,29 +93,31 @@ class TestGetAndInitParallelState: def setup_method(self): import xorl.distributed.parallel_state as ps_module + ps_module._PARALLEL_STATE = None def teardown_method(self): import xorl.distributed.parallel_state as ps_module + ps_module._PARALLEL_STATE = None - @patch('xorl.distributed.parallel_state.is_torch_version_greater_than', return_value=False) - @patch('xorl.distributed.parallel_state.dist.is_initialized', return_value=True) - @patch('xorl.distributed.parallel_state.dist.get_world_size', return_value=8) + @patch("xorl.distributed.parallel_state.is_torch_version_greater_than", return_value=False) + @patch("xorl.distributed.parallel_state.dist.is_initialized", return_value=True) + @patch("xorl.distributed.parallel_state.dist.get_world_size", return_value=8) def test_init_get_reinit_auto_dp_shard_default(self, mock_ws, mock_is_init, mock_version): """init sets state, get retrieves, re-init warns, auto dp_shard_size, get default when unset.""" init_parallel_state(dp_size=4, tp_size=2, dp_mode="fsdp2") state = get_parallel_state() assert state.dp_size == 4 and state.tp_size == 2 and state.dp_mode == "fsdp2" - with patch('xorl.distributed.parallel_state.logger.warning') as mock_warn: + with patch("xorl.distributed.parallel_state.logger.warning") as mock_warn: init_parallel_state(dp_size=8) mock_warn.assert_called_once_with("Parallel state has already been initialized.") assert get_parallel_state().dp_size == 4 - @patch('xorl.distributed.parallel_state.is_torch_version_greater_than', return_value=False) - @patch('xorl.distributed.parallel_state.dist.is_initialized', return_value=True) - @patch('xorl.distributed.parallel_state.dist.get_world_size', return_value=4) + @patch("xorl.distributed.parallel_state.is_torch_version_greater_than", return_value=False) + @patch("xorl.distributed.parallel_state.dist.is_initialized", return_value=True) + @patch("xorl.distributed.parallel_state.dist.get_world_size", return_value=4) def test_auto_dp_shard_and_default_uninitialized(self, mock_ws, mock_is_init, mock_version): """Auto dp_shard_size; device_type defaults; get_parallel_state returns default when unset.""" init_parallel_state(dp_size=4) diff --git a/tests/distributed/test_pipeline_parallel.py b/tests/distributed/test_pipeline_parallel.py index 3c746976..dd4ec4dd 100644 --- a/tests/distributed/test_pipeline_parallel.py +++ b/tests/distributed/test_pipeline_parallel.py @@ -8,6 +8,7 @@ from xorl.distributed.pipeline_parallel import generate_llm_fqn_per_model_part + pytestmark = [pytest.mark.distributed] @@ -46,7 +47,8 @@ def test_basic_stage_distribution(self): def test_qwen3_fqn_names_and_single_stage(self): """Qwen3-style nested FQN names; single stage contains all modules.""" result = generate_llm_fqn_per_model_part( - 2, 4, + 2, + 4, input_fqns=["model.embed_tokens"], layer_prefix="model.layers", output_fqns=["model.norm", "lm_head"], diff --git a/tests/distributed/test_qwen3_5_ulysses_cp.py b/tests/distributed/test_qwen3_5_ulysses_cp.py index f0a2d4eb..e877d047 100644 --- a/tests/distributed/test_qwen3_5_ulysses_cp.py +++ b/tests/distributed/test_qwen3_5_ulysses_cp.py @@ -16,12 +16,14 @@ from xorl.ops.linear_attention.ops.cp import build_linear_attention_cp_context from xorl.utils.device import get_nccl_backend + THIS_DIR = Path(__file__).resolve().parent if str(THIS_DIR) not in sys.path: sys.path.insert(0, str(THIS_DIR)) from distributed_utils import run_distributed_script, skip_if_gpu_count_less_than + pytestmark = [pytest.mark.distributed] @@ -151,6 +153,7 @@ def _main() -> None: if __name__ != "__main__": + @skip_if_gpu_count_less_than(2) def test_qwen35_ulysses_positive_smoke(): result = run_distributed_script( @@ -161,7 +164,6 @@ def test_qwen35_ulysses_positive_smoke(): ) result.assert_success("Qwen3.5 positive Ulysses smoke should pass") - @skip_if_gpu_count_less_than(2) def test_qwen35_ring_fla_negative_smoke(): result = run_distributed_script( diff --git a/tests/distributed/test_ring_attention.py b/tests/distributed/test_ring_attention.py index b726022d..e0f6b36b 100644 --- a/tests/distributed/test_ring_attention.py +++ b/tests/distributed/test_ring_attention.py @@ -6,13 +6,14 @@ import pytest import torch +from xorl.data.collators.sequence_shard_collator import ( + zigzag_reorder_packed_sequence, +) from xorl.distributed.sequence_parallel.ring_attention import ( _get_zigzag_step_section, _merge_attn_outputs, ) -from xorl.data.collators.sequence_shard_collator import ( - zigzag_reorder_packed_sequence, -) + pytestmark = [pytest.mark.distributed] @@ -91,10 +92,14 @@ def test_zigzag_sections_and_reorder(self): tensor = torch.arange(40).unsqueeze(0) position_ids = torch.arange(40).unsqueeze(0) reordered = zigzag_reorder_packed_sequence(tensor, position_ids, ringattn_size, dim=-1) - expected = torch.cat([ - torch.arange(0, 10), torch.arange(30, 40), - torch.arange(10, 20), torch.arange(20, 30), - ]).unsqueeze(0) + expected = torch.cat( + [ + torch.arange(0, 10), + torch.arange(30, 40), + torch.arange(10, 20), + torch.arange(20, 30), + ] + ).unsqueeze(0) assert torch.equal(reordered, expected) # Multi-doc reorder @@ -103,12 +108,18 @@ def test_zigzag_sections_and_reorder(self): tensor_m = torch.arange(total).unsqueeze(0) position_ids_m = torch.cat([torch.arange(doc_len), torch.arange(doc_len)]).unsqueeze(0) reordered_m = zigzag_reorder_packed_sequence(tensor_m, position_ids_m, ringattn_size, dim=-1) - expected_m = torch.cat([ - torch.arange(0, 5), torch.arange(15, 20), - torch.arange(20, 25), torch.arange(35, 40), - torch.arange(5, 10), torch.arange(10, 15), - torch.arange(25, 30), torch.arange(30, 35), - ]).unsqueeze(0) + expected_m = torch.cat( + [ + torch.arange(0, 5), + torch.arange(15, 20), + torch.arange(20, 25), + torch.arange(35, 40), + torch.arange(5, 10), + torch.arange(10, 15), + torch.arange(25, 30), + torch.arange(30, 35), + ] + ).unsqueeze(0) assert torch.equal(reordered_m, expected_m) # Position IDs doc boundaries per rank @@ -141,6 +152,4 @@ def test_zigzag_sections_and_reorder(self): # Invalid length raises with pytest.raises(ValueError, match="not divisible"): - zigzag_reorder_packed_sequence( - torch.arange(15).unsqueeze(0), torch.arange(15).unsqueeze(0), 2, dim=-1 - ) + zigzag_reorder_packed_sequence(torch.arange(15).unsqueeze(0), torch.arange(15).unsqueeze(0), 2, dim=-1) diff --git a/tests/distributed/test_sequence_parallel.py b/tests/distributed/test_sequence_parallel.py index a55f36da..9f028d57 100644 --- a/tests/distributed/test_sequence_parallel.py +++ b/tests/distributed/test_sequence_parallel.py @@ -11,6 +11,7 @@ ) from xorl.distributed.sequence_parallel.utils import pad_tensor, unpad_tensor + pytestmark = [pytest.mark.distributed] @@ -40,7 +41,9 @@ def test_pad_unpad_roundtrip_and_dims(self): # Roundtrip xr = torch.randn(3, 7, 5) - assert torch.allclose(xr, unpad_tensor(pad_tensor(xr, dim=1, padding_size=3, padding_value=0), dim=1, padding_size=3)) + assert torch.allclose( + xr, unpad_tensor(pad_tensor(xr, dim=1, padding_size=3, padding_value=0), dim=1, padding_size=3) + ) class TestSlicePositionEmbedding: diff --git a/tests/distributed/test_tensor_parallel.py b/tests/distributed/test_tensor_parallel.py index 0976d7cb..bc82a893 100644 --- a/tests/distributed/test_tensor_parallel.py +++ b/tests/distributed/test_tensor_parallel.py @@ -8,6 +8,7 @@ import pytest import torch + pytestmark = [pytest.mark.cpu, pytest.mark.distributed] @@ -51,7 +52,7 @@ def test_attention_and_mlp_unfuse(self): def test_unfused_forward_shape_and_model_level(self): """Unfused MLP forward produces same shape; model-level unfuse covers all layers.""" from xorl.models.transformers.qwen3.configuration_qwen3 import Qwen3Config - from xorl.models.transformers.qwen3.modeling_qwen3 import Qwen3MLP, Qwen3ForCausalLM + from xorl.models.transformers.qwen3.modeling_qwen3 import Qwen3ForCausalLM, Qwen3MLP config = Qwen3Config( hidden_size=256, @@ -87,6 +88,5 @@ def test_unfused_forward_shape_and_model_level(self): assert hasattr(layer.mlp, "up_proj") - if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/distributed/test_utils.py b/tests/distributed/test_utils.py index e4f79755..a1b33093 100644 --- a/tests/distributed/test_utils.py +++ b/tests/distributed/test_utils.py @@ -4,13 +4,14 @@ import torch.nn as nn from xorl.distributed.utils import ( - set_module_from_path, - get_module_from_path, check_all_fqn_match, check_any_fqn_match, check_fqn_match, + get_module_from_path, + set_module_from_path, ) + pytestmark = [pytest.mark.cpu, pytest.mark.distributed] @@ -21,11 +22,7 @@ def __init__(self): super().__init__() self.layer1 = nn.Linear(10, 20) self.layer2 = nn.Linear(20, 30) - self.nested = nn.Sequential( - nn.Linear(30, 40), - nn.ReLU(), - nn.Linear(40, 50) - ) + self.nested = nn.Sequential(nn.Linear(30, 40), nn.ReLU(), nn.Linear(40, 50)) class TestModulePaths: @@ -47,11 +44,11 @@ def test_set_get_roundtrip_and_nested(self): assert model.nested[0].out_features == 100 # Deeply nested via ModuleDict - model.deep = nn.ModuleDict({'a': nn.Sequential(nn.Linear(5, 10))}) + model.deep = nn.ModuleDict({"a": nn.Sequential(nn.Linear(5, 10))}) deep_layer = nn.Linear(5, 20) set_module_from_path(model, "deep.a.0", deep_layer) assert get_module_from_path(model, "deep.a.0") is deep_layer - assert model.deep['a'][0].out_features == 20 + assert model.deep["a"][0].out_features == 20 # Nonexistent paths raise with pytest.raises(AttributeError): @@ -108,10 +105,12 @@ def test_all_fqn_matching(self): assert check_all_fqn_match(["layer1.weight", "layer2.bias"], ["layer3.weight", "layer4.bias"]) is False # Multiple wildcards same number - assert check_all_fqn_match( - ["block*.layer*.weight", "block*.layer*.bias"], - ["block1.layer1.weight", "block1.layer1.bias"] - ) is True + assert ( + check_all_fqn_match( + ["block*.layer*.weight", "block*.layer*.bias"], ["block1.layer1.weight", "block1.layer1.bias"] + ) + is True + ) # Empty lists assert check_all_fqn_match([], []) is True diff --git a/tests/distributed/test_vocab_parallel_ce.py b/tests/distributed/test_vocab_parallel_ce.py index d7998b1b..37fe23fa 100644 --- a/tests/distributed/test_vocab_parallel_ce.py +++ b/tests/distributed/test_vocab_parallel_ce.py @@ -50,8 +50,12 @@ def check_correctness(rank, world_size, tp_group): for use_compile in [False, True]: par_ce = vocab_parallel_cross_entropy( - hidden_states, local_weight, labels, tp_group, - ignore_index=-100, use_compile=use_compile, + hidden_states, + local_weight, + labels, + tp_group, + ignore_index=-100, + use_compile=use_compile, ) err = (par_ce - ref_ce).abs().max().item() mode = "compiled" if use_compile else "eager" @@ -87,8 +91,12 @@ def check_backward(rank, world_size, tp_group): w_par = full_weight[rank * local_V : (rank + 1) * local_V].contiguous().requires_grad_(True) par_ce = vocab_parallel_cross_entropy( - h_par, w_par, labels, tp_group, - ignore_index=-100, use_compile=use_compile, + h_par, + w_par, + labels, + tp_group, + ignore_index=-100, + use_compile=use_compile, ) (par_ce.sum() / valid).backward() @@ -105,8 +113,7 @@ def check_backward(rank, world_size, tp_group): assert grad_w_err < 1e-2, f"grad_w error too large ({mode}): {grad_w_err}" -def _bench_one(hidden_states, local_weight, labels, tp_group, use_compile, - n_warmup=5, n_iter=20, fwd_only=False): +def _bench_one(hidden_states, local_weight, labels, tp_group, use_compile, n_warmup=5, n_iter=20, fwd_only=False): """Run warmup + timed iterations, return (ms/iter, peak_memory_MB).""" for _ in range(n_warmup): ce = vocab_parallel_cross_entropy(hidden_states, local_weight, labels, tp_group, use_compile=use_compile) @@ -164,8 +171,12 @@ def bench_perf(rank, world_size, tp_group): for use_compile in [False, True]: mode = "compiled" if use_compile else "eager" ms, peak_act_mb, peak_total_mb = _bench_one( - hidden_states, local_weight, labels, tp_group, - use_compile=use_compile, fwd_only=True, + hidden_states, + local_weight, + labels, + tp_group, + use_compile=use_compile, + fwd_only=True, ) if rank == 0: print(f"{mode:<12} {ms:<12.2f} {peak_act_mb:<16.1f} {peak_total_mb:.1f}") @@ -180,8 +191,12 @@ def bench_perf(rank, world_size, tp_group): for use_compile in [False, True]: mode = "compiled" if use_compile else "eager" ms, peak_act_mb, peak_total_mb = _bench_one( - hidden_states, local_weight, labels, tp_group, - use_compile=use_compile, fwd_only=False, + hidden_states, + local_weight, + labels, + tp_group, + use_compile=use_compile, + fwd_only=False, ) if rank == 0: print(f"{mode:<12} {ms:<12.2f} {peak_act_mb:<16.1f} {peak_total_mb:.1f}") @@ -223,6 +238,7 @@ def main(): if __name__ != "__main__": # Only define pytest tests when imported by pytest (not when run via torchrun) import pytest + from tests.distributed.distributed_utils import run_distributed_script, skip_if_gpu_count_less_than SCRIPT_PATH = os.path.abspath(__file__) diff --git a/tests/e2e/e2e_utils.py b/tests/e2e/e2e_utils.py index 7919696b..c27247bf 100644 --- a/tests/e2e/e2e_utils.py +++ b/tests/e2e/e2e_utils.py @@ -15,15 +15,13 @@ import socket import subprocess import sys -from dataclasses import dataclass, field -from pathlib import Path +from dataclasses import dataclass from typing import Any, Dict, List, Optional, Tuple +import pytest import torch import yaml -import pytest - # --------------------------------------------------------------------------- # Real model IDs (downloaded/cached via HF hub) @@ -36,6 +34,7 @@ # GPU helpers # --------------------------------------------------------------------------- + def gpu_count() -> int: if not torch.cuda.is_available(): return 0 @@ -65,19 +64,16 @@ def _has_quack() -> bool: return False -skip_if_no_flash_attn = pytest.mark.skipif( - not _has_flash_attn(), reason="flash_attn not installed" -) +skip_if_no_flash_attn = pytest.mark.skipif(not _has_flash_attn(), reason="flash_attn not installed") -skip_if_no_quack = pytest.mark.skipif( - not _has_quack(), reason="quack not installed" -) +skip_if_no_quack = pytest.mark.skipif(not _has_quack(), reason="quack not installed") # --------------------------------------------------------------------------- # Training result # --------------------------------------------------------------------------- + @dataclass class TrainingResult: """Captures the outcome of a training subprocess.""" @@ -132,16 +128,11 @@ def assert_loss_converged( "Need at least 2 loss values to check convergence" ) first, last = self.loss_history[0], self.loss_history[-1] - assert not math.isnan(first) and not math.isnan(last), ( - f"NaN in loss history: first={first}, last={last}" - ) - assert last < max_final_loss, ( - f"Final loss {last:.4f} >= {max_final_loss}" - ) + assert not math.isnan(first) and not math.isnan(last), f"NaN in loss history: first={first}, last={last}" + assert last < max_final_loss, f"Final loss {last:.4f} >= {max_final_loss}" drop = (first - last) / first assert drop >= min_drop_ratio, ( - f"Loss drop {drop:.2%} < required {min_drop_ratio:.0%} " - f"(first={first:.4f}, last={last:.4f})" + f"Loss drop {drop:.2%} < required {min_drop_ratio:.0%} (first={first:.4f}, last={last:.4f})" ) def assert_success(self, msg: str = ""): @@ -154,9 +145,7 @@ def assert_success(self, msg: str = ""): f"--- stdout (last 80 lines) ---\n{stdout_tail}\n" f"--- stderr (last 50 lines) ---\n{stderr_tail}" ) - assert self.metrics is not None, ( - f"Training exited 0 but no training_metrics.json found in {self.output_dir}" - ) + assert self.metrics is not None, f"Training exited 0 but no training_metrics.json found in {self.output_dir}" # --------------------------------------------------------------------------- @@ -252,6 +241,7 @@ def create_tiny_model_dir( def _save_random_weights(model_dir: str, model_config: dict): from transformers import AutoConfig, AutoModelForCausalLM + config = AutoConfig.from_pretrained(model_dir) model = AutoModelForCausalLM.from_config(config) model.save_pretrained(model_dir, safe_serialization=True) @@ -284,6 +274,7 @@ def _create_tokenizer_files(model_dir: str, vocab_size: int = 1024): # YAML config generation # --------------------------------------------------------------------------- + def generate_training_config( model_dir: str, output_dir: str, @@ -395,15 +386,19 @@ def generate_training_config( "enable_lora": True, "lora_rank": lora_rank, "lora_alpha": lora_alpha, - "lora_target_modules": lora_target_modules or [ - "q_proj", "k_proj", "v_proj", "o_proj", - "gate_proj", "up_proj", "down_proj", + "lora_target_modules": lora_target_modules + or [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", ], } if enable_qlora: - group_size = quant_group_size or ( - 16 if quant_format == "nvfp4" else 64 if quant_format == "nf4" else 128 - ) + group_size = quant_group_size or (16 if quant_format == "nvfp4" else 64 if quant_format == "nf4" else 128) config["lora"]["enable_qlora"] = True config["lora"]["quant_format"] = quant_format config["lora"]["quant_group_size"] = group_size @@ -423,6 +418,7 @@ def generate_training_config( # Training launcher # --------------------------------------------------------------------------- + def _get_free_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind(("", 0)) @@ -441,10 +437,15 @@ def run_training( env.update(extra_env) cmd = [ - sys.executable, "-m", "torch.distributed.run", - "--nproc_per_node", str(num_gpus), - "--master_port", str(_get_free_port()), - "-m", "xorl.cli.train", + sys.executable, + "-m", + "torch.distributed.run", + "--nproc_per_node", + str(num_gpus), + "--master_port", + str(_get_free_port()), + "-m", + "xorl.cli.train", config_path, ] @@ -457,10 +458,12 @@ def run_training( env=env, ) except subprocess.TimeoutExpired as e: + def _to_str(v, fallback): if isinstance(v, bytes): return v.decode(errors="replace") return v or fallback + return TrainingResult( exit_code=-1, stdout=_to_str(e.stdout, ""), @@ -498,6 +501,7 @@ def _read_metrics(output_dir: str) -> Optional[Dict[str, Any]]: # Shared data generation (for server vs local comparison tests) # --------------------------------------------------------------------------- + def generate_shared_token_data( num_samples: int, seq_len: int = 64, diff --git a/tests/e2e/qwen3_30b/compare_lora_qlora.py b/tests/e2e/qwen3_30b/compare_lora_qlora.py old mode 100644 new mode 100755 index 6659d268..6b30a7a9 --- a/tests/e2e/qwen3_30b/compare_lora_qlora.py +++ b/tests/e2e/qwen3_30b/compare_lora_qlora.py @@ -10,6 +10,7 @@ import sys import tempfile + ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) sys.path.insert(0, ROOT) @@ -88,7 +89,9 @@ def main(): **extra, ) results[name] = run_training( - config_path, num_gpus=num_gpus, timeout=3600, + config_path, + num_gpus=num_gpus, + timeout=3600, ) # --- Report --- diff --git a/tests/e2e/qwen3_30b/test_pp.py b/tests/e2e/qwen3_30b/test_pp.py index 9b691d4c..54e242f6 100644 --- a/tests/e2e/qwen3_30b/test_pp.py +++ b/tests/e2e/qwen3_30b/test_pp.py @@ -20,12 +20,13 @@ _create_full_weight_client, _get_free_port, _start_server_or_fail, - generate_server_config, + assert_loss_decreases, generate_random_sft_data, + generate_server_config, run_sft_steps, - assert_loss_decreases, ) + pytestmark = [pytest.mark.e2e, pytest.mark.gpu, pytest.mark.slow] # EP=4 and CP=4 are folded onto the same GPU axis: PP=2 * 4 = 8 GPUs total. @@ -36,7 +37,6 @@ class TestPP30BTrainer: - @skip_if_gpu_count_less_than(NUM_GPUS) def test_pp2_ep4_cp4_muon_loss_converges(self, tiny_moe_model_dir): """Trainer: PP=2 + EP=4 + CP=4 (folded, 8 GPUs) with Muon.""" @@ -68,7 +68,6 @@ def test_pp2_ep4_cp4_muon_loss_converges(self, tiny_moe_model_dir): class TestPP30BServer: - @skip_if_gpu_count_less_than(NUM_GPUS) def test_pp2_ep4_cp4_server_loss_decreases(self, tiny_moe_model_dir_with_weights): """Server: PP=2 + EP=4 + CP=4 (folded, 8 GPUs) — loss must be finite and decreasing.""" diff --git a/tests/e2e/qwen3_30b/test_server_moe.py b/tests/e2e/qwen3_30b/test_server_moe.py index 1c55d163..4beb9311 100644 --- a/tests/e2e/qwen3_30b/test_server_moe.py +++ b/tests/e2e/qwen3_30b/test_server_moe.py @@ -16,6 +16,7 @@ run_sft_steps, ) + pytestmark = [pytest.mark.e2e, pytest.mark.gpu, pytest.mark.server, pytest.mark.slow] @@ -49,9 +50,7 @@ def test_moe_ep2_fsdp4_lora_sft(self, tmp_workspace): try: _start_server_or_fail(server, timeout=600) - _, training_client = _create_lora_client( - server.base_url, model_dir, model_id="test-moe-ep2-fsdp4" - ) + _, training_client = _create_lora_client(server.base_url, model_dir, model_id="test-moe-ep2-fsdp4") data = generate_random_sft_data(num_samples=8, seq_len=64, vocab_size=151936) losses = run_sft_steps(training_client, data, num_steps=2) assert all(not math.isnan(l) for l in losses), f"NaN in losses: {losses}" diff --git a/tests/e2e/qwen3_8b/test_lora.py b/tests/e2e/qwen3_8b/test_lora.py index 8c2edeea..8292dc8e 100644 --- a/tests/e2e/qwen3_8b/test_lora.py +++ b/tests/e2e/qwen3_8b/test_lora.py @@ -11,11 +11,11 @@ skip_if_gpu_count_less_than, ) + pytestmark = [pytest.mark.e2e, pytest.mark.gpu, pytest.mark.slow] class TestLoRA1GPU: - @skip_if_gpu_count_less_than(1) def test_lora_loss_converges(self, tmp_workspace): """Qwen3-8B LoRA training shows strong loss convergence over 20 steps.""" @@ -39,7 +39,6 @@ def test_lora_loss_converges(self, tmp_workspace): class TestLoRA2GPU: - @skip_if_gpu_count_less_than(2) def test_lora_fsdp2(self, tmp_workspace): """Qwen3-8B LoRA + FSDP2 on 2 GPUs converges.""" diff --git a/tests/e2e/qwen3_8b/test_pp.py b/tests/e2e/qwen3_8b/test_pp.py index de20f9d1..d69944fd 100644 --- a/tests/e2e/qwen3_8b/test_pp.py +++ b/tests/e2e/qwen3_8b/test_pp.py @@ -32,6 +32,7 @@ run_sft_steps, ) + pytestmark = [pytest.mark.e2e, pytest.mark.gpu, pytest.mark.slow] @@ -39,8 +40,8 @@ # Trainer tests # --------------------------------------------------------------------------- -class TestPP2GPU: +class TestPP2GPU: @skip_if_gpu_count_less_than(2) def test_pp2_loss_converges(self, tiny_dense_model_dir): """Trainer: PP=2 on 2 GPUs — loss must be finite and decreasing.""" @@ -62,7 +63,6 @@ def test_pp2_loss_converges(self, tiny_dense_model_dir): class TestPP8GPU: - @skip_if_gpu_count_less_than(8) def test_pp2_fsdp4_loss_converges(self, tiny_dense_model_dir): """Trainer: PP=2 + FSDP=4 on 8 GPUs — validates fsdp_size=1 fix.""" @@ -108,8 +108,8 @@ def test_pp2_fsdp4_muon_loss_converges(self, tiny_dense_model_dir): # Server tests # --------------------------------------------------------------------------- -class TestPP2GPUServer: +class TestPP2GPUServer: @skip_if_gpu_count_less_than(2) def test_pp2_server_loss_decreases(self, tiny_dense_model_dir_with_weights): """Server: PP=2 on 2 GPUs — validates server PP padding + loss normalization fix.""" @@ -138,7 +138,6 @@ def test_pp2_server_loss_decreases(self, tiny_dense_model_dir_with_weights): class TestPP8GPUServer: - @skip_if_gpu_count_less_than(8) def test_pp2_fsdp4_server_loss_decreases(self, tiny_dense_model_dir_with_weights): """Server: PP=2 + FSDP=4 on 8 GPUs — validates reported_loss fix (no /fsdp_size).""" diff --git a/tests/e2e/qwen3_8b/test_tflops_threshold.py b/tests/e2e/qwen3_8b/test_tflops_threshold.py index e7517d64..94ec09b8 100644 --- a/tests/e2e/qwen3_8b/test_tflops_threshold.py +++ b/tests/e2e/qwen3_8b/test_tflops_threshold.py @@ -29,6 +29,7 @@ generate_server_config, ) + pytestmark = [pytest.mark.e2e, pytest.mark.gpu, pytest.mark.server, pytest.mark.benchmark] # --------------------------------------------------------------------------- @@ -58,6 +59,7 @@ # Helpers # --------------------------------------------------------------------------- + def _skip_if_no_qwen3_8b(): return pytest.mark.skipif( not os.path.isdir(QWEN3_8B_DIR), @@ -89,13 +91,9 @@ def _run_tflops_benchmark(num_gpus: int, dp_shard_size: int, tmp_path: str): try: _start_server_or_fail(server, timeout=300) - _, training_client = _create_lora_client( - server.base_url, QWEN3_8B_DIR, model_id=f"bench-{num_gpus}gpu" - ) + _, training_client = _create_lora_client(server.base_url, QWEN3_8B_DIR, model_id=f"bench-{num_gpus}gpu") - data = generate_random_sft_data( - num_samples=NUM_SAMPLES, seq_len=SEQ_LEN, vocab_size=VOCAB_SIZE - ) + data = generate_random_sft_data(num_samples=NUM_SAMPLES, seq_len=SEQ_LEN, vocab_size=VOCAB_SIZE) adam_params = xorl_client.AdamParams(learning_rate=1e-4, beta1=0.9, beta2=0.95, eps=1e-8) # Warmup @@ -128,15 +126,15 @@ def _run_tflops_benchmark(num_gpus: int, dp_shard_size: int, tmp_path: str): device_peak_tflops = get_device_flops(unit="T") mfu = flops_achieved / device_peak_tflops if device_peak_tflops > 0 else 0.0 - print(f"\n{'─'*60}") + print(f"\n{'─' * 60}") print(f"Qwen3-8B LoRA SFT — {num_gpus} GPU(s)") - print(f"{'─'*60}") + print(f"{'─' * 60}") print(f" Loss: {losses[0]:.4f} → {losses[-1]:.4f}") print(f" Avg step: {avg_step:.3f}s") print(f" Tokens/sec: {tokens_per_sec:.0f}") print(f" TFLOPS: {flops_achieved:.2f} / {device_peak_tflops:.0f} peak") print(f" MFU: {mfu:.2%}") - print(f"{'─'*60}") + print(f"{'─' * 60}") return { "tflops": flops_achieved, @@ -153,6 +151,7 @@ def _run_tflops_benchmark(num_gpus: int, dp_shard_size: int, tmp_path: str): # Tests # --------------------------------------------------------------------------- + class TestQwen3_8B_TFLOPS: """Assert Qwen3-8B LoRA training meets TFLOPS thresholds on H100.""" @@ -161,9 +160,7 @@ class TestQwen3_8B_TFLOPS: def test_1gpu_tflops(self, tmp_path): """1-GPU Qwen3-8B LoRA must achieve >= 358 TFLOPS on H100.""" result = _run_tflops_benchmark(num_gpus=1, dp_shard_size=1, tmp_path=str(tmp_path)) - assert result["tflops"] >= MIN_TFLOPS[1], ( - f"1-GPU TFLOPS {result['tflops']:.1f} below threshold {MIN_TFLOPS[1]}" - ) + assert result["tflops"] >= MIN_TFLOPS[1], f"1-GPU TFLOPS {result['tflops']:.1f} below threshold {MIN_TFLOPS[1]}" assert result["losses"][-1] < result["losses"][0], "Loss should decrease" @skip_if_gpu_count_less_than(2) @@ -171,9 +168,7 @@ def test_1gpu_tflops(self, tmp_path): def test_2gpu_tflops(self, tmp_path): """2-GPU Qwen3-8B LoRA FSDP2 must achieve >= 336 TFLOPS on H100.""" result = _run_tflops_benchmark(num_gpus=2, dp_shard_size=2, tmp_path=str(tmp_path)) - assert result["tflops"] >= MIN_TFLOPS[2], ( - f"2-GPU TFLOPS {result['tflops']:.1f} below threshold {MIN_TFLOPS[2]}" - ) + assert result["tflops"] >= MIN_TFLOPS[2], f"2-GPU TFLOPS {result['tflops']:.1f} below threshold {MIN_TFLOPS[2]}" assert result["losses"][-1] < result["losses"][0], "Loss should decrease" @skip_if_gpu_count_less_than(4) @@ -181,7 +176,5 @@ def test_2gpu_tflops(self, tmp_path): def test_4gpu_tflops(self, tmp_path): """4-GPU Qwen3-8B LoRA FSDP2 must achieve >= 329 TFLOPS on H100.""" result = _run_tflops_benchmark(num_gpus=4, dp_shard_size=4, tmp_path=str(tmp_path)) - assert result["tflops"] >= MIN_TFLOPS[4], ( - f"4-GPU TFLOPS {result['tflops']:.1f} below threshold {MIN_TFLOPS[4]}" - ) + assert result["tflops"] >= MIN_TFLOPS[4], f"4-GPU TFLOPS {result['tflops']:.1f} below threshold {MIN_TFLOPS[4]}" assert result["losses"][-1] < result["losses"][0], "Loss should decrease" diff --git a/tests/e2e/server_utils.py b/tests/e2e/server_utils.py index 7fd41566..6a515fec 100644 --- a/tests/e2e/server_utils.py +++ b/tests/e2e/server_utils.py @@ -14,6 +14,7 @@ import requests import yaml + def generate_server_config( model_dir: str, output_dir: str, @@ -60,8 +61,13 @@ def generate_server_config( if lora_target_modules is None: lora_target_modules = [ - "q_proj", "k_proj", "v_proj", "o_proj", - "gate_proj", "up_proj", "down_proj", + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", ] # Build config matching examples/server/ flat format @@ -135,6 +141,7 @@ def generate_server_config( # Server process management # --------------------------------------------------------------------------- + def _get_free_port() -> int: """Find a free port that is not in TIME_WAIT state.""" with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: @@ -199,11 +206,17 @@ def start(self, timeout: float = 180.0) -> bool: master_port = _get_free_port() cmd = [ - sys.executable, "-m", "xorl.server.launcher", - "--mode", "auto", - "--config", self.config_path, - "--api-port", str(self.api_port), - "--master-port", str(master_port), + sys.executable, + "-m", + "xorl.server.launcher", + "--mode", + "auto", + "--config", + self.config_path, + "--api-port", + str(self.api_port), + "--master-port", + str(master_port), ] env = os.environ.copy() @@ -271,6 +284,7 @@ def get_log(self) -> str: # Random data generation # --------------------------------------------------------------------------- + def generate_random_sft_data( num_samples: int, seq_len: int = 64, @@ -340,6 +354,7 @@ def extract_loss(fwd_bwd_result) -> float: # Training helpers # --------------------------------------------------------------------------- + def run_sft_steps(training_client, data, num_steps=5, lr=1e-3) -> list: """Run SFT training steps and return loss history.""" import xorl_client @@ -376,10 +391,7 @@ def _start_server_or_fail(server, timeout=180.0): healthy = server.start(timeout=timeout) if not healthy: log_tail = "\n".join(server.get_log().splitlines()[-50:]) - pytest.fail( - f"Server failed to become healthy within {timeout}s.\n" - f"--- Log (last 50 lines) ---\n{log_tail}" - ) + pytest.fail(f"Server failed to become healthy within {timeout}s.\n--- Log (last 50 lines) ---\n{log_tail}") def _create_lora_client(base_url, model_dir, model_id="test", rank=8): @@ -415,4 +427,3 @@ def _create_full_weight_client(base_url, model_dir): # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- - diff --git a/tests/models/test_moe_experts_lora.py b/tests/models/test_moe_experts_lora.py index 0355242f..a84d17fb 100644 --- a/tests/models/test_moe_experts_lora.py +++ b/tests/models/test_moe_experts_lora.py @@ -2,20 +2,22 @@ import pytest + pytestmark = [pytest.mark.cpu, pytest.mark.gpu] import torch import torch.nn as nn +from xorl.lora.mapping import can_apply_lora, get_lora_class_for_module +from xorl.models.layers.moe import MOE_EXPERT_BACKENDS, MoEBlock, MoEExperts, MoEExpertsLoRA, MoELoRAConfig from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import ( - Qwen3MoeTritonExperts, Qwen3MoeSparseExperts, + Qwen3MoeTritonExperts, ) -from xorl.models.layers.moe import MoEExperts, MoEExpertsLoRA, MoELoRAConfig, MoEBlock, MOE_EXPERT_BACKENDS -from xorl.lora.mapping import can_apply_lora, get_lora_class_for_module class MockConfig: """Mock config for testing.""" + def __init__( self, num_experts=4, @@ -196,8 +198,11 @@ def test_eager_forward_backward_and_moe_block(self): # MoEBlock end-to-end block = MoEBlock( - hidden_size=32, num_experts=4, top_k=2, - intermediate_size=64, moe_implementation="eager", + hidden_size=32, + num_experts=4, + top_k=2, + intermediate_size=64, + moe_implementation="eager", ) nn.init.xavier_normal_(block.experts.gate_proj.data) nn.init.xavier_normal_(block.experts.up_proj.data) @@ -261,8 +266,11 @@ def test_native_via_moe_block(self): """Test native LoRA works end-to-end through MoEBlock.""" device = "cuda" block = MoEBlock( - hidden_size=32, num_experts=4, top_k=2, - intermediate_size=64, moe_implementation="native", + hidden_size=32, + num_experts=4, + top_k=2, + intermediate_size=64, + moe_implementation="native", ) nn.init.xavier_normal_(block.experts.gate_proj.data) nn.init.xavier_normal_(block.experts.up_proj.data) @@ -288,7 +296,14 @@ def test_native_via_moe_block(self): def _make_lora_block( - backend, num_experts, hidden_dim, intermediate, r, lora_alpha, device, dtype, + backend, + num_experts, + hidden_dim, + intermediate, + r, + lora_alpha, + device, + dtype, ): """Create a MoEBlock with LoRA on the given backend, with deterministic init.""" block = MoEBlock( @@ -342,12 +357,24 @@ class TestCrossBackendConsistency: def _make_pair(self, ref_backend, test_backend, device): """Create a reference and test block with identical weights.""" ref = _make_lora_block( - ref_backend, self.NUM_EXPERTS, self.HIDDEN_DIM, self.INTERMEDIATE, - self.R, self.LORA_ALPHA, device, self.DTYPE, + ref_backend, + self.NUM_EXPERTS, + self.HIDDEN_DIM, + self.INTERMEDIATE, + self.R, + self.LORA_ALPHA, + device, + self.DTYPE, ) test = _make_lora_block( - test_backend, self.NUM_EXPERTS, self.HIDDEN_DIM, self.INTERMEDIATE, - self.R, self.LORA_ALPHA, device, self.DTYPE, + test_backend, + self.NUM_EXPERTS, + self.HIDDEN_DIM, + self.INTERMEDIATE, + self.R, + self.LORA_ALPHA, + device, + self.DTYPE, ) _copy_block_weights(ref, test) return ref, test @@ -358,8 +385,11 @@ def test_zero_lora_matches_base(self, backend): device = "cuda" # Base block (no LoRA) base_block = MoEBlock( - hidden_size=self.HIDDEN_DIM, num_experts=self.NUM_EXPERTS, top_k=2, - intermediate_size=self.INTERMEDIATE, moe_implementation=backend, + hidden_size=self.HIDDEN_DIM, + num_experts=self.NUM_EXPERTS, + top_k=2, + intermediate_size=self.INTERMEDIATE, + moe_implementation=backend, ) torch.manual_seed(42) nn.init.xavier_normal_(base_block.experts.gate_proj.data) @@ -370,8 +400,11 @@ def test_zero_lora_matches_base(self, backend): # LoRA block with lora_B = 0 lora_block = MoEBlock( - hidden_size=self.HIDDEN_DIM, num_experts=self.NUM_EXPERTS, top_k=2, - intermediate_size=self.INTERMEDIATE, moe_implementation=backend, + hidden_size=self.HIDDEN_DIM, + num_experts=self.NUM_EXPERTS, + top_k=2, + intermediate_size=self.INTERMEDIATE, + moe_implementation=backend, ) torch.manual_seed(42) nn.init.xavier_normal_(lora_block.experts.gate_proj.data) @@ -388,15 +421,21 @@ def test_zero_lora_matches_base(self, backend): lora_out, _ = lora_block(hidden) torch.testing.assert_close( - lora_out, base_out, atol=1e-3, rtol=1e-2, + lora_out, + base_out, + atol=1e-3, + rtol=1e-2, msg=f"[{backend}] Zero-LoRA output should match base model", ) - @pytest.mark.parametrize("ref_backend,test_backend", [ - ("eager", "native"), - ("eager", "triton"), - ("triton", "native"), - ]) + @pytest.mark.parametrize( + "ref_backend,test_backend", + [ + ("eager", "native"), + ("eager", "triton"), + ("triton", "native"), + ], + ) def test_cross_backend_output_and_gradients(self, ref_backend, test_backend): """Cross-backend outputs and LoRA gradients should match.""" ref, test = self._make_pair(ref_backend, test_backend, "cuda") @@ -413,7 +452,10 @@ def test_cross_backend_output_and_gradients(self, ref_backend, test_backend): # Output agreement torch.testing.assert_close( - test_out, ref_out, atol=0.05, rtol=0.02, + test_out, + ref_out, + atol=0.05, + rtol=0.02, msg=f"{ref_backend} vs {test_backend} output mismatch", ) @@ -428,11 +470,17 @@ def test_cross_backend_output_and_gradients(self, ref_backend, test_backend): assert test_grad_A is not None, f"test {proj}_lora_A grad is None" torch.testing.assert_close( - test_grad_A, ref_grad_A, atol=0.05, rtol=0.05, + test_grad_A, + ref_grad_A, + atol=0.05, + rtol=0.05, msg=f"Gradient mismatch: {proj}_lora_A ({ref_backend} vs {test_backend})", ) torch.testing.assert_close( - test_grad_B, ref_grad_B, atol=0.05, rtol=0.05, + test_grad_B, + ref_grad_B, + atol=0.05, + rtol=0.05, msg=f"Gradient mismatch: {proj}_lora_B ({ref_backend} vs {test_backend})", ) @@ -441,8 +489,14 @@ def test_nonzero_lora_changes_output(self, backend): """Non-zero LoRA weights must produce a different output from base.""" device = "cuda" block = _make_lora_block( - backend, self.NUM_EXPERTS, self.HIDDEN_DIM, self.INTERMEDIATE, - self.R, self.LORA_ALPHA, device, self.DTYPE, + backend, + self.NUM_EXPERTS, + self.HIDDEN_DIM, + self.INTERMEDIATE, + self.R, + self.LORA_ALPHA, + device, + self.DTYPE, ) # Set lora_B to non-zero with torch.no_grad(): @@ -451,10 +505,17 @@ def test_nonzero_lora_changes_output(self, backend): nn.init.xavier_normal_(lora_B) # Base block (no LoRA) with same base weights - base_block = MoEBlock( - hidden_size=self.HIDDEN_DIM, num_experts=self.NUM_EXPERTS, top_k=2, - intermediate_size=self.INTERMEDIATE, moe_implementation=backend, - ).to(device).to(self.DTYPE) + base_block = ( + MoEBlock( + hidden_size=self.HIDDEN_DIM, + num_experts=self.NUM_EXPERTS, + top_k=2, + intermediate_size=self.INTERMEDIATE, + moe_implementation=backend, + ) + .to(device) + .to(self.DTYPE) + ) with torch.no_grad(): base_block.gate.weight.copy_(block.gate.weight) base_block.experts.gate_proj.copy_(block.experts.gate_proj) @@ -467,9 +528,7 @@ def test_nonzero_lora_changes_output(self, backend): lora_out, _ = block(hidden) diff = (lora_out - base_out).abs().max().item() - assert diff > 1e-3, ( - f"[{backend}] Non-zero LoRA should change the output, but max diff={diff}" - ) + assert diff > 1e-3, f"[{backend}] Non-zero LoRA should change the output, but max diff={diff}" # --------------------------------------------------------------------------- @@ -549,13 +608,14 @@ def test_from_module_with_qwen3_subclass(self): def test_injection_error_handling(self): """Test error cases and valid injection for inject_lora_into_model.""" - from xorl.lora import inject_lora_into_model, LoraLinear + from xorl.lora import LoraLinear, inject_lora_into_model # No matching modules class ModelA(nn.Module): def __init__(self): super().__init__() self.layer1 = nn.Linear(64, 64) + with pytest.raises(ValueError, match="No modules found matching target_modules"): inject_lora_into_model(ModelA(), r=8, lora_alpha=16, target_modules=["nonexistent_proj"]) @@ -564,6 +624,7 @@ class UnsupportedModule(nn.Module): def __init__(self): super().__init__() self.weight = nn.Parameter(torch.randn(64, 64)) + def forward(self, x): return x @ self.weight @@ -571,6 +632,7 @@ class ModelB(nn.Module): def __init__(self): super().__init__() self.custom_layer = UnsupportedModule() + with pytest.raises(ValueError, match="No modules found matching target_modules"): inject_lora_into_model(ModelB(), r=8, lora_alpha=16, target_modules=["custom_layer"]) @@ -580,6 +642,7 @@ def __init__(self): super().__init__() self.q_proj = nn.Linear(64, 64) self.v_proj = nn.Linear(64, 64) + model = ModelC() inject_lora_into_model(model, r=8, lora_alpha=16, target_modules=["q_proj", "v_proj"]) assert isinstance(model.q_proj, LoraLinear) diff --git a/tests/models/test_moe_routing_cache.py b/tests/models/test_moe_routing_cache.py index 7420663f..4c7085fe 100644 --- a/tests/models/test_moe_routing_cache.py +++ b/tests/models/test_moe_routing_cache.py @@ -10,16 +10,16 @@ import pytest import torch import torch.nn as nn -from functools import partial from torch.utils.checkpoint import checkpoint + pytestmark = [pytest.mark.gpu] try: from xorl.models.layers.moe.moe_block import ( MoEBlock, - moe_routing_context_fn, _routing_cache_mode, + moe_routing_context_fn, ) from xorl.models.layers.moe.router import TopKRouter except ImportError: @@ -37,7 +37,8 @@ def _checkpoint_with_routing(fn, *args, **kwargs): """Checkpoint wrapper that uses moe_routing_context_fn.""" return checkpoint( - fn, *args, + fn, + *args, use_reentrant=False, context_fn=moe_routing_context_fn, **kwargs, @@ -87,6 +88,7 @@ class TestContextFnMechanism: def test_forward_sets_mode(self): """During checkpoint forward, mode is 'forward'.""" import xorl.models.layers.moe.moe_block as mb + observed = [] class _Observer(nn.Module): @@ -110,6 +112,7 @@ def forward(self, x): def test_mode_none_without_context_fn(self): """Without context_fn, mode stays None — no caching.""" import xorl.models.layers.moe.moe_block as mb + observed = [] class _Observer(nn.Module): @@ -207,10 +210,7 @@ def test_pp_four_microbatches_1f1b(self): layer.train() moe = layer.mlp - xs = [ - torch.randn(1, 8 + i, 64, device="cuda", requires_grad=True) - for i in range(4) - ] + xs = [torch.randn(1, 8 + i, 64, device="cuda", requires_grad=True) for i in range(4)] outs = [] # Warmup: 2 forwards @@ -242,9 +242,7 @@ def test_multi_layer_pp(self): class _Model(nn.Module): def __init__(self, num_layers=3): super().__init__() - self.layers = nn.ModuleList( - [_SimpleDecoderLayer(hidden_size) for _ in range(num_layers)] - ) + self.layers = nn.ModuleList([_SimpleDecoderLayer(hidden_size) for _ in range(num_layers)]) def forward(self, x): for layer in self.layers: @@ -262,23 +260,17 @@ def forward(self, x): out2 = model(x2) for i, layer in enumerate(model.layers): - assert len(layer.mlp._routing_cache) == 2, ( - f"Layer {i}: expected 2 cached entries" - ) + assert len(layer.mlp._routing_cache) == 2, f"Layer {i}: expected 2 cached entries" # Backward MB1 out1.sum().backward() for i, layer in enumerate(model.layers): - assert len(layer.mlp._routing_cache) == 1, ( - f"Layer {i}: expected 1 after MB1 backward" - ) + assert len(layer.mlp._routing_cache) == 1, f"Layer {i}: expected 1 after MB1 backward" # Backward MB2 out2.sum().backward() for i, layer in enumerate(model.layers): - assert len(layer.mlp._routing_cache) == 0, ( - f"Layer {i}: expected 0 after MB2 backward" - ) + assert len(layer.mlp._routing_cache) == 0, f"Layer {i}: expected 0 after MB2 backward" def test_nondeterministic_attention_routing_preserved(self): """With non-deterministic attention, cached routing is replayed on recompute.""" @@ -387,7 +379,6 @@ class TestBaseModelIntegration: def test_context_fn_injected_for_moe(self): """gradient_checkpointing_enable adds context_fn when MoE blocks present.""" - import xorl.models.layers.moe.moe_block as mb class _FakeConfig: pass diff --git a/tests/models/test_moe_routing_replay.py b/tests/models/test_moe_routing_replay.py index 21bb9445..0b8ba6ee 100644 --- a/tests/models/test_moe_routing_replay.py +++ b/tests/models/test_moe_routing_replay.py @@ -10,18 +10,17 @@ import pytest import torch import torch.nn as nn -from functools import partial from torch.utils.checkpoint import checkpoint + pytestmark = [pytest.mark.gpu] +from xorl.models.layers.moe.moe_block import MoEBlock from xorl.models.layers.moe.routing_replay import ( RoutingReplay, get_replay_stage, set_replay_stage, ) -from xorl.models.layers.moe.moe_block import MoEBlock -from xorl.models.layers.moe.router import TopKRouter # --------------------------------------------------------------------------- @@ -212,8 +211,11 @@ def test_record_replay_no_replay_and_router_training(self): router training/detach, and regather correctness.""" # --- Record stores expert indices correctly --- moe = MoEBlock( - hidden_size=64, num_experts=4, top_k=2, - intermediate_size=128, moe_implementation="eager", + hidden_size=64, + num_experts=4, + top_k=2, + intermediate_size=128, + moe_implementation="eager", ).cuda() moe.train() replay = RoutingReplay() @@ -249,8 +251,11 @@ def test_record_replay_no_replay_and_router_training(self): # --- No-recording conditions --- moe2 = MoEBlock( - hidden_size=64, num_experts=4, top_k=2, - intermediate_size=128, moe_implementation="eager", + hidden_size=64, + num_experts=4, + top_k=2, + intermediate_size=128, + moe_implementation="eager", ).cuda() moe2.train() replay4 = RoutingReplay() @@ -263,8 +268,11 @@ def test_record_replay_no_replay_and_router_training(self): # No _routing_replay => stage doesn't matter moe3 = MoEBlock( - hidden_size=64, num_experts=4, top_k=2, - intermediate_size=128, moe_implementation="eager", + hidden_size=64, + num_experts=4, + top_k=2, + intermediate_size=128, + moe_implementation="eager", ).cuda() moe3.train() assert moe3._routing_replay is None @@ -273,8 +281,11 @@ def test_record_replay_no_replay_and_router_training(self): # Eval mode with stage=None moe4 = MoEBlock( - hidden_size=64, num_experts=4, top_k=2, - intermediate_size=128, moe_implementation="eager", + hidden_size=64, + num_experts=4, + top_k=2, + intermediate_size=128, + moe_implementation="eager", ).cuda() moe4.eval() replay5 = RoutingReplay() @@ -286,8 +297,11 @@ def test_record_replay_no_replay_and_router_training(self): # --- Router training and regather --- moe5 = MoEBlock( - hidden_size=64, num_experts=4, top_k=2, - intermediate_size=128, moe_implementation="eager", + hidden_size=64, + num_experts=4, + top_k=2, + intermediate_size=128, + moe_implementation="eager", norm_topk_prob=True, ).cuda() moe5.train() @@ -296,9 +310,7 @@ def test_record_replay_no_replay_and_router_training(self): hidden_states = x5.view(-1, 64) router_logits = moe5.gate(hidden_states) orig_weights, orig_experts = moe5.router(router_logits, x5.dtype) - regathered_experts, regathered_weights = moe5._regather_routing( - router_logits, orig_experts, x5.dtype - ) + regathered_experts, regathered_weights = moe5._regather_routing(router_logits, orig_experts, x5.dtype) assert torch.equal(regathered_experts, orig_experts) assert torch.allclose(regathered_weights, orig_weights, atol=1e-6) @@ -318,8 +330,11 @@ def test_record_replay_no_replay_and_router_training(self): # train_router=False => no gradient from expert path moe6 = MoEBlock( - hidden_size=64, num_experts=4, top_k=2, - intermediate_size=128, moe_implementation="eager", + hidden_size=64, + num_experts=4, + top_k=2, + intermediate_size=128, + moe_implementation="eager", train_router=False, ).cuda() moe6.train() @@ -346,9 +361,7 @@ def _make_model(self, num_layers=3, hidden_size=64): class _Model(nn.Module): def __init__(self): super().__init__() - self.layers = nn.ModuleList( - [_SimpleDecoderLayer(hidden_size) for _ in range(num_layers)] - ) + self.layers = nn.ModuleList([_SimpleDecoderLayer(hidden_size) for _ in range(num_layers)]) def forward(self, x): for layer in self.layers: @@ -498,8 +511,11 @@ class _FakeConfig: # --- R3 Preload --- moe = MoEBlock( - hidden_size=64, num_experts=4, top_k=2, - intermediate_size=128, moe_implementation="eager", + hidden_size=64, + num_experts=4, + top_k=2, + intermediate_size=128, + moe_implementation="eager", ).cuda() moe.eval() diff --git a/tests/models/test_moe_weight_auto_merge.py b/tests/models/test_moe_weight_auto_merge.py index 9375f61d..6986832d 100644 --- a/tests/models/test_moe_weight_auto_merge.py +++ b/tests/models/test_moe_weight_auto_merge.py @@ -17,6 +17,7 @@ import pytest import torch + pytestmark = [pytest.mark.cpu] @@ -24,13 +25,9 @@ # Copy of the implementation for testing (avoids heavy import dependencies) # ============================================================================= -_EXPERT_KEY_PATTERN = re.compile( - r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.weight$" -) +_EXPERT_KEY_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.weight$") -_FUSED_EXPERT_PATTERN = re.compile( - r"^model\.layers\.\d+\.mlp\.experts\.(gate|up|down)_proj$" -) +_FUSED_EXPERT_PATTERN = re.compile(r"^model\.layers\.\d+\.mlp\.experts\.(gate|up|down)_proj$") def parse_expert_key(key: str) -> Optional[Tuple[int, int, str]]: @@ -57,9 +54,7 @@ def add(self, layer_idx: int, expert_idx: int, proj: str, tensor: torch.Tensor) key = (layer_idx, proj) if key not in self._stacked_buffers: stacked_shape = (self.num_experts,) + tensor.shape - self._stacked_buffers[key] = torch.empty( - stacked_shape, dtype=tensor.dtype, device="cpu" - ) + self._stacked_buffers[key] = torch.empty(stacked_shape, dtype=tensor.dtype, device="cpu") self._stacked_buffers[key][expert_idx].copy_(tensor) self._filled_experts[key].add(expert_idx) @@ -74,8 +69,7 @@ def pop_stacked(self, layer_idx: int, proj: str) -> torch.Tensor: filled = self._filled_experts.pop(key) if len(filled) != self.num_experts: raise ValueError( - f"Incomplete experts for layer {layer_idx}, {proj}_proj: " - f"got {len(filled)}, expected {self.num_experts}" + f"Incomplete experts for layer {layer_idx}, {proj}_proj: got {len(filled)}, expected {self.num_experts}" ) return self._stacked_buffers.pop(key) @@ -127,20 +121,30 @@ def test_parse_expert_key_and_merging_detection(self): assert parse_expert_key("model.layers.0.mlp.experts.5.gate_proj") is None # Fused format model needs merging - assert _model_needs_expert_merging({ - "model.layers.0.mlp.experts.gate_proj", - "model.layers.0.mlp.experts.up_proj", - "model.layers.0.mlp.experts.down_proj", - "model.layers.0.self_attn.q_proj.weight", - }) is True + assert ( + _model_needs_expert_merging( + { + "model.layers.0.mlp.experts.gate_proj", + "model.layers.0.mlp.experts.up_proj", + "model.layers.0.mlp.experts.down_proj", + "model.layers.0.self_attn.q_proj.weight", + } + ) + is True + ) # Non-MoE model does not - assert _model_needs_expert_merging({ - "model.layers.0.mlp.gate_proj.weight", - "model.layers.0.mlp.up_proj.weight", - "model.layers.0.mlp.down_proj.weight", - "model.layers.0.self_attn.q_proj.weight", - }) is False + assert ( + _model_needs_expert_merging( + { + "model.layers.0.mlp.gate_proj.weight", + "model.layers.0.mlp.up_proj.weight", + "model.layers.0.mlp.down_proj.weight", + "model.layers.0.self_attn.q_proj.weight", + } + ) + is False + ) # Empty assert _model_needs_expert_merging(set()) is False @@ -268,6 +272,7 @@ def test_buffer_lifecycle_layers_errors_and_edge_cases(self): # Tests for checkpoint format detection and loading logic # ============================================================================= + def _checkpoint_has_per_expert_weights(checkpoint_keys): for key in checkpoint_keys: if _EXPERT_KEY_PATTERN.match(key): @@ -281,32 +286,52 @@ class TestCheckpointFormatAndLoading: def test_checkpoint_format_detection_and_loading(self): """Test detection of per-expert/fused/non-MoE/empty/mixed formats, and loading both.""" # Per-expert format - assert _checkpoint_has_per_expert_weights({ - "model.layers.0.mlp.experts.0.gate_proj.weight", - "model.layers.0.mlp.experts.1.gate_proj.weight", - "model.layers.0.self_attn.q_proj.weight", - }) is True + assert ( + _checkpoint_has_per_expert_weights( + { + "model.layers.0.mlp.experts.0.gate_proj.weight", + "model.layers.0.mlp.experts.1.gate_proj.weight", + "model.layers.0.self_attn.q_proj.weight", + } + ) + is True + ) # Fused format - assert _checkpoint_has_per_expert_weights({ - "model.layers.0.mlp.experts.gate_proj", - "model.layers.0.mlp.experts.up_proj", - }) is False + assert ( + _checkpoint_has_per_expert_weights( + { + "model.layers.0.mlp.experts.gate_proj", + "model.layers.0.mlp.experts.up_proj", + } + ) + is False + ) # Non-MoE - assert _checkpoint_has_per_expert_weights({ - "model.layers.0.mlp.gate_proj.weight", - "model.layers.0.self_attn.q_proj.weight", - }) is False + assert ( + _checkpoint_has_per_expert_weights( + { + "model.layers.0.mlp.gate_proj.weight", + "model.layers.0.self_attn.q_proj.weight", + } + ) + is False + ) # Empty assert _checkpoint_has_per_expert_weights(set()) is False # Mixed - assert _checkpoint_has_per_expert_weights({ - "model.layers.0.mlp.experts.gate_proj", - "model.layers.1.mlp.experts.0.gate_proj.weight", - }) is True + assert ( + _checkpoint_has_per_expert_weights( + { + "model.layers.0.mlp.experts.gate_proj", + "model.layers.1.mlp.experts.0.gate_proj.weight", + } + ) + is True + ) # --- Loading both formats --- model_params = { diff --git a/tests/models/test_moe_weight_loading_integration.py b/tests/models/test_moe_weight_loading_integration.py index 5bd448d7..a2328943 100644 --- a/tests/models/test_moe_weight_loading_integration.py +++ b/tests/models/test_moe_weight_loading_integration.py @@ -5,33 +5,28 @@ into a model that expects fused (stacked) expert format. """ -import os -import tempfile from typing import Dict, Iterator, Tuple import pytest import torch import torch.nn as nn + pytestmark = [pytest.mark.cpu] # Re-implement the core logic to test without full xorl dependencies import re from collections import defaultdict -from typing import Dict, Optional, Set, Tuple +from typing import Optional, Set # ============================================================================= # Copy of core implementation for testing # ============================================================================= -_EXPERT_KEY_PATTERN = re.compile( - r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.weight$" -) +_EXPERT_KEY_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.weight$") -_FUSED_EXPERT_PATTERN = re.compile( - r"^model\.layers\.\d+\.mlp\.experts\.(gate|up|down)_proj$" -) +_FUSED_EXPERT_PATTERN = re.compile(r"^model\.layers\.\d+\.mlp\.experts\.(gate|up|down)_proj$") def parse_expert_key(key: str) -> Optional[Tuple[int, int, str]]: @@ -62,9 +57,7 @@ def add(self, layer_idx: int, expert_idx: int, proj: str, tensor: torch.Tensor) tensor = tensor.t().contiguous() if key not in self._stacked_buffers: stacked_shape = (self.num_experts,) + tensor.shape - self._stacked_buffers[key] = torch.empty( - stacked_shape, dtype=tensor.dtype, device="cpu" - ) + self._stacked_buffers[key] = torch.empty(stacked_shape, dtype=tensor.dtype, device="cpu") # Copy directly into the slice (streaming) self._stacked_buffers[key][expert_idx].copy_(tensor) self._filled_experts[key].add(expert_idx) @@ -80,8 +73,7 @@ def pop_stacked(self, layer_idx: int, proj: str) -> torch.Tensor: filled = self._filled_experts.pop(key) if len(filled) != self.num_experts: raise ValueError( - f"Incomplete experts for layer {layer_idx}, {proj}_proj: " - f"got {len(filled)}, expected {self.num_experts}" + f"Incomplete experts for layer {layer_idx}, {proj}_proj: got {len(filled)}, expected {self.num_experts}" ) return self._stacked_buffers.pop(key) @@ -105,15 +97,9 @@ def __init__(self, num_experts: int, hidden_size: int, intermediate_size: int): super().__init__() self.num_experts = num_experts # Fused format (G, K, N): [num_experts, in_features, out_features] - self.gate_proj = nn.Parameter( - torch.empty(num_experts, hidden_size, intermediate_size) - ) - self.up_proj = nn.Parameter( - torch.empty(num_experts, hidden_size, intermediate_size) - ) - self.down_proj = nn.Parameter( - torch.empty(num_experts, intermediate_size, hidden_size) - ) + self.gate_proj = nn.Parameter(torch.empty(num_experts, hidden_size, intermediate_size)) + self.up_proj = nn.Parameter(torch.empty(num_experts, hidden_size, intermediate_size)) + self.down_proj = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size)) class MockMoeLayer(nn.Module): @@ -129,12 +115,12 @@ class MockModelInner(nn.Module): def __init__(self, num_layers: int, num_experts: int, hidden_size: int, intermediate_size: int): super().__init__() - self.layers = nn.ModuleList([ - nn.ModuleDict({ - "mlp": MockMoeLayer(num_experts, hidden_size, intermediate_size) - }) - for _ in range(num_layers) - ]) + self.layers = nn.ModuleList( + [ + nn.ModuleDict({"mlp": MockMoeLayer(num_experts, hidden_size, intermediate_size)}) + for _ in range(num_layers) + ] + ) class MockMoeModel(nn.Module): @@ -174,13 +160,16 @@ def create_per_expert_state_dict( for layer_idx in range(num_layers): for expert_idx in range(num_experts): # gate_proj and up_proj: [intermediate_size, hidden_size] - state_dict[f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.gate_proj.weight"] = \ - torch.randn(intermediate_size, hidden_size) - state_dict[f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.up_proj.weight"] = \ - torch.randn(intermediate_size, hidden_size) + state_dict[f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.gate_proj.weight"] = torch.randn( + intermediate_size, hidden_size + ) + state_dict[f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.up_proj.weight"] = torch.randn( + intermediate_size, hidden_size + ) # down_proj: [hidden_size, intermediate_size] - state_dict[f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight"] = \ - torch.randn(hidden_size, intermediate_size) + state_dict[f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight"] = torch.randn( + hidden_size, intermediate_size + ) return state_dict @@ -250,13 +239,9 @@ def test_load_and_verify_shapes_and_order(self): intermediate_size = 128 model = MockMoeModel(num_layers, num_experts, hidden_size, intermediate_size) - per_expert_state_dict = create_per_expert_state_dict( - num_layers, num_experts, hidden_size, intermediate_size - ) + per_expert_state_dict = create_per_expert_state_dict(num_layers, num_experts, hidden_size, intermediate_size) - loaded_weights = simulate_load_model_weights( - model, iter(per_expert_state_dict.items()), num_experts - ) + loaded_weights = simulate_load_model_weights(model, iter(per_expert_state_dict.items()), num_experts) # All fused parameters created with correct shapes for layer_idx in range(num_layers): @@ -278,14 +263,14 @@ def test_load_and_verify_shapes_and_order(self): model2 = MockMoeModel(1, num_experts, 8, 16) state_dict2 = {} for expert_idx in range(num_experts): - state_dict2[f"model.layers.0.mlp.experts.{expert_idx}.gate_proj.weight"] = \ - torch.full((16, 8), float(expert_idx)) - state_dict2[f"model.layers.0.mlp.experts.{expert_idx}.up_proj.weight"] = \ - torch.randn(16, 8) - state_dict2[f"model.layers.0.mlp.experts.{expert_idx}.down_proj.weight"] = \ - torch.randn(8, 16) + state_dict2[f"model.layers.0.mlp.experts.{expert_idx}.gate_proj.weight"] = torch.full( + (16, 8), float(expert_idx) + ) + state_dict2[f"model.layers.0.mlp.experts.{expert_idx}.up_proj.weight"] = torch.randn(16, 8) + state_dict2[f"model.layers.0.mlp.experts.{expert_idx}.down_proj.weight"] = torch.randn(8, 16) import random + items = list(state_dict2.items()) random.seed(42) random.shuffle(items) @@ -293,8 +278,7 @@ def test_load_and_verify_shapes_and_order(self): loaded2 = simulate_load_model_weights(model2, iter(items), num_experts) gate_proj = loaded2["model.layers.0.mlp.experts.gate_proj"] for expert_idx in range(num_experts): - assert torch.all(gate_proj[expert_idx] == float(expert_idx)), \ - f"Expert {expert_idx} not in correct position" + assert torch.all(gate_proj[expert_idx] == float(expert_idx)), f"Expert {expert_idx} not in correct position" def test_streaming_and_edge_cases(self): """Test sharded streaming load, large expert count, single expert, and random order.""" @@ -305,9 +289,7 @@ def test_streaming_and_edge_cases(self): intermediate_size = 32 model = MockMoeModel(num_layers, num_experts, hidden_size, intermediate_size) - state_dict = create_per_expert_state_dict( - num_layers, num_experts, hidden_size, intermediate_size - ) + state_dict = create_per_expert_state_dict(num_layers, num_experts, hidden_size, intermediate_size) items = list(state_dict.items()) shard1 = [(k, v) for k, v in items if any(f".experts.{i}." in k for i in range(4))] shard2 = [(k, v) for k, v in items if any(f".experts.{i}." in k for i in range(4, 8))] @@ -321,21 +303,18 @@ def test_streaming_and_edge_cases(self): # Large expert count (128) model_large = MockMoeModel(1, 128, 16, 32) state_dict_large = create_per_expert_state_dict(1, 128, 16, 32) - loaded_large = simulate_load_model_weights( - model_large, iter(state_dict_large.items()), 128 - ) + loaded_large = simulate_load_model_weights(model_large, iter(state_dict_large.items()), 128) assert loaded_large["model.layers.0.mlp.experts.gate_proj"].shape == (128, 16, 32) # Single expert model_single = MockMoeModel(1, 1, 8, 16) state_dict_single = create_per_expert_state_dict(1, 1, 8, 16) - loaded_single = simulate_load_model_weights( - model_single, iter(state_dict_single.items()), 1 - ) + loaded_single = simulate_load_model_weights(model_single, iter(state_dict_single.items()), 1) assert loaded_single["model.layers.0.mlp.experts.gate_proj"].shape == (1, 8, 16) # Random order loading import random + model_rand = MockMoeModel(2, 4, 8, 16) state_dict_rand = create_per_expert_state_dict(2, 4, 8, 16) items_rand = list(state_dict_rand.items()) @@ -352,12 +331,9 @@ def test_incomplete_experts_raises_error(self): state_dict = {} for expert_idx in range(3): # Missing expert 3 - state_dict[f"model.layers.0.mlp.experts.{expert_idx}.gate_proj.weight"] = \ - torch.randn(16, 8) - state_dict[f"model.layers.0.mlp.experts.{expert_idx}.up_proj.weight"] = \ - torch.randn(16, 8) - state_dict[f"model.layers.0.mlp.experts.{expert_idx}.down_proj.weight"] = \ - torch.randn(8, 16) + state_dict[f"model.layers.0.mlp.experts.{expert_idx}.gate_proj.weight"] = torch.randn(16, 8) + state_dict[f"model.layers.0.mlp.experts.{expert_idx}.up_proj.weight"] = torch.randn(16, 8) + state_dict[f"model.layers.0.mlp.experts.{expert_idx}.down_proj.weight"] = torch.randn(8, 16) with pytest.raises(RuntimeError, match="Incomplete expert weights"): simulate_load_model_weights(model, iter(state_dict.items()), 4) diff --git a/tests/models/test_qwen3_moe_fused_lora.py b/tests/models/test_qwen3_moe_fused_lora.py index 1becdb78..53c9c3e3 100644 --- a/tests/models/test_qwen3_moe_fused_lora.py +++ b/tests/models/test_qwen3_moe_fused_lora.py @@ -4,16 +4,17 @@ import torch import torch.nn as nn +from xorl.lora.mapping import can_apply_lora, get_lora_class_for_module +from xorl.models.layers.moe import MOE_EXPERT_BACKENDS, MoEBlock, MoEExperts, MoEExpertsLoRA, MoELoRAConfig from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import ( - Qwen3MoeTritonExperts, Qwen3MoeSparseExperts, + Qwen3MoeTritonExperts, ) -from xorl.models.layers.moe import MoEExperts, MoEExpertsLoRA, MoELoRAConfig, MoEBlock, MOE_EXPERT_BACKENDS -from xorl.lora.mapping import can_apply_lora, get_lora_class_for_module class MockConfig: """Mock config for testing.""" + def __init__( self, num_experts=4, @@ -178,8 +179,11 @@ def test_eager_forward_backward_and_moe_block(self): # MoEBlock end-to-end block = MoEBlock( - hidden_size=32, num_experts=4, top_k=2, - intermediate_size=64, moe_implementation="eager", + hidden_size=32, + num_experts=4, + top_k=2, + intermediate_size=64, + moe_implementation="eager", ) nn.init.xavier_normal_(block.experts.gate_proj.data) nn.init.xavier_normal_(block.experts.up_proj.data) @@ -242,8 +246,11 @@ def test_forward_backward_and_moe_block(self, backend): # MoEBlock end-to-end (test only native to avoid duplicate) if backend == "native": block = MoEBlock( - hidden_size=32, num_experts=4, top_k=2, - intermediate_size=64, moe_implementation="native", + hidden_size=32, + num_experts=4, + top_k=2, + intermediate_size=64, + moe_implementation="native", ) nn.init.xavier_normal_(block.experts.gate_proj.data) nn.init.xavier_normal_(block.experts.up_proj.data) @@ -268,7 +275,14 @@ def test_forward_backward_and_moe_block(self, backend): def _make_lora_block( - backend, num_experts, hidden_dim, intermediate, r, lora_alpha, device, dtype, + backend, + num_experts, + hidden_dim, + intermediate, + r, + lora_alpha, + device, + dtype, ): """Create a MoEBlock with LoRA on the given backend, with deterministic init.""" block = MoEBlock( @@ -317,12 +331,24 @@ class TestCrossBackendConsistency: def _make_pair(self, ref_backend, test_backend, device): ref = _make_lora_block( - ref_backend, self.NUM_EXPERTS, self.HIDDEN_DIM, self.INTERMEDIATE, - self.R, self.LORA_ALPHA, device, self.DTYPE, + ref_backend, + self.NUM_EXPERTS, + self.HIDDEN_DIM, + self.INTERMEDIATE, + self.R, + self.LORA_ALPHA, + device, + self.DTYPE, ) test = _make_lora_block( - test_backend, self.NUM_EXPERTS, self.HIDDEN_DIM, self.INTERMEDIATE, - self.R, self.LORA_ALPHA, device, self.DTYPE, + test_backend, + self.NUM_EXPERTS, + self.HIDDEN_DIM, + self.INTERMEDIATE, + self.R, + self.LORA_ALPHA, + device, + self.DTYPE, ) _copy_block_weights(ref, test) return ref, test @@ -333,8 +359,11 @@ def test_zero_lora_and_nonzero_lora(self): for backend in ["eager", "triton", "native"]: # --- zero LoRA matches base --- base_block = MoEBlock( - hidden_size=self.HIDDEN_DIM, num_experts=self.NUM_EXPERTS, top_k=2, - intermediate_size=self.INTERMEDIATE, moe_implementation=backend, + hidden_size=self.HIDDEN_DIM, + num_experts=self.NUM_EXPERTS, + top_k=2, + intermediate_size=self.INTERMEDIATE, + moe_implementation=backend, ) torch.manual_seed(42) nn.init.xavier_normal_(base_block.experts.gate_proj.data) @@ -344,8 +373,11 @@ def test_zero_lora_and_nonzero_lora(self): base_block = base_block.to(device).to(self.DTYPE) lora_block = MoEBlock( - hidden_size=self.HIDDEN_DIM, num_experts=self.NUM_EXPERTS, top_k=2, - intermediate_size=self.INTERMEDIATE, moe_implementation=backend, + hidden_size=self.HIDDEN_DIM, + num_experts=self.NUM_EXPERTS, + top_k=2, + intermediate_size=self.INTERMEDIATE, + moe_implementation=backend, ) torch.manual_seed(42) nn.init.xavier_normal_(lora_block.experts.gate_proj.data) @@ -360,24 +392,40 @@ def test_zero_lora_and_nonzero_lora(self): base_out, _ = base_block(hidden) lora_out, _ = lora_block(hidden) torch.testing.assert_close( - lora_out, base_out, atol=1e-3, rtol=1e-2, + lora_out, + base_out, + atol=1e-3, + rtol=1e-2, msg=f"[{backend}] Zero-LoRA output should match base model", ) # --- nonzero LoRA changes output --- block = _make_lora_block( - backend, self.NUM_EXPERTS, self.HIDDEN_DIM, self.INTERMEDIATE, - self.R, self.LORA_ALPHA, device, self.DTYPE, + backend, + self.NUM_EXPERTS, + self.HIDDEN_DIM, + self.INTERMEDIATE, + self.R, + self.LORA_ALPHA, + device, + self.DTYPE, ) with torch.no_grad(): for proj in ["gate_proj", "up_proj", "down_proj"]: lora_B = getattr(block.experts, f"{proj}_lora_B") nn.init.xavier_normal_(lora_B) - base_block2 = MoEBlock( - hidden_size=self.HIDDEN_DIM, num_experts=self.NUM_EXPERTS, top_k=2, - intermediate_size=self.INTERMEDIATE, moe_implementation=backend, - ).to(device).to(self.DTYPE) + base_block2 = ( + MoEBlock( + hidden_size=self.HIDDEN_DIM, + num_experts=self.NUM_EXPERTS, + top_k=2, + intermediate_size=self.INTERMEDIATE, + moe_implementation=backend, + ) + .to(device) + .to(self.DTYPE) + ) with torch.no_grad(): base_block2.gate.weight.copy_(block.gate.weight) base_block2.experts.gate_proj.copy_(block.experts.gate_proj) @@ -390,9 +438,7 @@ def test_zero_lora_and_nonzero_lora(self): lora_out2, _ = block(hidden2) diff = (lora_out2 - base_out2).abs().max().item() - assert diff > 1e-3, ( - f"[{backend}] Non-zero LoRA should change the output, but max diff={diff}" - ) + assert diff > 1e-3, f"[{backend}] Non-zero LoRA should change the output, but max diff={diff}" def test_cross_backend_output_and_gradients(self): """Cross-backend outputs and LoRA gradients should match.""" @@ -409,7 +455,10 @@ def test_cross_backend_output_and_gradients(self): test_out.sum().backward() torch.testing.assert_close( - test_out, ref_out, atol=0.05, rtol=0.02, + test_out, + ref_out, + atol=0.05, + rtol=0.02, msg=f"{ref_backend} vs {test_backend} output mismatch", ) @@ -423,11 +472,17 @@ def test_cross_backend_output_and_gradients(self): assert test_grad_A is not None, f"test {proj}_lora_A grad is None" torch.testing.assert_close( - test_grad_A, ref_grad_A, atol=0.05, rtol=0.05, + test_grad_A, + ref_grad_A, + atol=0.05, + rtol=0.05, msg=f"Gradient mismatch: {proj}_lora_A ({ref_backend} vs {test_backend})", ) torch.testing.assert_close( - test_grad_B, ref_grad_B, atol=0.05, rtol=0.05, + test_grad_B, + ref_grad_B, + atol=0.05, + rtol=0.05, msg=f"Gradient mismatch: {proj}_lora_B ({ref_backend} vs {test_backend})", ) @@ -506,12 +561,13 @@ def __init__(self, config, backend): assert torch.allclose(lora_exp.gate_proj, base.gate_proj) # Error handling - from xorl.lora import inject_lora_into_model, LoraLinear + from xorl.lora import LoraLinear, inject_lora_into_model class ModelA(nn.Module): def __init__(self): super().__init__() self.layer1 = nn.Linear(64, 64) + with pytest.raises(ValueError, match="No modules found matching target_modules"): inject_lora_into_model(ModelA(), r=8, lora_alpha=16, target_modules=["nonexistent_proj"]) @@ -519,6 +575,7 @@ class UnsupportedModule(nn.Module): def __init__(self): super().__init__() self.weight = nn.Parameter(torch.randn(64, 64)) + def forward(self, x): return x @ self.weight @@ -526,6 +583,7 @@ class ModelB(nn.Module): def __init__(self): super().__init__() self.custom_layer = UnsupportedModule() + with pytest.raises(ValueError, match="No modules found matching target_modules"): inject_lora_into_model(ModelB(), r=8, lora_alpha=16, target_modules=["custom_layer"]) @@ -534,6 +592,7 @@ def __init__(self): super().__init__() self.q_proj = nn.Linear(64, 64) self.v_proj = nn.Linear(64, 64) + model = ModelC() inject_lora_into_model(model, r=8, lora_alpha=16, target_modules=["q_proj", "v_proj"]) assert isinstance(model.q_proj, LoraLinear) diff --git a/tests/ops/__init__.py b/tests/ops/__init__.py index b1974871..f6bd0a5a 100644 --- a/tests/ops/__init__.py +++ b/tests/ops/__init__.py @@ -1,3 +1 @@ """Tests for xorl.ops module.""" - - diff --git a/tests/ops/loss/test_drgrpo_loss.py b/tests/ops/loss/test_drgrpo_loss.py index 89109853..e8d0ca98 100644 --- a/tests/ops/loss/test_drgrpo_loss.py +++ b/tests/ops/loss/test_drgrpo_loss.py @@ -9,15 +9,14 @@ import pytest import torch -from xorl.ops.loss import drgrpo_loss_function - from tests.ops.loss.conftest import assert_close +from xorl.ops.loss import drgrpo_loss_function @pytest.fixture def inputs(): """Test inputs for DRGRPO loss. - + Note: xorl's drgrpo_loss_function takes hidden_states and weight instead of logits. We create hidden_states and weight such that hidden_states @ weight.T produces logits with similar scale to torch.randn(B, S, V). @@ -29,9 +28,9 @@ def inputs(): # Create hidden_states and weight matrix # Scale by 1/sqrt(H) so that logits = hidden_states @ weight.T has variance ~1 # (similar to torch.randn logits in torchforge tests) - hidden_states = torch.randn(B, S, H) / (H ** 0.5) + hidden_states = torch.randn(B, S, H) / (H**0.5) weight = torch.randn(V, H) - + # Compute effective logits for reference (should have variance ~1) logits = hidden_states @ weight.T @@ -76,13 +75,13 @@ def inputs(): class TestDRGRPOLoss: """Tests for drgrpo_loss_function. - + Note: test_forward and test_backward contain regression tests with exact expected values. If the implementation changes intentionally, update the expected values by running the tests with pytest -v and recording the actual values: - + pytest tests/ops/loss/test_drgrpo_loss.py -v -k "test_forward or test_backward" - + Then update the assert_close(...) calls with the new values. """ diff --git a/tests/ops/test_attention.py b/tests/ops/test_attention.py index 47614ca5..967f338b 100644 --- a/tests/ops/test_attention.py +++ b/tests/ops/test_attention.py @@ -1,14 +1,17 @@ """Tests for attention backend functions.""" +from unittest.mock import Mock, patch + import pytest import torch -from unittest.mock import Mock, patch -from xorl.models.layers.attention.utils import repeat_kv from xorl.models.layers.attention.backend.eager import eager_attention_forward +from xorl.models.layers.attention.utils import repeat_kv + try: from xorl.models.layers.attention.backend.flash_attention import flash_attention_forward + _FLASH_ATTN_IMPORT_ERROR = None except ImportError as exc: flash_attention_forward = None @@ -72,20 +75,29 @@ def test_flash_attention_api_behavior(self): key = torch.randn(batch, seqlen, num_heads, head_dim) value = torch.randn(batch, seqlen, num_heads, head_dim) - with patch('xorl.models.layers.attention.backend.flash_attention.flash_attn_func') as mock_fa: + with patch("xorl.models.layers.attention.backend.flash_attention.flash_attn_func") as mock_fa: mock_fa.return_value = torch.zeros(batch, seqlen, num_heads, head_dim) # output_attentions warning - with patch('xorl.models.layers.attention.backend.flash_attention.logger') as mock_logger: + with patch("xorl.models.layers.attention.backend.flash_attention.logger") as mock_logger: flash_attention_forward( - module, query, key, value, attention_mask=None, output_attentions=True, + module, + query, + key, + value, + attention_mask=None, + output_attentions=True, ) assert mock_logger.warning_once.called # head_mask warning - with patch('xorl.models.layers.attention.backend.flash_attention.logger') as mock_logger: + with patch("xorl.models.layers.attention.backend.flash_attention.logger") as mock_logger: flash_attention_forward( - module, query, key, value, attention_mask=None, + module, + query, + key, + value, + attention_mask=None, head_mask=torch.ones(num_heads), ) assert mock_logger.warning_once.called @@ -93,29 +105,48 @@ def test_flash_attention_api_behavior(self): # is_causal kwarg popped, module.is_causal used module.is_causal = False flash_attention_forward( - module, query, key, value, attention_mask=None, is_causal=True, + module, + query, + key, + value, + attention_mask=None, + is_causal=True, ) - assert mock_fa.call_args[1]['causal'] == False + assert mock_fa.call_args[1]["causal"] == False # Returns None attention weights module.is_causal = True result, attn_weights = flash_attention_forward( - module, query, key, value, attention_mask=None, + module, + query, + key, + value, + attention_mask=None, ) assert result.shape == (batch, seqlen, num_heads, head_dim) assert attn_weights is None # Scaling passed as softmax_scale flash_attention_forward( - module, query, key, value, attention_mask=None, scaling=0.125, + module, + query, + key, + value, + attention_mask=None, + scaling=0.125, ) - assert mock_fa.call_args[1]['softmax_scale'] == 0.125 + assert mock_fa.call_args[1]["softmax_scale"] == 0.125 # Sliding window -> window_size tuple flash_attention_forward( - module, query, key, value, attention_mask=None, sliding_window=128, + module, + query, + key, + value, + attention_mask=None, + sliding_window=128, ) - assert mock_fa.call_args[1]['window_size'] == (128, 0) + assert mock_fa.call_args[1]["window_size"] == (128, 0) def test_varlen_path_with_cu_seqlens(self): """cu_seqlens kwargs trigger the varlen path.""" @@ -128,15 +159,21 @@ def test_varlen_path_with_cu_seqlens(self): value = torch.randn(1, total_tokens, num_heads, head_dim) cu_seqlens = torch.tensor([0, 16, 32], dtype=torch.int64) - with patch('xorl.models.layers.attention.backend.flash_attention.flash_attn_varlen_func') as mock_varlen: + with patch("xorl.models.layers.attention.backend.flash_attention.flash_attn_varlen_func") as mock_varlen: mock_varlen.return_value = torch.zeros(total_tokens, num_heads, head_dim) result, _ = flash_attention_forward( - module, query, key, value, attention_mask=None, - cu_seq_lens_q=cu_seqlens, cu_seq_lens_k=cu_seqlens, - max_length_q=16, max_length_k=16, + module, + query, + key, + value, + attention_mask=None, + cu_seq_lens_q=cu_seqlens, + cu_seq_lens_k=cu_seqlens, + max_length_q=16, + max_length_k=16, ) assert mock_varlen.called - assert mock_varlen.call_args[1]['cu_seqlens_q'].dtype == torch.int32 + assert mock_varlen.call_args[1]["cu_seqlens_q"].dtype == torch.int32 assert result.shape == (1, total_tokens, num_heads, head_dim) @pytest.mark.gpu @@ -177,8 +214,13 @@ def test_eager_attention_head_layout(self): value = torch.randn(batch, seq, kv_heads, head_dim) attn_output, attn_weights = eager_attention_forward( - module=module, query=query, key=key, value=value, - attention_mask=None, scaling=head_dim**-0.5, dropout=0.0, + module=module, + query=query, + key=key, + value=value, + attention_mask=None, + scaling=head_dim**-0.5, + dropout=0.0, ) assert attn_output.shape == (batch, seq, q_heads, head_dim) assert attn_weights.shape == (batch, q_heads, seq, seq) @@ -187,7 +229,10 @@ def test_eager_attention_head_layout(self): with pytest.raises(RuntimeError, match="query_heads=3 is not divisible by kv_heads=2"): eager_attention_forward( module=module, - query=torch.randn(1, 4, 3, 8), key=torch.randn(1, 4, 2, 8), + query=torch.randn(1, 4, 3, 8), + key=torch.randn(1, 4, 2, 8), value=torch.randn(1, 4, 2, 8), - attention_mask=None, scaling=8**-0.5, dropout=0.0, + attention_mask=None, + scaling=8**-0.5, + dropout=0.0, ) diff --git a/tests/ops/test_block_fp8.py b/tests/ops/test_block_fp8.py index 785bf96d..3f7fe12b 100644 --- a/tests/ops/test_block_fp8.py +++ b/tests/ops/test_block_fp8.py @@ -6,16 +6,23 @@ import pytest import torch -import numpy as np + # Try to import the block_fp8 module try: from xorl.ops.quantize import ( - block_fp8_quantize as block_fp8_quant, block_fp8_dequantize as block_fp8_dequant, + ) + from xorl.ops.quantize import ( block_fp8_dequantize_gkn as block_fp8_weight_dequant, + ) + from xorl.ops.quantize import ( block_fp8_gemm, ) + from xorl.ops.quantize import ( + block_fp8_quantize as block_fp8_quant, + ) + HAS_BLOCK_FP8 = True except ImportError: HAS_BLOCK_FP8 = False diff --git a/tests/ops/test_block_fp8_gkn.py b/tests/ops/test_block_fp8_gkn.py index 56db59d9..c0e8757b 100644 --- a/tests/ops/test_block_fp8_gkn.py +++ b/tests/ops/test_block_fp8_gkn.py @@ -8,7 +8,8 @@ import torch import triton -from xorl.ops.quantize import block_fp8_quantize_gkn, block_fp8_dequantize_gkn +from xorl.ops.quantize import block_fp8_dequantize_gkn, block_fp8_quantize_gkn + pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -109,6 +110,7 @@ def test_dequantize_output_and_requirements(self): # Bandwidth benchmarks # --------------------------------------------------------------------------- + @pytest.mark.benchmark class TestBlockFP8GKNBandwidth: """Bandwidth tests targeting >2000 GB/s on H100.""" @@ -132,7 +134,7 @@ def test_quantize_dequantize_gkn_bandwidth(self): elapsed_ms = start.elapsed_time(end) / 100 total_bytes = K * N * (4 + 1) bw_quant = total_bytes / (elapsed_ms * 1e-3) / 1e9 - print(f"\nblock_fp8_quantize_gkn: {bw_quant:.0f} GB/s ({elapsed_ms*1000:.1f} us)") + print(f"\nblock_fp8_quantize_gkn: {bw_quant:.0f} GB/s ({elapsed_ms * 1000:.1f} us)") assert bw_quant > 2000, f"Quant bandwidth {bw_quant:.0f} GB/s < 2000 GB/s" # --- Dequantize bandwidth --- @@ -150,5 +152,5 @@ def test_quantize_dequantize_gkn_bandwidth(self): elapsed_ms2 = start2.elapsed_time(end2) / 100 total_bytes2 = K * N * (1 + 4) bw_dequant = total_bytes2 / (elapsed_ms2 * 1e-3) / 1e9 - print(f"\nblock_fp8_dequantize_gkn: {bw_dequant:.0f} GB/s ({elapsed_ms2*1000:.1f} us)") + print(f"\nblock_fp8_dequantize_gkn: {bw_dequant:.0f} GB/s ({elapsed_ms2 * 1000:.1f} us)") assert bw_dequant > 2000, f"Dequant bandwidth {bw_dequant:.0f} GB/s < 2000 GB/s" diff --git a/tests/ops/test_eager_vs_native_moe.py b/tests/ops/test_eager_vs_native_moe.py index 7e759887..46faee23 100644 --- a/tests/ops/test_eager_vs_native_moe.py +++ b/tests/ops/test_eager_vs_native_moe.py @@ -9,6 +9,7 @@ import torch import torch.nn as nn + DEVICE = "cuda" DTYPE = torch.bfloat16 @@ -17,11 +18,13 @@ # Helpers (lazy-import to avoid torchvision env crash at module level) # --------------------------------------------------------------------------- + def _import_moe(): """Import MoE layers; returns (MoEBlock, MoEExperts) or skips.""" try: - from xorl.models.layers.moe.moe_block import MoEBlock from xorl.models.layers.moe.experts import MoEExperts + from xorl.models.layers.moe.moe_block import MoEBlock + return MoEBlock, MoEExperts except Exception as e: pytest.skip(f"Cannot import MoE layers: {e}") @@ -57,10 +60,10 @@ def _make_pair(num_experts, hidden_dim, intermediate, top_k, seed=42): # (num_experts, hidden_dim, intermediate, top_k, batch, seq) (4, 64, 128, 2, 2, 8), (8, 128, 256, 2, 4, 16), - (4, 64, 128, 1, 2, 8), # top_k=1 - (8, 128, 256, 4, 2, 16), # top_k=4 - (16, 64, 128, 2, 1, 4), # 16 experts - (4, 64, 128, 2, 1, 1), # minimal seq + (4, 64, 128, 1, 2, 8), # top_k=1 + (8, 128, 256, 4, 2, 16), # top_k=4 + (16, 64, 128, 2, 1, 4), # 16 experts + (4, 64, 128, 2, 1, 1), # minimal seq ] @@ -80,7 +83,10 @@ def test_forward_and_backward_agreement(ne, hd, inter, topk, bs, seq): torch.testing.assert_close(eager_logits, native_logits, atol=0, rtol=0) max_diff = (eager_out - native_out).abs().max().item() torch.testing.assert_close( - native_out, eager_out, atol=0.05, rtol=0.02, + native_out, + eager_out, + atol=0.05, + rtol=0.02, msg=f"Forward mismatch: max_diff={max_diff:.6f}", ) @@ -104,8 +110,11 @@ def test_forward_and_backward_agreement(ne, hd, inter, topk, bs, seq): assert native_grad is not None, f"native {name} grad is None" torch.testing.assert_close(native_grad, eager_grad, atol=atol, rtol=rtol, msg=f"{name} gradient mismatch") torch.testing.assert_close( - native_block.gate.weight.grad, eager_block.gate.weight.grad, - atol=atol, rtol=rtol, msg="Gate weight gradient mismatch", + native_block.gate.weight.grad, + eager_block.gate.weight.grad, + atol=atol, + rtol=rtol, + msg="Gate weight gradient mismatch", ) @@ -113,6 +122,7 @@ def test_forward_and_backward_agreement(ne, hd, inter, topk, bs, seq): # Test 2: Determinism + edge case (all tokens same expert) # --------------------------------------------------------------------------- + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_determinism_and_edge_cases(): """Determinism: same input produces identical output. Edge case: all tokens to same expert.""" @@ -152,13 +162,21 @@ def test_determinism_and_edge_cases(): with torch.no_grad(): native_out = native_expert_forward( - x_edge, routing_weights, selected_experts, - experts.gate_proj, experts.up_proj, experts.down_proj, num_experts=ne, + x_edge, + routing_weights, + selected_experts, + experts.gate_proj, + experts.up_proj, + experts.down_proj, + num_experts=ne, ) eager_out = experts(x_edge, expert_idx=0) torch.testing.assert_close( - native_out, eager_out, atol=0.01, rtol=0.01, + native_out, + eager_out, + atol=0.01, + rtol=0.01, msg="Same-expert output mismatch", ) @@ -167,6 +185,7 @@ def test_determinism_and_edge_cases(): # Test 3: Large scale forward + backward # --------------------------------------------------------------------------- + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_large_scale(): """Forward + backward agreement at larger dimensions (E=64, H=512, I=1024, K=8).""" @@ -180,7 +199,10 @@ def test_large_scale(): eager_out, _ = eager_block(x) native_out, _ = native_block(x) torch.testing.assert_close( - native_out, eager_out, atol=0.1, rtol=0.05, + native_out, + eager_out, + atol=0.1, + rtol=0.05, msg="Large scale forward mismatch", ) diff --git a/tests/ops/test_group_gemm.py b/tests/ops/test_group_gemm.py index 2d632d85..7772b630 100644 --- a/tests/ops/test_group_gemm.py +++ b/tests/ops/test_group_gemm.py @@ -6,15 +6,18 @@ import pytest import torch -from typing import List + # Mark all tests as GPU since group_gemm requires CUDA pytestmark = pytest.mark.gpu def naive_group_gemm_same_nk( - a: torch.Tensor, b: torch.Tensor, cumsum_M: torch.Tensor, - transpose_a: bool = False, transpose_b: bool = False, + a: torch.Tensor, + b: torch.Tensor, + cumsum_M: torch.Tensor, + transpose_a: bool = False, + transpose_b: bool = False, ) -> torch.Tensor: """Naive PyTorch implementation of grouped GEMM with same N, K.""" G = b.shape[0] @@ -47,8 +50,13 @@ def naive_group_gemm_same_nk( def naive_group_gemm_same_mn( - a: torch.Tensor, b: torch.Tensor, cumsum_K: torch.Tensor, - M: int, N: int, transpose_a: bool = False, transpose_b: bool = False, + a: torch.Tensor, + b: torch.Tensor, + cumsum_K: torch.Tensor, + M: int, + N: int, + transpose_a: bool = False, + transpose_b: bool = False, ) -> torch.Tensor: """Naive PyTorch implementation of grouped GEMM with same M, N.""" G = cumsum_K.shape[0] @@ -92,7 +100,7 @@ def test_same_nk_comprehensive(self): G, K, N = 4, 128, 256 group_sizes = [10, 20, 15, 25] total_M = sum(group_sizes) - cumsum_M = torch.tensor([sum(group_sizes[:i+1]) for i in range(G)], dtype=torch.int32).cuda() + cumsum_M = torch.tensor([sum(group_sizes[: i + 1]) for i in range(G)], dtype=torch.int32).cuda() max_M = max(group_sizes) a = torch.randn(total_M, K, dtype=torch.bfloat16).cuda() @@ -107,7 +115,7 @@ def test_same_nk_comprehensive(self): G2, K2, N2 = 8, 64, 128 gs2 = [5, 100, 2, 50, 30, 8, 45, 20] total_M2 = sum(gs2) - cumsum_M2 = torch.tensor([sum(gs2[:i+1]) for i in range(G2)], dtype=torch.int32).cuda() + cumsum_M2 = torch.tensor([sum(gs2[: i + 1]) for i in range(G2)], dtype=torch.int32).cuda() a2 = torch.randn(total_M2, K2, dtype=torch.float16).cuda() b2 = torch.randn(G2, K2, N2, dtype=torch.float16).cuda() out2 = group_gemm_same_nk(a2, b2, cumsum_M2, max(gs2)) @@ -125,7 +133,7 @@ def test_same_nk_comprehensive(self): G4, K4, N4 = 8, 4096, 14336 gs4 = [512, 480, 520, 490, 510, 505, 495, 488] total_M4 = sum(gs4) - cumsum_M4 = torch.tensor([sum(gs4[:i+1]) for i in range(G4)], dtype=torch.int32).cuda() + cumsum_M4 = torch.tensor([sum(gs4[: i + 1]) for i in range(G4)], dtype=torch.int32).cuda() a4 = torch.randn(total_M4, K4, dtype=torch.bfloat16).cuda() b4 = torch.randn(G4, K4, N4, dtype=torch.bfloat16).cuda() out4 = group_gemm_same_nk(a4, b4, cumsum_M4, max(gs4)) @@ -158,7 +166,7 @@ def test_same_mn_comprehensive(self): G, M, N = 4, 128, 256 group_Ks = [64, 128, 96, 112] total_K = sum(group_Ks) - cumsum_K = torch.tensor([sum(group_Ks[:i+1]) for i in range(G)], dtype=torch.int32).cuda() + cumsum_K = torch.tensor([sum(group_Ks[: i + 1]) for i in range(G)], dtype=torch.int32).cuda() a = torch.randn(total_K, M, dtype=torch.bfloat16).cuda() b = torch.randn(total_K, N, dtype=torch.bfloat16).cuda() @@ -172,7 +180,7 @@ def test_same_mn_comprehensive(self): G2, M2, N2 = 8, 64, 128 gKs2 = [10, 200, 5, 100, 50, 15, 90, 30] total_K2 = sum(gKs2) - cumsum_K2 = torch.tensor([sum(gKs2[:i+1]) for i in range(G2)], dtype=torch.int32).cuda() + cumsum_K2 = torch.tensor([sum(gKs2[: i + 1]) for i in range(G2)], dtype=torch.int32).cuda() a2 = torch.randn(total_K2, M2, dtype=torch.float16).cuda() b2 = torch.randn(total_K2, N2, dtype=torch.float16).cuda() c2 = torch.empty(G2, M2, N2, dtype=torch.float16).cuda() @@ -184,7 +192,7 @@ def test_same_mn_comprehensive(self): G3, M3, N3 = 4, 64, 128 gKs3 = [64, 0, 96, 32] total_K3 = sum(gKs3) - cumsum_K3 = torch.tensor([sum(gKs3[:i+1]) for i in range(G3)], dtype=torch.int32).cuda() + cumsum_K3 = torch.tensor([sum(gKs3[: i + 1]) for i in range(G3)], dtype=torch.int32).cuda() a3 = torch.randn(total_K3, M3, dtype=torch.bfloat16).cuda() b3 = torch.randn(total_K3, N3, dtype=torch.bfloat16).cuda() c3 = torch.empty(G3, M3, N3, dtype=torch.bfloat16).cuda() @@ -219,7 +227,7 @@ def test_properties(self): G, K, N = 2, 64, 128 group_sizes = [32, 32] total_M = sum(group_sizes) - cumsum_M = torch.tensor([sum(group_sizes[:i+1]) for i in range(G)], dtype=torch.int32).cuda() + cumsum_M = torch.tensor([sum(group_sizes[: i + 1]) for i in range(G)], dtype=torch.int32).cuda() max_M = max(group_sizes) # Dtype support diff --git a/tests/ops/test_moe_act.py b/tests/ops/test_moe_act.py index 69bc2495..1c0a3e58 100644 --- a/tests/ops/test_moe_act.py +++ b/tests/ops/test_moe_act.py @@ -17,6 +17,7 @@ import torch import torch.nn as nn + DEVICE = "cuda" DTYPE = torch.bfloat16 @@ -25,9 +26,11 @@ # Helpers # --------------------------------------------------------------------------- + def _available_moe_act_backends(): """Return backends that have moe_act local variants registered.""" from xorl.models.layers.moe.backend import MOE_EXPERT_BACKENDS_MOE_ACT + return list(MOE_EXPERT_BACKENDS_MOE_ACT.keys()) @@ -67,11 +70,11 @@ def _make_block_pair(ne, hd, inter, topk, backend, seed=42): CORRECTNESS_CONFIGS = [ # (num_experts, hidden, intermediate, top_k, batch, seq) - (4, 64, 128, 2, 2, 8), - (8, 128, 256, 2, 4, 16), - (4, 64, 128, 1, 2, 8), # top_k=1 - (8, 128, 256, 4, 2, 16), # top_k=4 - (4, 64, 128, 2, 1, 1), # minimal + (4, 64, 128, 2, 2, 8), + (8, 128, 256, 2, 4, 16), + (4, 64, 128, 1, 2, 8), # top_k=1 + (8, 128, 256, 4, 2, 16), # top_k=4 + (4, 64, 128, 2, 1, 1), # minimal ] @@ -94,7 +97,10 @@ def test_forward_correctness(backend, ne, hd, inter, topk, bs, seq): max_diff = (act_out - std_out).abs().max().item() torch.testing.assert_close( - act_out, std_out, atol=0.05, rtol=0.02, + act_out, + std_out, + atol=0.05, + rtol=0.02, msg=f"[{backend}] Forward mismatch: max_diff={max_diff:.6f}", ) @@ -104,7 +110,7 @@ def test_forward_correctness(backend, ne, hd, inter, topk, bs, seq): # --------------------------------------------------------------------------- BACKWARD_CONFIGS = [ - (4, 64, 128, 2, 2, 8), + (4, 64, 128, 2, 2, 8), (8, 128, 256, 2, 4, 16), ] @@ -129,7 +135,10 @@ def test_backward_correctness(backend, ne, hd, inter, topk, bs, seq): act_out.sum().backward() torch.testing.assert_close( - x_act.grad, x_std.grad, atol=atol, rtol=rtol, + x_act.grad, + x_std.grad, + atol=atol, + rtol=rtol, msg=f"[{backend}] Input gradient mismatch", ) for name in ["gate_proj", "up_proj", "down_proj"]: @@ -138,11 +147,17 @@ def test_backward_correctness(backend, ne, hd, inter, topk, bs, seq): assert g_std is not None, f"[{backend}] std {name}.grad is None" assert g_act is not None, f"[{backend}] act {name}.grad is None" torch.testing.assert_close( - g_act, g_std, atol=atol, rtol=rtol, + g_act, + g_std, + atol=atol, + rtol=rtol, msg=f"[{backend}] {name} gradient mismatch", ) torch.testing.assert_close( - act.gate.weight.grad, std.gate.weight.grad, atol=atol, rtol=rtol, + act.gate.weight.grad, + std.gate.weight.grad, + atol=atol, + rtol=rtol, msg=f"[{backend}] Gate weight gradient mismatch", ) @@ -151,6 +166,7 @@ def test_backward_correctness(backend, ne, hd, inter, topk, bs, seq): # Test 3: Activation memory savings # --------------------------------------------------------------------------- + def _measure_fwd_bwd_peak_memory(block, x, warmup=3): """Peak GPU memory (bytes) for one forward+backward call.""" for _ in range(warmup): @@ -189,16 +205,15 @@ def test_memory_savings(backend): savings_mb = (mem_std - mem_act) / 1024**2 print( - f"\n[{backend}] Memory: std={mem_std/1024**2:.1f} MB " - f"moe_act={mem_act/1024**2:.1f} MB " + f"\n[{backend}] Memory: std={mem_std / 1024**2:.1f} MB " + f"moe_act={mem_act / 1024**2:.1f} MB " f"savings={savings_mb:+.1f} MB" ) # moe_act should not use significantly more memory than standard # (allow 5% overhead for checkpoint bookkeeping) assert mem_act <= mem_std * 1.05, ( - f"[{backend}] moe_act used more memory than standard: " - f"{mem_act/1024**2:.1f} MB vs {mem_std/1024**2:.1f} MB" + f"[{backend}] moe_act used more memory than standard: {mem_act / 1024**2:.1f} MB vs {mem_std / 1024**2:.1f} MB" ) @@ -206,6 +221,7 @@ def test_memory_savings(backend): # Benchmark: TFLOPS standard vs moe_act per backend # --------------------------------------------------------------------------- + def _moe_flops(bs, seq, hd, inter, topk): """Forward FLOPs: 3 GEMMs × 2 (matmul count) × tokens × top_k.""" return bs * seq * topk * 6 * hd * inter @@ -256,26 +272,20 @@ def bench_moe_act_tflops(seq_len): t_act = _benchmark(act, x) results[backend] = { - "std_ms": t_std * 1000, - "act_ms": t_act * 1000, + "std_ms": t_std * 1000, + "act_ms": t_act * 1000, "std_tflops": flops / t_std / 1e12, "act_tflops": flops / t_act / 1e12, - "overhead": t_act / t_std, + "overhead": t_act / t_std, } del std, act print("\n" + "=" * 90) - print( - f" moe_act TFLOPS (bs={bs}, seq={seq_len}, hidden={hd}, " - f"inter={inter}, E={ne}, top_k={topk})" - ) - print(f" FLOPs/fwd: {flops/1e9:.1f} GFLOP | warmup=20, iters=40") + print(f" moe_act TFLOPS (bs={bs}, seq={seq_len}, hidden={hd}, inter={inter}, E={ne}, top_k={topk})") + print(f" FLOPs/fwd: {flops / 1e9:.1f} GFLOP | warmup=20, iters=40") print("=" * 90) - print( - f" {'Backend':<10} {'Std (ms)':>10} {'Act (ms)':>10} " - f"{'Std TF':>10} {'Act TF':>10} {'Overhead':>10}" - ) + print(f" {'Backend':<10} {'Std (ms)':>10} {'Act (ms)':>10} {'Std TF':>10} {'Act TF':>10} {'Overhead':>10}") print("-" * 90) for backend, r in results.items(): print( @@ -290,13 +300,14 @@ def bench_moe_act_tflops(seq_len): # Test 4: moe_act works correctly inside gradient_checkpointing_enable # --------------------------------------------------------------------------- + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("backend", AVAILABLE_BACKENDS) def test_moe_act_via_gradient_checkpointing_enable(backend): """gradient_checkpointing_enable(moe_checkpoint_method='moe_act') sets _moe_act correctly.""" + from xorl.models.layers.moe.experts import MoEExperts from xorl.models.transformers.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import Qwen3MoeForCausalLM - from xorl.models.layers.moe.experts import MoEExperts config = Qwen3MoeConfig( vocab_size=1000, @@ -319,11 +330,13 @@ def test_moe_act_via_gradient_checkpointing_enable(backend): model = Qwen3MoeForCausalLM(config).to(DEVICE, DTYPE) # Enable selective GC with moe_act - model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={ - "use_reentrant": False, - "recompute_modules": ["self_attn", "mlp"], - "moe_checkpoint_method": "moe_act", - }) + model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={ + "use_reentrant": False, + "recompute_modules": ["self_attn", "mlp"], + "moe_checkpoint_method": "moe_act", + } + ) # Verify _moe_act is set on all MoEExperts modules moe_experts_modules = [m for m in model.modules() if isinstance(m, MoEExperts)] @@ -344,6 +357,7 @@ def test_moe_act_via_gradient_checkpointing_enable(backend): # Test 5: moe_act + torch.compile correctness # --------------------------------------------------------------------------- + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("backend", AVAILABLE_BACKENDS) def test_moe_act_compile(backend): @@ -368,10 +382,8 @@ def test_moe_act_compile(backend): comp_out, _ = compiled(x2) comp_out.sum().backward() - torch.testing.assert_close(comp_out, ref_out, atol=0.05, rtol=0.02, - msg=f"[{backend}] compile forward mismatch") - torch.testing.assert_close(x2.grad, ref_grad, atol=0.05, rtol=0.05, - msg=f"[{backend}] compile backward mismatch") + torch.testing.assert_close(comp_out, ref_out, atol=0.05, rtol=0.02, msg=f"[{backend}] compile forward mismatch") + torch.testing.assert_close(x2.grad, ref_grad, atol=0.05, rtol=0.05, msg=f"[{backend}] compile backward mismatch") torch._dynamo.reset() diff --git a/tests/ops/test_moe_gkn_format.py b/tests/ops/test_moe_gkn_format.py index bb0b9418..ae678955 100644 --- a/tests/ops/test_moe_gkn_format.py +++ b/tests/ops/test_moe_gkn_format.py @@ -20,9 +20,15 @@ # Reference implementation -- per-expert loop with nn.Linear weights # --------------------------------------------------------------------------- + def reference_moe_forward( - hidden_states, routing_weights, selected_experts, - gate_proj_gkn, up_proj_gkn, down_proj_gkn, num_experts, + hidden_states, + routing_weights, + selected_experts, + gate_proj_gkn, + up_proj_gkn, + down_proj_gkn, + num_experts, ): """Naive per-expert loop MoE forward using (G,K,N) weights directly.""" num_tokens = hidden_states.shape[0] @@ -54,6 +60,7 @@ def make_gkn_weights_from_linear(experts_list): class ExpertMLP(nn.Module): """Single expert MLP for reference.""" + def __init__(self, hidden_size, intermediate_size): super().__init__() self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) @@ -103,8 +110,13 @@ def test_gkn_format_correctness(self): routing_weights = torch.softmax(torch.randn(num_tokens, top_k), dim=-1) output_gkn = reference_moe_forward( - hidden_states, routing_weights, selected_experts, - gate_gkn, up_gkn, down_gkn, num_experts2, + hidden_states, + routing_weights, + selected_experts, + gate_gkn, + up_gkn, + down_gkn, + num_experts2, ) output_hf = torch.zeros_like(hidden_states) @@ -127,6 +139,7 @@ class TestCheckpointLoadingGKN: def _get_buffer_class(): try: from xorl.models.checkpoint_handlers.buffers import ExpertWeightBuffer + return ExpertWeightBuffer except (ImportError, ModuleNotFoundError): pytest.skip("checkpoint_handlers import requires transformers") @@ -166,8 +179,13 @@ def test_checkpoint_loading_and_output(self): routing_weights = torch.softmax(torch.randn(num_tokens, top_k), dim=-1) output_gkn = reference_moe_forward( - hidden_states, routing_weights, selected_experts, - gate_stacked, up_stacked, down_stacked, num_experts, + hidden_states, + routing_weights, + selected_experts, + gate_stacked, + up_stacked, + down_stacked, + num_experts, ) output_hf = torch.zeros_like(hidden_states) @@ -223,7 +241,13 @@ def test_backends_match_reference_and_agree(self): rw = torch.softmax(torch.randn(num_tokens, top_k, device=device, dtype=dtype), dim=-1) native_out = native_expert_forward( - hidden_states, rw, selected, gate_cuda, up_cuda, down_cuda, num_experts, + hidden_states, + rw, + selected, + gate_cuda, + up_cuda, + down_cuda, + num_experts, ) ref_out = reference_moe_forward(hidden_states, rw, selected, gate_cuda, up_cuda, down_cuda, num_experts) torch.testing.assert_close(native_out, ref_out, atol=0.02, rtol=0.02) @@ -231,19 +255,32 @@ def test_backends_match_reference_and_agree(self): # --- Triton backend --- try: from xorl.utils.import_utils import is_fused_moe_available + if not is_fused_moe_available(): raise ImportError from xorl.ops.moe.triton import TritonMoeExpertsFunction triton_out = TritonMoeExpertsFunction.apply( - num_experts, rw, selected, hidden_states, gate_cuda, up_cuda, down_cuda, + num_experts, + rw, + selected, + hidden_states, + gate_cuda, + up_cuda, + down_cuda, ) torch.testing.assert_close(triton_out, ref_out, atol=0.01, rtol=0.01) # Triton backward: gradients exist and non-zero - gate_g = torch.randn(num_experts, hidden_size, intermediate_size, device=device, dtype=dtype, requires_grad=True) - up_g = torch.randn(num_experts, hidden_size, intermediate_size, device=device, dtype=dtype, requires_grad=True) - down_g = torch.randn(num_experts, intermediate_size, hidden_size, device=device, dtype=dtype, requires_grad=True) + gate_g = torch.randn( + num_experts, hidden_size, intermediate_size, device=device, dtype=dtype, requires_grad=True + ) + up_g = torch.randn( + num_experts, hidden_size, intermediate_size, device=device, dtype=dtype, requires_grad=True + ) + down_g = torch.randn( + num_experts, intermediate_size, hidden_size, device=device, dtype=dtype, requires_grad=True + ) h_g = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype, requires_grad=True) out_g = TritonMoeExpertsFunction.apply(num_experts, rw, selected, h_g, gate_g, up_g, down_g) out_g.sum().backward() @@ -254,7 +291,7 @@ def test_backends_match_reference_and_agree(self): # --- MoEBlock all backends agree --- try: - from xorl.models.layers.moe import MoEBlock, MOE_EXPERT_BACKENDS + from xorl.models.layers.moe import MOE_EXPERT_BACKENDS, MoEBlock except (ImportError, ModuleNotFoundError): return # skip if not importable @@ -266,8 +303,11 @@ def test_backends_match_reference_and_agree(self): for backend in MOE_EXPERT_BACKENDS: block = MoEBlock( - hidden_size=hidden_size, num_experts=num_experts, top_k=2, - intermediate_size=intermediate_size, moe_implementation=backend, + hidden_size=hidden_size, + num_experts=num_experts, + top_k=2, + intermediate_size=intermediate_size, + moe_implementation=backend, ) with torch.no_grad(): block.experts.gate_proj.copy_(g_gkn) @@ -282,8 +322,10 @@ def test_backends_match_reference_and_agree(self): for i in range(len(backends)): for j in range(i + 1, len(backends)): torch.testing.assert_close( - outputs[backends[i]], outputs[backends[j]], - atol=0.05, rtol=0.02, + outputs[backends[i]], + outputs[backends[j]], + atol=0.05, + rtol=0.02, msg=f"Backend mismatch: {backends[i]} vs {backends[j]}", ) diff --git a/tests/ops/test_moe_ops.py b/tests/ops/test_moe_ops.py index 72fb9da2..0ba1caf5 100644 --- a/tests/ops/test_moe_ops.py +++ b/tests/ops/test_moe_ops.py @@ -7,6 +7,7 @@ import pytest import torch + # Mark all tests as GPU since MoE ops require CUDA pytestmark = pytest.mark.gpu @@ -82,9 +83,7 @@ def test_histogram_and_index_compute(self): assert torch.equal(output_i64.cpu(), naive_expert_histogram(input_i64, 4).cpu()) # --- Index compute: basic + ordering --- - experts_for_tokens = torch.tensor( - [[0, 1], [1, 2], [0, 2], [2, 3]], dtype=torch.int32 - ).cuda() + experts_for_tokens = torch.tensor([[0, 1], [1, 2], [0, 2], [2, 3]], dtype=torch.int32).cuda() histogram = expert_histogram(experts_for_tokens.flatten(), 4) cumsum = torch.cumsum(histogram, dim=0).int().cuda() indices = moe_index_compute(experts_for_tokens, cumsum) @@ -94,9 +93,7 @@ def test_histogram_and_index_compute(self): assert indices.flatten().unique().numel() == experts_for_tokens.numel() # Ordering: same expert gets sequential indices - experts_same = torch.tensor( - [[0, 0], [0, 0], [1, 1], [1, 1]], dtype=torch.int32 - ).cuda() + experts_same = torch.tensor([[0, 0], [0, 0], [1, 1], [1, 1]], dtype=torch.int32).cuda() hist2 = expert_histogram(experts_same.flatten(), 2) cumsum2 = torch.cumsum(hist2, dim=0).int().cuda() indices2 = moe_index_compute(experts_same, cumsum2) @@ -114,7 +111,7 @@ def test_gather_scatter_add_gather(self): if not torch.cuda.is_available(): pytest.skip("CUDA not available") try: - from xorl.ops.group_gemm.kernel.moe import moe_gather, moe_scatter, moe_add_gather + from xorl.ops.group_gemm.kernel.moe import moe_add_gather, moe_gather, moe_scatter except ImportError: pytest.skip("moe ops not available") @@ -131,8 +128,10 @@ def test_gather_scatter_add_gather(self): x2 = torch.randn(M2 * topk2, N2, dtype=torch.bfloat16).cuda() index2 = torch.tensor([[0, 1], [1, 2], [0, 2]], dtype=torch.int32).cuda() assert torch.allclose( - moe_gather(x2, index2).float(), naive_moe_gather(x2, index2).float(), - rtol=1e-2, atol=1e-2, + moe_gather(x2, index2).float(), + naive_moe_gather(x2, index2).float(), + rtol=1e-2, + atol=1e-2, ) # Gather: large dimensions @@ -140,8 +139,10 @@ def test_gather_scatter_add_gather(self): x3 = torch.randn(M3 * topk3, N3, dtype=torch.float16).cuda() index3 = torch.arange(M3 * topk3, dtype=torch.int32).cuda().reshape(M3, topk3) assert torch.allclose( - moe_gather(x3, index3).float(), naive_moe_gather(x3, index3).float(), - rtol=1e-2, atol=1e-2, + moe_gather(x3, index3).float(), + naive_moe_gather(x3, index3).float(), + rtol=1e-2, + atol=1e-2, ) # --- Scatter: basic --- @@ -174,8 +175,10 @@ def test_gather_scatter_add_gather(self): y5 = torch.randn(M5 * topk5, N5, dtype=torch.bfloat16).cuda() idx5 = torch.arange(M5 * topk5, dtype=torch.int32).cuda().reshape(M5, topk5) assert torch.allclose( - moe_add_gather(x5, y5, idx5), moe_gather(x5 + y5, idx5), - rtol=1e-3, atol=1e-3, + moe_add_gather(x5, y5, idx5), + moe_gather(x5 + y5, idx5), + rtol=1e-3, + atol=1e-3, ) @@ -188,7 +191,10 @@ def test_full_moe_pipeline(self): pytest.skip("CUDA not available") try: from xorl.ops.group_gemm.kernel.moe import ( - expert_histogram, moe_scatter, moe_gather, moe_index_compute, + expert_histogram, + moe_gather, + moe_index_compute, + moe_scatter, ) except ImportError: pytest.skip("moe ops not available") diff --git a/tests/ops/test_moe_torch_compile.py b/tests/ops/test_moe_torch_compile.py index 9e0e8b39..a22b310e 100644 --- a/tests/ops/test_moe_torch_compile.py +++ b/tests/ops/test_moe_torch_compile.py @@ -13,10 +13,10 @@ """ import pytest -import time import torch import torch.nn as nn + DEVICE = "cuda" DTYPE = torch.bfloat16 @@ -50,6 +50,7 @@ def _tiny_moe_config(**overrides): def _make_position_embeddings(config, seq_len, device, dtype): """Create position_embeddings (cos, sin) for decoder layer tests.""" from xorl.models.layers.rope import RotaryEmbedding + rotary = RotaryEmbedding(config=config).to(device) dummy_hidden = torch.randn(1, seq_len, config.hidden_size, device=device, dtype=dtype) position_ids = torch.arange(seq_len, device=device).unsqueeze(0) @@ -60,6 +61,7 @@ def _make_position_embeddings(config, seq_len, device, dtype): def _make_moe_block(moe_backend, hidden_size=128, num_experts=4, top_k=2, intermediate=128): """Create an MoEBlock with xavier init for numerical stability.""" from xorl.models.layers.moe.moe_block import MoEBlock + block = MoEBlock( hidden_size=hidden_size, num_experts=num_experts, @@ -79,12 +81,12 @@ def _available_backends(): backends = ["native", "eager"] try: from xorl.utils.import_utils import is_fused_moe_available + if is_fused_moe_available(): backends.append("triton") except Exception: pass try: - from xorl.ops.group_gemm.kernel.quack import quack_group_gemm_same_nk backends.append("quack") except Exception: pass @@ -98,6 +100,7 @@ def _available_backends(): # Test 1: MoEBlock compile -- aot_eager + inductor + fullgraph + correctness # --------------------------------------------------------------------------- + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("moe_backend", AVAILABLE_BACKENDS) def test_moe_block_compile(moe_backend): @@ -148,6 +151,7 @@ def test_moe_block_compile(moe_backend): # Test 2: Qwen3MoeDecoderLayer compile (aot_eager + inductor) # --------------------------------------------------------------------------- + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("moe_backend", AVAILABLE_BACKENDS) def test_decoder_layer_compile(moe_backend): @@ -181,19 +185,24 @@ def test_decoder_layer_compile(moe_backend): # Test 3: Full model per-layer compile (torchtitan style) # --------------------------------------------------------------------------- + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("moe_backend,compile_backend", [ - ("native", "aot_eager"), - ("native", "inductor"), - ("eager", "aot_eager"), - ("eager", "inductor"), -] + ([("triton", "aot_eager"), ("triton", "inductor")] if "triton" in AVAILABLE_BACKENDS else []) - + ([("quack", "aot_eager"), ("quack", "inductor")] if "quack" in AVAILABLE_BACKENDS else [])) +@pytest.mark.parametrize( + "moe_backend,compile_backend", + [ + ("native", "aot_eager"), + ("native", "inductor"), + ("eager", "aot_eager"), + ("eager", "inductor"), + ] + + ([("triton", "aot_eager"), ("triton", "inductor")] if "triton" in AVAILABLE_BACKENDS else []) + + ([("quack", "aot_eager"), ("quack", "inductor")] if "quack" in AVAILABLE_BACKENDS else []), +) def test_full_model_per_layer_compile(moe_backend, compile_backend): """Apply torch.compile to each decoder layer, run forward + backward.""" from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import ( - Qwen3MoeForCausalLM, Qwen3MoeDecoderLayer, + Qwen3MoeForCausalLM, ) config = _tiny_moe_config(_moe_implementation=moe_backend) @@ -220,6 +229,7 @@ def test_full_model_per_layer_compile(moe_backend, compile_backend): # Benchmark: TFLOPS measurement compiled vs uncompiled # --------------------------------------------------------------------------- + def _moe_flops(batch, seq, hidden, intermediate, num_experts, top_k): """Estimate FLOPs for one MoE forward pass.""" tokens = batch * seq @@ -276,39 +286,51 @@ def bench_tflops(seq_len): torch._dynamo.reset() compiled_block = torch.compile( - block, fullgraph=False, backend="inductor", dynamic=False, + block, + fullgraph=False, + backend="inductor", + dynamic=False, ) try: t_compiled = _benchmark_moe_block(compiled_block, x) except Exception as e: print(f" {backend} inductor compile failed: {type(e).__name__}") results[backend] = { - "base_ms": t_base * 1000, "compiled_ms": float("nan"), - "base_tflops": tflops_base, "compiled_tflops": float("nan"), - "speedup": float("nan"), "compile_backend": "inductor (FAILED)", + "base_ms": t_base * 1000, + "compiled_ms": float("nan"), + "base_tflops": tflops_base, + "compiled_tflops": float("nan"), + "speedup": float("nan"), + "compile_backend": "inductor (FAILED)", } torch._dynamo.reset() continue tflops_compiled = flops / t_compiled / 1e12 speedup = t_base / t_compiled results[backend] = { - "base_ms": t_base * 1000, "compiled_ms": t_compiled * 1000, - "base_tflops": tflops_base, "compiled_tflops": tflops_compiled, - "speedup": speedup, "compile_backend": "inductor", + "base_ms": t_base * 1000, + "compiled_ms": t_compiled * 1000, + "base_tflops": tflops_base, + "compiled_tflops": tflops_compiled, + "speedup": speedup, + "compile_backend": "inductor", } print("\n" + "=" * 90) - print(f" MoE TFLOPS Benchmark (batch={batch}, seq={seq_len}, hidden={hidden}, " - f"E={num_experts}, top_k={top_k})") + print(f" MoE TFLOPS Benchmark (batch={batch}, seq={seq_len}, hidden={hidden}, E={num_experts}, top_k={top_k})") print(f" FLOPs per fwd: {flops / 1e9:.1f} GFLOP | dynamic=False, warmup=30, iters=50") print("=" * 95) - print(f" {'Backend':<10} {'Compiler':<10} {'Base (ms)':>10} {'Compiled (ms)':>14} " - f"{'Base TFLOPS':>12} {'Comp TFLOPS':>12} {'Speedup':>8}") + print( + f" {'Backend':<10} {'Compiler':<10} {'Base (ms)':>10} {'Compiled (ms)':>14} " + f"{'Base TFLOPS':>12} {'Comp TFLOPS':>12} {'Speedup':>8}" + ) print("-" * 95) for backend, r in results.items(): - print(f" {backend:<10} {r['compile_backend']:<10} {r['base_ms']:>10.2f} {r['compiled_ms']:>14.2f} " - f"{r['base_tflops']:>12.2f} {r['compiled_tflops']:>12.2f} " - f"{r['speedup']:>7.2f}x") + print( + f" {backend:<10} {r['compile_backend']:<10} {r['base_ms']:>10.2f} {r['compiled_ms']:>14.2f} " + f"{r['base_tflops']:>12.2f} {r['compiled_tflops']:>12.2f} " + f"{r['speedup']:>7.2f}x" + ) print("=" * 95) @@ -316,6 +338,7 @@ def bench_tflops(seq_len): # Benchmark: Peak memory usage compiled vs uncompiled # --------------------------------------------------------------------------- + def _measure_peak_memory(fn, warmup=5): """Run fn, return peak GPU memory allocated in bytes.""" for _ in range(warmup): @@ -357,7 +380,10 @@ def run_base(): torch._dynamo.reset() compiled_block = torch.compile( - block, fullgraph=False, backend="inductor", dynamic=False, + block, + fullgraph=False, + backend="inductor", + dynamic=False, ) def run_compiled(): @@ -370,8 +396,10 @@ def run_compiled(): except Exception as e: print(f" {backend} inductor compile failed: {type(e).__name__}") results[backend] = { - "base_mb": mem_base / 1024**2, "compiled_mb": float("nan"), - "diff_mb": float("nan"), "ratio": float("nan"), + "base_mb": mem_base / 1024**2, + "compiled_mb": float("nan"), + "diff_mb": float("nan"), + "ratio": float("nan"), "compile_backend": "inductor (FAILED)", } del compiled_block, x @@ -380,7 +408,8 @@ def run_compiled(): continue results[backend] = { - "base_mb": mem_base / 1024**2, "compiled_mb": mem_compiled / 1024**2, + "base_mb": mem_base / 1024**2, + "compiled_mb": mem_compiled / 1024**2, "diff_mb": (mem_compiled - mem_base) / 1024**2, "ratio": mem_compiled / mem_base if mem_base > 0 else float("inf"), "compile_backend": "inductor", @@ -391,16 +420,16 @@ def run_compiled(): torch.cuda.empty_cache() print("\n" + "=" * 95) - print(f" MoE Peak Memory (batch={batch}, seq={seq_len}, hidden={hidden}, " - f"E={num_experts}, top_k={top_k})") + print(f" MoE Peak Memory (batch={batch}, seq={seq_len}, hidden={hidden}, E={num_experts}, top_k={top_k})") print("=" * 95) - print(f" {'Backend':<10} {'Compiler':<10} {'Base (MB)':>10} {'Compiled (MB)':>14} " - f"{'Delta (MB)':>12} {'Ratio':>8}") + print(f" {'Backend':<10} {'Compiler':<10} {'Base (MB)':>10} {'Compiled (MB)':>14} {'Delta (MB)':>12} {'Ratio':>8}") print("-" * 95) for backend, r in results.items(): - print(f" {backend:<10} {r['compile_backend']:<10} {r['base_mb']:>10.1f} " - f"{r['compiled_mb']:>14.1f} {r['diff_mb']:>+12.1f} " - f"{r['ratio']:>7.2f}x") + print( + f" {backend:<10} {r['compile_backend']:<10} {r['base_mb']:>10.1f} " + f"{r['compiled_mb']:>14.1f} {r['diff_mb']:>+12.1f} " + f"{r['ratio']:>7.2f}x" + ) print("=" * 95) @@ -408,6 +437,7 @@ def run_compiled(): # Benchmark: Full decoder layer (attention + MoE + norms + residuals) # --------------------------------------------------------------------------- + def _decoder_layer_flops(batch, seq, hidden, intermediate, num_heads, num_kv_heads, num_experts, top_k): """Estimate FLOPs for one Qwen3MoeDecoderLayer forward pass.""" tokens = batch * seq @@ -422,7 +452,8 @@ def _benchmark_decoder_layer(layer, x, position_ids, position_embeddings, warmup """Benchmark decoder layer fwd+bwd, return median GPU time in seconds.""" for _ in range(warmup): outputs = layer( - hidden_states=x, position_ids=position_ids, + hidden_states=x, + position_ids=position_ids, position_embeddings=position_embeddings, ) outputs[0].sum().backward() @@ -437,7 +468,8 @@ def _benchmark_decoder_layer(layer, x, position_ids, position_embeddings, warmup end_event = torch.cuda.Event(enable_timing=True) start_event.record() outputs = layer( - hidden_states=x, position_ids=position_ids, + hidden_states=x, + position_ids=position_ids, position_embeddings=position_embeddings, ) outputs[0].sum().backward() @@ -456,7 +488,8 @@ def _measure_decoder_layer_peak_memory(layer, x, position_ids, position_embeddin """Measure peak GPU memory for one decoder layer fwd+bwd.""" for _ in range(warmup): outputs = layer( - hidden_states=x, position_ids=position_ids, + hidden_states=x, + position_ids=position_ids, position_embeddings=position_embeddings, ) outputs[0].sum().backward() @@ -469,7 +502,8 @@ def _measure_decoder_layer_peak_memory(layer, x, position_ids, position_embeddin torch.cuda.empty_cache() outputs = layer( - hidden_states=x, position_ids=position_ids, + hidden_states=x, + position_ids=position_ids, position_embeddings=position_embeddings, ) outputs[0].sum().backward() @@ -497,18 +531,29 @@ def test_decoder_layer_benchmark(seq_len): batch = 4 flops = _decoder_layer_flops( - batch, seq_len, hidden, intermediate, num_heads, num_kv_heads, num_experts, top_k, + batch, + seq_len, + hidden, + intermediate, + num_heads, + num_kv_heads, + num_experts, + top_k, ) results = {} for backend in AVAILABLE_BACKENDS: config = _tiny_moe_config( - hidden_size=hidden, intermediate_size=hidden * 4, + hidden_size=hidden, + intermediate_size=hidden * 4, moe_intermediate_size=intermediate, - num_attention_heads=num_heads, num_key_value_heads=num_kv_heads, - num_experts=num_experts, num_experts_per_tok=top_k, - _moe_implementation=backend, _attn_implementation="sdpa", + num_attention_heads=num_heads, + num_key_value_heads=num_kv_heads, + num_experts=num_experts, + num_experts_per_tok=top_k, + _moe_implementation=backend, + _attn_implementation="sdpa", ) layer = Qwen3MoeDecoderLayer(config, layer_idx=0).to(DEVICE, DTYPE) @@ -527,23 +572,32 @@ def test_decoder_layer_benchmark(seq_len): t_compiled = _benchmark_decoder_layer(compiled_layer, x, position_ids, position_embeddings) tflops_compiled = flops / t_compiled / 1e12 mem_compiled = _measure_decoder_layer_peak_memory( - compiled_layer, x, position_ids, position_embeddings, + compiled_layer, + x, + position_ids, + position_embeddings, ) results[backend] = { - "base_ms": t_base * 1000, "compiled_ms": t_compiled * 1000, - "base_tflops": tflops_base, "compiled_tflops": tflops_compiled, + "base_ms": t_base * 1000, + "compiled_ms": t_compiled * 1000, + "base_tflops": tflops_base, + "compiled_tflops": tflops_compiled, "speedup": t_base / t_compiled, - "base_mb": mem_base / 1024**2, "compiled_mb": mem_compiled / 1024**2, + "base_mb": mem_base / 1024**2, + "compiled_mb": mem_compiled / 1024**2, "mem_diff_mb": (mem_compiled - mem_base) / 1024**2, "compile_backend": "inductor", } except Exception as e: print(f" {backend} inductor compile failed: {type(e).__name__}") results[backend] = { - "base_ms": t_base * 1000, "compiled_ms": float("nan"), - "base_tflops": tflops_base, "compiled_tflops": float("nan"), + "base_ms": t_base * 1000, + "compiled_ms": float("nan"), + "base_tflops": tflops_base, + "compiled_tflops": float("nan"), "speedup": float("nan"), - "base_mb": mem_base / 1024**2, "compiled_mb": float("nan"), + "base_mb": mem_base / 1024**2, + "compiled_mb": float("nan"), "mem_diff_mb": float("nan"), "compile_backend": "inductor (FAILED)", } @@ -553,21 +607,27 @@ def test_decoder_layer_benchmark(seq_len): torch.cuda.empty_cache() print("\n" + "=" * 110) - print(f" Decoder Layer Benchmark (batch={batch}, seq={seq_len}, hidden={hidden}, " - f"heads={num_heads}/{num_kv_heads}, E={num_experts}, top_k={top_k})") + print( + f" Decoder Layer Benchmark (batch={batch}, seq={seq_len}, hidden={hidden}, " + f"heads={num_heads}/{num_kv_heads}, E={num_experts}, top_k={top_k})" + ) print(f" FLOPs per fwd: {flops / 1e9:.1f} GFLOP | dynamic=False, warmup=30, iters=50") print("=" * 110) - print(f" {'Backend':<10} {'Compiler':<10} {'Base ms':>8} {'Comp ms':>8} " - f"{'Base TF':>8} {'Comp TF':>8} {'Speed':>6} " - f"{'Base MB':>8} {'Comp MB':>8} {'Mem D':>8}") + print( + f" {'Backend':<10} {'Compiler':<10} {'Base ms':>8} {'Comp ms':>8} " + f"{'Base TF':>8} {'Comp TF':>8} {'Speed':>6} " + f"{'Base MB':>8} {'Comp MB':>8} {'Mem D':>8}" + ) print("-" * 110) for backend, r in results.items(): - print(f" {backend:<10} {r['compile_backend']:<10} " - f"{r['base_ms']:>8.2f} {r['compiled_ms']:>8.2f} " - f"{r['base_tflops']:>8.2f} {r['compiled_tflops']:>8.2f} " - f"{r['speedup']:>5.2f}x " - f"{r['base_mb']:>8.1f} {r['compiled_mb']:>8.1f} " - f"{r['mem_diff_mb']:>+8.1f}") + print( + f" {backend:<10} {r['compile_backend']:<10} " + f"{r['base_ms']:>8.2f} {r['compiled_ms']:>8.2f} " + f"{r['base_tflops']:>8.2f} {r['compiled_tflops']:>8.2f} " + f"{r['speedup']:>5.2f}x " + f"{r['base_mb']:>8.1f} {r['compiled_mb']:>8.1f} " + f"{r['mem_diff_mb']:>+8.1f}" + ) print("=" * 110) diff --git a/tests/ops/test_nf4.py b/tests/ops/test_nf4.py index 6794a6da..b914ff9b 100644 --- a/tests/ops/test_nf4.py +++ b/tests/ops/test_nf4.py @@ -3,14 +3,17 @@ Validates correctness (roundtrip accuracy, codec, shapes) and measures effective memory bandwidth of the dequantization kernels. """ + +import time + import pytest import torch -import time + try: - from xorl.ops.quantize import nf4_quantize, nf4_dequantize - from xorl.ops.quantize import nf4_quantize_gkn, nf4_dequantize_gkn - from xorl.ops.quantize.nf4_codec import NF4_TABLE, NF4_MIN_STEP + from xorl.ops.quantize import nf4_dequantize, nf4_dequantize_gkn, nf4_quantize, nf4_quantize_gkn + from xorl.ops.quantize.nf4_codec import NF4_MIN_STEP, NF4_TABLE + HAS_NF4 = True except ImportError: HAS_NF4 = False @@ -47,9 +50,7 @@ def test_roundtrip_exact_table_values(self): for i in range(16): expected = NF4_TABLE[i] actual = out[i, 0].item() - assert abs(actual - expected) < 0.05, ( - f"Code {i}: expected {expected:.4f}, got {actual:.4f}" - ) + assert abs(actual - expected) < 0.05, f"Code {i}: expected {expected:.4f}, got {actual:.4f}" class TestNF4Quantize1D: @@ -214,10 +215,13 @@ def _measure_bandwidth_gkn(self, packed, scales, K, N, group_size, num_iters=200 bw = total_bytes * num_iters / elapsed / 1e9 return bw - @pytest.mark.parametrize("size", [ - (4096, 4096), # ~16M elements - (8192, 8192), # ~67M elements - ]) + @pytest.mark.parametrize( + "size", + [ + (4096, 4096), # ~16M elements + (8192, 8192), # ~67M elements + ], + ) def test_dequant_1d_bandwidth(self, size): """1D dequant should achieve >2000 GB/s on large tensors.""" M, K = size @@ -230,10 +234,13 @@ def test_dequant_1d_bandwidth(self, size): if bw < 2000: pytest.skip(f"Bandwidth {bw:.0f} GB/s below 2000 GB/s target (may vary by GPU)") - @pytest.mark.parametrize("size", [ - (4096, 4096), - (2048, 7168), # MoE expert dimensions - ]) + @pytest.mark.parametrize( + "size", + [ + (4096, 4096), + (2048, 7168), # MoE expert dimensions + ], + ) def test_dequant_gkn_bandwidth(self, size): """GKN dequant should achieve >2000 GB/s on large tensors.""" K, N = size diff --git a/tests/ops/test_prequant_gnk_to_gkn.py b/tests/ops/test_prequant_gnk_to_gkn.py index 58d75ae1..f36782ba 100644 --- a/tests/ops/test_prequant_gnk_to_gkn.py +++ b/tests/ops/test_prequant_gnk_to_gkn.py @@ -18,12 +18,12 @@ import triton from xorl.ops.quantize import ( - block_fp8_quantize_gkn, block_fp8_dequantize_gkn, + block_fp8_quantize_gkn, nvfp4_quantize, - nvfp4_dequantize, ) -from xorl.ops.quantize.nvfp4_gkn_quantize import nvfp4_quantize_gkn, nvfp4_dequantize_gkn +from xorl.ops.quantize.nvfp4_gkn_quantize import nvfp4_dequantize_gkn, nvfp4_quantize_gkn + pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -32,6 +32,7 @@ # Block FP8 -- GNK -> GKN # ========================================================================= + class TestBlockFP8PrequantGNKtoGKN: """Verify block_fp8 quantized data can be correctly transposed from GNK to GKN.""" @@ -97,6 +98,7 @@ def test_block_fp8_transpose_roundtrip_and_stacking(self): # NVFP4 -- GNK -> GKN # ========================================================================= + class TestNVFP4PrequantGNKtoGKN: """Verify nvfp4 quantized data can be correctly transposed from GNK to GKN.""" diff --git a/tests/qlora/test_detect_prequantized.py b/tests/qlora/test_detect_prequantized.py index d5753ff7..04732822 100644 --- a/tests/qlora/test_detect_prequantized.py +++ b/tests/qlora/test_detect_prequantized.py @@ -8,17 +8,16 @@ """ import json -import os -import tempfile import pytest import torch + pytestmark = [pytest.mark.cpu] from xorl.models.checkpoint_handlers.buffers import ( - detect_prequantized_checkpoint, detect_prequantized_block_fp8_checkpoint, + detect_prequantized_checkpoint, get_prequantized_exclude_modules, ) from xorl.models.transformers.qwen3.checkpoint_handler import Qwen3CheckpointHandler @@ -33,10 +32,13 @@ def test_nvfp4_detection(self, tmp_path): # -- Positive cases -- # modelopt nested format with open(tmp_path / "hf_quant_config.json", "w") as f: - json.dump({ - "producer": {"name": "modelopt", "version": "0.34.1"}, - "quantization": {"quant_algo": "NVFP4", "group_size": 16, "exclude_modules": ["lm_head"]}, - }, f) + json.dump( + { + "producer": {"name": "modelopt", "version": "0.34.1"}, + "quantization": {"quant_algo": "NVFP4", "group_size": 16, "exclude_modules": ["lm_head"]}, + }, + f, + ) assert detect_prequantized_checkpoint(str(tmp_path)) is True # flat format @@ -93,30 +95,47 @@ def test_block_fp8_detection(self, tmp_path): d1 = tmp_path / "fp8" d1.mkdir() with open(d1 / "config.json", "w") as f: - json.dump({"quantization_config": { - "quant_method": "fp8", "fmt": "e4m3", "weight_block_size": [128, 128], - }}, f) + json.dump( + { + "quantization_config": { + "quant_method": "fp8", + "fmt": "e4m3", + "weight_block_size": [128, 128], + } + }, + f, + ) assert detect_prequantized_block_fp8_checkpoint(str(d1)) is True # safetensors index with weight_scale_inv d2 = tmp_path / "safe" d2.mkdir() with open(d2 / "model.safetensors.index.json", "w") as f: - json.dump({"weight_map": { - "model.layers.0.self_attn.q_proj.weight": "s.safetensors", - "model.layers.0.self_attn.q_proj.weight_scale_inv": "s.safetensors", - }}, f) + json.dump( + { + "weight_map": { + "model.layers.0.self_attn.q_proj.weight": "s.safetensors", + "model.layers.0.self_attn.q_proj.weight_scale_inv": "s.safetensors", + } + }, + f, + ) assert detect_prequantized_block_fp8_checkpoint(str(d2)) is True # NVFP4 index (weight_scale + weight_scale_2) not FP8 d3 = tmp_path / "nvfp4" d3.mkdir() with open(d3 / "model.safetensors.index.json", "w") as f: - json.dump({"weight_map": { - "m.q_proj.weight": "s.safetensors", - "m.q_proj.weight_scale": "s.safetensors", - "m.q_proj.weight_scale_2": "s.safetensors", - }}, f) + json.dump( + { + "weight_map": { + "m.q_proj.weight": "s.safetensors", + "m.q_proj.weight_scale": "s.safetensors", + "m.q_proj.weight_scale_2": "s.safetensors", + } + }, + f, + ) assert detect_prequantized_block_fp8_checkpoint(str(d3)) is False d4 = tmp_path / "wrong_bs" @@ -167,8 +186,11 @@ class TestQwen3DenseCheckpointHandlerPrequantized: def _make_handler(self, is_prequantized, **kwargs): return Qwen3CheckpointHandler( - num_attention_heads=64, num_key_value_heads=8, head_dim=128, - is_prequantized=is_prequantized, **kwargs, + num_attention_heads=64, + num_key_value_heads=8, + head_dim=128, + is_prequantized=is_prequantized, + **kwargs, ) def test_prequantized_skip_and_on_load(self): @@ -273,8 +295,11 @@ class TestCheckpointHandlerExcludeModules: def test_dense_handler_exclude_modules(self): """Excluded modules not skipped in skip_fn and on_load_weight; non-excluded skipped; empty/no exclude normal.""" handler = Qwen3CheckpointHandler( - num_attention_heads=64, num_key_value_heads=8, head_dim=128, - is_prequantized=True, exclude_modules={"down_proj"}, + num_attention_heads=64, + num_key_value_heads=8, + head_dim=128, + is_prequantized=True, + exclude_modules={"down_proj"}, ) skip_fn = handler.get_skip_key_fn() assert not skip_fn("model.layers.0.mlp.down_proj.weight") @@ -288,8 +313,11 @@ def test_dense_handler_exclude_modules(self): # Aux keys for excluded module pass through handler2 = Qwen3CheckpointHandler( - num_attention_heads=64, num_key_value_heads=8, head_dim=128, - is_prequantized=True, exclude_modules={"o_proj"}, + num_attention_heads=64, + num_key_value_heads=8, + head_dim=128, + is_prequantized=True, + exclude_modules={"o_proj"}, ) sfn2 = handler2.get_skip_key_fn() assert not sfn2("model.layers.0.self_attn.o_proj.weight") @@ -300,23 +328,34 @@ def test_dense_handler_exclude_modules(self): # Empty/no exclude_modules: all quant weight keys skipped for kwargs in [{"exclude_modules": set()}, {}]: h = Qwen3CheckpointHandler( - num_attention_heads=64, num_key_value_heads=8, head_dim=128, - is_prequantized=True, **kwargs, + num_attention_heads=64, + num_key_value_heads=8, + head_dim=128, + is_prequantized=True, + **kwargs, ) assert h.get_skip_key_fn()("model.layers.0.mlp.down_proj.weight") def test_moe_handler_exclude_modules(self): """MoE handler: excluded keys pass through, non-excluded skipped; on_load_weight consistent.""" handler = Qwen3MoeCheckpointHandler( - num_experts=64, num_attention_heads=64, num_key_value_heads=8, head_dim=128, - is_prequantized=True, exclude_modules={"gate"}, + num_experts=64, + num_attention_heads=64, + num_key_value_heads=8, + head_dim=128, + is_prequantized=True, + exclude_modules={"gate"}, ) assert not handler.get_skip_key_fn()("model.layers.0.mlp.gate.weight") assert handler.get_skip_key_fn()("model.layers.0.self_attn.o_proj.weight") handler2 = Qwen3MoeCheckpointHandler( - num_experts=64, num_attention_heads=64, num_key_value_heads=8, head_dim=128, - is_prequantized=True, exclude_modules={"down_proj"}, + num_experts=64, + num_attention_heads=64, + num_key_value_heads=8, + head_dim=128, + is_prequantized=True, + exclude_modules={"down_proj"}, ) tensor = torch.randn(1024, 512) result = handler2.on_load_weight("model.layers.0.mlp.shared_expert.down_proj.weight", tensor) diff --git a/tests/qlora/test_moe_load_experts.py b/tests/qlora/test_moe_load_experts.py index 19359292..63f9afe7 100644 --- a/tests/qlora/test_moe_load_experts.py +++ b/tests/qlora/test_moe_load_experts.py @@ -6,13 +6,13 @@ """ import time + import pytest import torch -import torch.nn.functional as F -from xorl.qlora.modules.moe_experts import NvFP4QLoRAMoeExperts -from xorl.qlora.modules.moe_experts import BlockFP8QLoRAMoeExperts from xorl.ops.quantize.fp4_codec import FP4_E2M1_MAX, FP8_E4M3_MAX +from xorl.qlora.modules.moe_experts import BlockFP8QLoRAMoeExperts, NvFP4QLoRAMoeExperts + pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -23,6 +23,7 @@ # Helpers: reference (old) implementations # --------------------------------------------------------------------------- + def _ref_load_nvfp4(packed_hf_list, bs_list, gs_list, device): """Reference: per-expert CPU transpose + stack + H2D (old implementation).""" packed_out, scales_out, amax_out = [], [], [] @@ -57,6 +58,7 @@ def _ref_load_block_fp8(fp8_list, scales_list, device): # Synthetic data generators # --------------------------------------------------------------------------- + def _make_nvfp4_tensors(E, N, K, block_size=16, seed=0): """Generate synthetic HF-format NVFP4 tensors for E experts.""" torch.manual_seed(seed) @@ -87,6 +89,7 @@ def _make_block_fp8_tensors(E, N, K, block_size=128, seed=0): # Helper: invoke _load_experts via mock _load_tensor # --------------------------------------------------------------------------- + def _run_nvfp4_load_experts(module, packed_list, bs_list, gs_list): """Feed synthetic tensors through NvFP4QLoRAMoeExperts._load_experts.""" E = len(packed_list) @@ -95,7 +98,7 @@ def _run_nvfp4_load_experts(module, packed_list, bs_list, gs_list): for proj_name, hf_name in [("gate", "gate_proj"), ("up", "up_proj"), ("down", "down_proj")]: for i in range(E): fqn = f"layer.{i}.{hf_name}" - data[f"{fqn}.weight"] = packed_list[i] + data[f"{fqn}.weight"] = packed_list[i] data[f"{fqn}.weight_scale"] = bs_list[i] data[f"{fqn}.weight_scale_2"] = gs_list[i] @@ -113,7 +116,7 @@ def _run_block_fp8_load_experts(module, fp8_list, scales_list): for proj_name, hf_name in [("gate", "gate_proj"), ("up", "up_proj"), ("down", "down_proj")]: for i in range(E): fqn = f"layer.{i}.{hf_name}" - data[f"{fqn}.weight"] = fp8_list[i] + data[f"{fqn}.weight"] = fp8_list[i] data[f"{fqn}.weight_scale_inv"] = scales_list[i] module._source_fqn = "layer" @@ -127,6 +130,7 @@ def _run_block_fp8_load_experts(module, fp8_list, scales_list): # Tests: correctness vs reference # --------------------------------------------------------------------------- + class TestNvFP4LoadExperts: """NvFP4QLoRAMoeExperts._load_experts produces correct GKN buffers.""" @@ -136,9 +140,13 @@ def test_packed_matches_reference(self, E, N, K): packed_list, bs_list, gs_list = _make_nvfp4_tensors(E, N, K) module = NvFP4QLoRAMoeExperts( - num_local_experts=E, num_experts=E, - intermediate_size=N, hidden_size=K, - r=4, lora_alpha=4, device=DEVICE, + num_local_experts=E, + num_experts=E, + intermediate_size=N, + hidden_size=K, + r=4, + lora_alpha=4, + device=DEVICE, ) _run_nvfp4_load_experts(module, packed_list, bs_list, gs_list) @@ -156,9 +164,13 @@ def test_block_scales_match_reference(self, E, N, K, block_size=16): packed_list, bs_list, gs_list = _make_nvfp4_tensors(E, N, K, block_size) module = NvFP4QLoRAMoeExperts( - num_local_experts=E, num_experts=E, - intermediate_size=N, hidden_size=K, - r=4, lora_alpha=4, device=DEVICE, + num_local_experts=E, + num_experts=E, + intermediate_size=N, + hidden_size=K, + r=4, + lora_alpha=4, + device=DEVICE, ) _run_nvfp4_load_experts(module, packed_list, bs_list, gs_list) @@ -174,9 +186,13 @@ def test_ema_amax_matches_reference(self, E, N, K): packed_list, bs_list, gs_list = _make_nvfp4_tensors(E, N, K) module = NvFP4QLoRAMoeExperts( - num_local_experts=E, num_experts=E, - intermediate_size=N, hidden_size=K, - r=4, lora_alpha=4, device=DEVICE, + num_local_experts=E, + num_experts=E, + intermediate_size=N, + hidden_size=K, + r=4, + lora_alpha=4, + device=DEVICE, ) _run_nvfp4_load_experts(module, packed_list, bs_list, gs_list) @@ -192,9 +208,13 @@ def test_all_projections_loaded(self, E, N, K): packed_list, bs_list, gs_list = _make_nvfp4_tensors(E, N, K) module = NvFP4QLoRAMoeExperts( - num_local_experts=E, num_experts=E, - intermediate_size=N, hidden_size=K, - r=4, lora_alpha=4, device=DEVICE, + num_local_experts=E, + num_experts=E, + intermediate_size=N, + hidden_size=K, + r=4, + lora_alpha=4, + device=DEVICE, ) _run_nvfp4_load_experts(module, packed_list, bs_list, gs_list) @@ -207,8 +227,7 @@ def test_all_projections_loaded(self, E, N, K): assert gs_buf is not None # global scale should be 1.0 (absorbed into block scales) gs_val = module._recover_tensor(gs_buf, torch.float32) - assert torch.allclose(gs_val, torch.ones_like(gs_val)), \ - f"{proj} global_scale not 1.0 after absorption" + assert torch.allclose(gs_val, torch.ones_like(gs_val)), f"{proj} global_scale not 1.0 after absorption" @pytest.mark.parametrize("E,N,K", [(4, 768, 2048)]) def test_dequant_shape_and_dtype(self, E, N, K): @@ -216,9 +235,13 @@ def test_dequant_shape_and_dtype(self, E, N, K): packed_list, bs_list, gs_list = _make_nvfp4_tensors(E, N, K) module = NvFP4QLoRAMoeExperts( - num_local_experts=E, num_experts=E, - intermediate_size=N, hidden_size=K, - r=4, lora_alpha=4, device=DEVICE, + num_local_experts=E, + num_experts=E, + intermediate_size=N, + hidden_size=K, + r=4, + lora_alpha=4, + device=DEVICE, ) _run_nvfp4_load_experts(module, packed_list, bs_list, gs_list) @@ -236,9 +259,13 @@ def test_packed_matches_reference(self, E, N, K, block_size=128): fp8_list, scales_list = _make_block_fp8_tensors(E, N, K, block_size) module = BlockFP8QLoRAMoeExperts( - num_local_experts=E, num_experts=E, - intermediate_size=N, hidden_size=K, - r=4, lora_alpha=4, device=DEVICE, + num_local_experts=E, + num_experts=E, + intermediate_size=N, + hidden_size=K, + r=4, + lora_alpha=4, + device=DEVICE, ) _run_block_fp8_load_experts(module, fp8_list, scales_list) @@ -254,9 +281,13 @@ def test_scales_match_reference(self, E, N, K, block_size=128): fp8_list, scales_list = _make_block_fp8_tensors(E, N, K, block_size) module = BlockFP8QLoRAMoeExperts( - num_local_experts=E, num_experts=E, - intermediate_size=N, hidden_size=K, - r=4, lora_alpha=4, device=DEVICE, + num_local_experts=E, + num_experts=E, + intermediate_size=N, + hidden_size=K, + r=4, + lora_alpha=4, + device=DEVICE, ) _run_block_fp8_load_experts(module, fp8_list, scales_list) @@ -272,9 +303,13 @@ def test_all_projections_loaded(self, E, N, K): fp8_list, scales_list = _make_block_fp8_tensors(E, N, K) module = BlockFP8QLoRAMoeExperts( - num_local_experts=E, num_experts=E, - intermediate_size=N, hidden_size=K, - r=4, lora_alpha=4, device=DEVICE, + num_local_experts=E, + num_experts=E, + intermediate_size=N, + hidden_size=K, + r=4, + lora_alpha=4, + device=DEVICE, ) _run_block_fp8_load_experts(module, fp8_list, scales_list) @@ -287,15 +322,20 @@ def test_all_projections_loaded(self, E, N, K): # Benchmark (not a pytest test — run directly or with -s -k bench) # --------------------------------------------------------------------------- + def _bench_nvfp4_load(E=128, N=768, K=2048, n_reps=5): """Benchmark NvFP4 expert loading: new GPU-transpose vs old CPU-transpose.""" packed_list, bs_list, gs_list = _make_nvfp4_tensors(E, N, K) def make_module(): return NvFP4QLoRAMoeExperts( - num_local_experts=E, num_experts=E, - intermediate_size=N, hidden_size=K, - r=4, lora_alpha=4, device=DEVICE, + num_local_experts=E, + num_experts=E, + intermediate_size=N, + hidden_size=K, + r=4, + lora_alpha=4, + device=DEVICE, ) # New implementation (GPU-transpose) @@ -335,9 +375,13 @@ def make_module(): for E, N, K in [(4, 768, 2048), (8, 512, 1024), (1, 256, 512)]: packed_list, bs_list, gs_list = _make_nvfp4_tensors(E, N, K) module = NvFP4QLoRAMoeExperts( - num_local_experts=E, num_experts=E, - intermediate_size=N, hidden_size=K, - r=4, lora_alpha=4, device=DEVICE, + num_local_experts=E, + num_experts=E, + intermediate_size=N, + hidden_size=K, + r=4, + lora_alpha=4, + device=DEVICE, ) _run_nvfp4_load_experts(module, packed_list, bs_list, gs_list) ref_p, ref_s, _ = _ref_load_nvfp4(packed_list, bs_list, gs_list, DEVICE) @@ -348,9 +392,13 @@ def make_module(): for E, N, K in [(4, 768, 2048), (8, 512, 1024)]: fp8_list, scales_list = _make_block_fp8_tensors(E, N, K) module = BlockFP8QLoRAMoeExperts( - num_local_experts=E, num_experts=E, - intermediate_size=N, hidden_size=K, - r=4, lora_alpha=4, device=DEVICE, + num_local_experts=E, + num_experts=E, + intermediate_size=N, + hidden_size=K, + r=4, + lora_alpha=4, + device=DEVICE, ) _run_block_fp8_load_experts(module, fp8_list, scales_list) ref_p, ref_s = _ref_load_block_fp8(fp8_list, scales_list, DEVICE) diff --git a/tests/qlora/test_qlora.py b/tests/qlora/test_qlora.py index 51016e9c..7749a4d4 100644 --- a/tests/qlora/test_qlora.py +++ b/tests/qlora/test_qlora.py @@ -5,13 +5,14 @@ import torch.nn as nn import torch.nn.functional as F -from xorl.qlora.modules.linear import QLoRALinear +from xorl.ops.quantize import block_fp8_quantize_gkn as block_fp8_weight_quant +from xorl.ops.quantize import nvfp4_dequantize, nvfp4_quantize +from xorl.ops.quantize.fp4_codec import FP4_E2M1_MAX, FP8_E4M3_MAX from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear +from xorl.qlora.modules.linear import QLoRALinear from xorl.qlora.utils import inject_qlora_into_model, maybe_requant_qlora -from xorl.ops.quantize import nvfp4_quantize, nvfp4_dequantize -from xorl.ops.quantize.fp4_codec import FP4_E2M1_MAX, FP8_E4M3_MAX -from xorl.ops.quantize import block_fp8_quantize_gkn as block_fp8_weight_quant, block_fp8_dequantize_gkn as block_fp8_weight_dequant -from xorl.trainers.training_utils import reset_lora_optimizer_states, maybe_merge_lora +from xorl.trainers.training_utils import maybe_merge_lora, reset_lora_optimizer_states + pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -20,14 +21,17 @@ # Helpers # --------------------------------------------------------------------------- + def _make_model(): """Simple model with named linear layers for testing.""" + class MLP(nn.Module): def __init__(self): super().__init__() self.gate_proj = nn.Linear(256, 512, bias=False) self.up_proj = nn.Linear(256, 512, bias=False) self.down_proj = nn.Linear(512, 256, bias=False) + def forward(self, x): return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) @@ -38,6 +42,7 @@ def __init__(self): self.k_proj = nn.Linear(256, 256, bias=False) self.v_proj = nn.Linear(256, 256, bias=False) self.o_proj = nn.Linear(256, 256, bias=False) + def forward(self, x): q = self.q_proj(x) k = self.k_proj(x) @@ -54,6 +59,7 @@ def __init__(self): super().__init__() self.self_attn = Attn() self.mlp = MLP() + def forward(self, x): return x + self.mlp(self.self_attn(x)) @@ -61,6 +67,7 @@ class Model(nn.Module): def __init__(self): super().__init__() self.layers = nn.ModuleList([Layer()]) + def forward(self, x): for layer in self.layers: x = layer(x) @@ -76,6 +83,7 @@ def _quantize_injected_model(model): This helper fills them with properly quantized random weights for testing. """ from xorl.qlora.modules.linear import QLoRALinear + for m in model.modules(): if isinstance(m, QLoRALinear) and m._is_prequantized: w = torch.randn(m.out_features, m.in_features, device="cuda", dtype=torch.bfloat16) @@ -93,13 +101,13 @@ def _make_fp8_data(out_features, in_features, block_size=128): # 1. Quantize bf16 -> packed_f32: forward, backward, memory, dequant (both formats) # --------------------------------------------------------------------------- + @pytest.mark.parametrize("quant_format,gs", [("nvfp4", 16), ("block_fp8", 128)]) def test_quantize_forward_backward_memory(quant_format, gs): """bf16 quantization, forward, backward (only LoRA gets grad), memory savings.""" linear = nn.Linear(256, 512, bias=False, device="cuda", dtype=torch.bfloat16) bf16_bytes = linear.weight.numel() * 2 - qlora = QLoRALinear.from_module(linear, r=16, lora_alpha=16, - quant_format=quant_format, quant_group_size=gs) + qlora = QLoRALinear.from_module(linear, r=16, lora_alpha=16, quant_format=quant_format, quant_group_size=gs) # Quantized storage assert qlora.packed_weight_f32 is not None @@ -129,8 +137,7 @@ def test_dequantize_roundtrip(): """Dequantized weight should be close to original.""" linear = nn.Linear(256, 512, bias=False, device="cuda", dtype=torch.bfloat16) w_orig = linear.weight.detach().clone() - qlora = QLoRALinear.from_module(linear, r=16, lora_alpha=16, - quant_format="nvfp4", quant_group_size=16) + qlora = QLoRALinear.from_module(linear, r=16, lora_alpha=16, quant_format="nvfp4", quant_group_size=16) w_deq = qlora._dequantize_weight().to(torch.bfloat16) assert torch.allclose(w_orig, w_deq, atol=0.05, rtol=0.05) @@ -139,15 +146,21 @@ def test_dequantize_roundtrip(): # 2. Pre-quantized nvfp4 loading # --------------------------------------------------------------------------- + def test_prequantized_nvfp4_loading(): """from_quantized() loads pre-packed nvfp4 weights; forward+backward work.""" w = torch.randn(512, 256, device="cuda", dtype=torch.bfloat16) packed, block_scales, global_scale = nvfp4_quantize(w, 16) qlora = QLoRALinear.from_quantized( - packed_weight=packed, weight_block_scales=block_scales, - weight_global_scale=global_scale, in_features=256, out_features=512, - quant_format="nvfp4", quant_group_size=16, device="cuda", + packed_weight=packed, + weight_block_scales=block_scales, + weight_global_scale=global_scale, + in_features=256, + out_features=512, + quant_format="nvfp4", + quant_group_size=16, + device="cuda", ) assert qlora.packed_weight_f32.dtype == torch.float32 @@ -162,14 +175,14 @@ def test_prequantized_nvfp4_loading(): # 3. EMA amax + NVFP4 scale convention # --------------------------------------------------------------------------- + def test_ema_amax_and_scale_convention(): """EMA amax: init from bf16, update on merge, global_scale formula. Scale convention: block_scales use full fp8 range, dequant roundtrip accuracy.""" linear = nn.Linear(256, 512, bias=False, device="cuda", dtype=torch.bfloat16) expected_amax = linear.weight.float().abs().max().item() - qlora = QLoRALinear.from_module(linear, r=16, lora_alpha=16, - quant_format="nvfp4", quant_group_size=16) + qlora = QLoRALinear.from_module(linear, r=16, lora_alpha=16, quant_format="nvfp4", quant_group_size=16) # Init from bf16 assert qlora._ema_amax is not None and qlora._ema_amax.shape == (1,) @@ -184,9 +197,7 @@ def test_ema_amax_and_scale_convention(): assert qlora._ema_amax.item() != amax_before # Global scale reflects EMA - gs = qlora._recover_tensor( - qlora.weight_global_scale, qlora._scale_dtypes["weight_global_scale"] - ).item() + gs = qlora._recover_tensor(qlora.weight_global_scale, qlora._scale_dtypes["weight_global_scale"]).item() expected_gs = qlora._ema_amax.item() / (FP4_E2M1_MAX * FP8_E4M3_MAX) assert abs(gs - expected_gs) / max(abs(expected_gs), 1e-12) < 0.01 @@ -211,11 +222,11 @@ def test_ema_amax_and_scale_convention(): # 4. Merge weights + maybe_requant # --------------------------------------------------------------------------- + def test_merge_weights_and_requant(): """merge_weights folds LoRA into base; maybe_requant_qlora merges+resets+EMA updates.""" linear = nn.Linear(256, 512, bias=False, device="cuda", dtype=torch.bfloat16) - qlora = QLoRALinear.from_module(linear, r=16, lora_alpha=16, - quant_format="nvfp4", quant_group_size=16) + qlora = QLoRALinear.from_module(linear, r=16, lora_alpha=16, quant_format="nvfp4", quant_group_size=16) # Merge weights with torch.no_grad(): @@ -230,8 +241,7 @@ def test_merge_weights_and_requant(): # Requant incorporates LoRA delta linear2 = nn.Linear(256, 512, bias=False, device="cuda", dtype=torch.bfloat16) - qlora2 = QLoRALinear.from_module(linear2, r=16, lora_alpha=16, - quant_format="nvfp4", quant_group_size=16) + qlora2 = QLoRALinear.from_module(linear2, r=16, lora_alpha=16, quant_format="nvfp4", quant_group_size=16) with torch.no_grad(): qlora2.lora_A.fill_(0.05) qlora2.lora_B.fill_(0.05) @@ -248,8 +258,7 @@ def test_merge_weights_and_requant(): # maybe_requant_qlora linear3 = nn.Linear(256, 512, bias=False, device="cuda", dtype=torch.bfloat16) - qlora3 = QLoRALinear.from_module(linear3, r=16, lora_alpha=16, - quant_format="nvfp4", quant_group_size=16) + qlora3 = QLoRALinear.from_module(linear3, r=16, lora_alpha=16, quant_format="nvfp4", quant_group_size=16) amax_before = qlora3._ema_amax.item() with torch.no_grad(): qlora3.lora_A.fill_(1.0) @@ -267,12 +276,18 @@ def test_merge_weights_and_requant(): # 5. Injection + training # --------------------------------------------------------------------------- + def test_injection_and_training(): """inject_qlora replaces target modules; forward/backward work; loss decreases.""" model = _make_model() - inject_qlora_into_model(model, r=16, lora_alpha=16, - quant_format="nvfp4", quant_group_size=16, - target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]) + inject_qlora_into_model( + model, + r=16, + lora_alpha=16, + quant_format="nvfp4", + quant_group_size=16, + target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], + ) _quantize_injected_model(model) layer = model.layers[0] assert isinstance(layer.self_attn.q_proj, QLoRALinear) @@ -290,8 +305,7 @@ def test_injection_and_training(): def test_training_step_converges(quant_format, gs): """Loss decreases over training steps for both quant formats.""" model = _make_model() - inject_qlora_into_model(model, r=16, lora_alpha=16, - quant_format=quant_format, quant_group_size=gs) + inject_qlora_into_model(model, r=16, lora_alpha=16, quant_format=quant_format, quant_group_size=gs) _quantize_injected_model(model) trainable = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.AdamW(trainable, lr=1e-3) @@ -312,11 +326,11 @@ def test_training_step_converges(quant_format, gs): # 6. Step-based requant training (end-to-end) # --------------------------------------------------------------------------- + def test_step_based_requant_training(): """Train with periodic requant: loss decreases, requant triggered, continues after reset.""" model = _make_model() - inject_qlora_into_model(model, r=16, lora_alpha=16, - quant_format="nvfp4", quant_group_size=16) + inject_qlora_into_model(model, r=16, lora_alpha=16, quant_format="nvfp4", quant_group_size=16) _quantize_injected_model(model) model.train() trainable = [p for p in model.parameters() if p.requires_grad] @@ -340,8 +354,7 @@ def test_step_based_requant_training(): # Also verify single-module requant continues training after reset linear = nn.Linear(256, 512, bias=False, device="cuda", dtype=torch.bfloat16) - qlora = QLoRALinear.from_module(linear, r=16, lora_alpha=16, - quant_format="nvfp4", quant_group_size=16) + qlora = QLoRALinear.from_module(linear, r=16, lora_alpha=16, quant_format="nvfp4", quant_group_size=16) qlora.train() opt = torch.optim.AdamW([qlora.lora_A, qlora.lora_B], lr=1e-3) x2 = torch.randn(4, 16, 256, device="cuda", dtype=torch.bfloat16) @@ -363,8 +376,11 @@ def test_inject_with_checkpoint_quant_format(): """inject_qlora with checkpoint_quant_format sets _source_quant_format.""" model = _make_model() inject_qlora_into_model( - model, r=16, lora_alpha=16, - quant_format="block_fp8", quant_group_size=128, + model, + r=16, + lora_alpha=16, + quant_format="block_fp8", + quant_group_size=128, checkpoint_quant_format="block_fp8", ) for m in model.modules(): @@ -377,6 +393,7 @@ def test_inject_with_checkpoint_quant_format(): # 8. Pre-quantized block FP8 loading (HF checkpoint path) # --------------------------------------------------------------------------- + def test_prequantized_block_fp8_load_and_forward(): """Load FP8 single module + merged qkv; forward/backward work; dequant roundtrip.""" M, K = 256, 256 @@ -483,11 +500,11 @@ def test_prequantized_block_fp8_merge_and_training(): # 9. ReLoRA optimizer reset # --------------------------------------------------------------------------- + def test_reset_lora_optimizer_states_clears(): """Verify ReLoRA reset clears optimizer states for LoRA params.""" linear = nn.Linear(256, 512, bias=False, device="cuda", dtype=torch.bfloat16) - qlora = QLoRALinear.from_module(linear, r=16, lora_alpha=16, - quant_format="block_fp8", quant_group_size=128) + qlora = QLoRALinear.from_module(linear, r=16, lora_alpha=16, quant_format="block_fp8", quant_group_size=128) qlora.train() opt = torch.optim.AdamW([qlora.lora_A, qlora.lora_B], lr=1e-3) @@ -525,9 +542,14 @@ def test_reset_lora_optimizer_states_clears(): def test_reset_ignores_non_lora_params(): """Optimizer reset only touches LoRA params, not other trainable params.""" model = _make_model() - inject_qlora_into_model(model, r=16, lora_alpha=16, - quant_format="block_fp8", quant_group_size=128, - target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]) + inject_qlora_into_model( + model, + r=16, + lora_alpha=16, + quant_format="block_fp8", + quant_group_size=128, + target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], + ) _quantize_injected_model(model) # MLP layers are nn.Linear (not QLoRA), so add them as trainable too all_params = list(model.parameters()) @@ -553,8 +575,9 @@ def test_reset_ignores_non_lora_params(): for name, p in model.named_parameters(): if name in non_lora_states_before: torch.testing.assert_close( - opt.state[p]["exp_avg"], non_lora_states_before[name], - msg=f"Non-LoRA param {name} was modified by reset" + opt.state[p]["exp_avg"], + non_lora_states_before[name], + msg=f"Non-LoRA param {name} was modified by reset", ) @@ -584,8 +607,7 @@ def _train(merge_interval, reset_opt, total_steps, r=4, lr=1e-3): linear = nn.Linear(dim, dim, bias=False, device="cuda", dtype=torch.bfloat16) with torch.no_grad(): linear.weight.copy_(W_base) - qlora = QLoRALinear.from_module(linear, r=r, lora_alpha=r, - quant_format="block_fp8", quant_group_size=128) + qlora = QLoRALinear.from_module(linear, r=r, lora_alpha=r, quant_format="block_fp8", quant_group_size=128) qlora.train() opt = torch.optim.AdamW([qlora.lora_A, qlora.lora_B], lr=lr) model = nn.ModuleList([qlora]) @@ -609,13 +631,15 @@ def _train(merge_interval, reset_opt, total_steps, r=4, lr=1e-3): losses_merge_reset = _train(merge_interval=40, reset_opt=True, total_steps=200) # Merge + reset should converge - assert losses_merge_reset[-1] < losses_merge_reset[0], \ + assert losses_merge_reset[-1] < losses_merge_reset[0], ( f"Merge+reset didn't converge: {losses_merge_reset[0]:.4f} -> {losses_merge_reset[-1]:.4f}" + ) # Merge + reset should reach lower loss than no-merge # (accumulated rank from 5 merges of rank-4 > single rank-4) - assert losses_merge_reset[-1] < losses_no_merge[-1], \ + assert losses_merge_reset[-1] < losses_no_merge[-1], ( f"Merge+reset ({losses_merge_reset[-1]:.4f}) should beat no-merge ({losses_no_merge[-1]:.4f})" + ) @pytest.mark.parametrize("quant_format,gs", [("nvfp4", 16), ("block_fp8", 128)]) @@ -624,8 +648,7 @@ def test_merge_with_optimizer_reset_still_converges(quant_format, gs): torch.manual_seed(42) dim = 256 linear = nn.Linear(dim, dim, bias=False, device="cuda", dtype=torch.bfloat16) - qlora = QLoRALinear.from_module(linear, r=8, lora_alpha=8, - quant_format=quant_format, quant_group_size=gs) + qlora = QLoRALinear.from_module(linear, r=8, lora_alpha=8, quant_format=quant_format, quant_group_size=gs) qlora.train() opt = torch.optim.AdamW([qlora.lora_A, qlora.lora_B], lr=1e-3) model = nn.ModuleList([qlora]) @@ -641,8 +664,7 @@ def test_merge_with_optimizer_reset_still_converges(quant_format, gs): if step % 20 == 0: maybe_requant_qlora(model) reset_lora_optimizer_states(model, opt) - assert losses[-1] < losses[0], \ - f"{quant_format} didn't converge: {losses[0]:.4f} -> {losses[-1]:.4f}" + assert losses[-1] < losses[0], f"{quant_format} didn't converge: {losses[0]:.4f} -> {losses[-1]:.4f}" def test_merge_with_reset_vs_without_reset(): @@ -665,8 +687,7 @@ def _train(reset_opt, total_steps=120, merge_interval=30): linear = nn.Linear(dim, dim, bias=False, device="cuda", dtype=torch.bfloat16) with torch.no_grad(): linear.weight.copy_(W_base) - qlora = QLoRALinear.from_module(linear, r=8, lora_alpha=8, - quant_format="block_fp8", quant_group_size=128) + qlora = QLoRALinear.from_module(linear, r=8, lora_alpha=8, quant_format="block_fp8", quant_group_size=128) qlora.train() opt = torch.optim.AdamW([qlora.lora_A, qlora.lora_B], lr=1e-3) model = nn.ModuleList([qlora]) @@ -687,21 +708,23 @@ def _train(reset_opt, total_steps=120, merge_interval=30): losses_merge_reset = _train(reset_opt=True) # Both should converge (loss decreases) - assert losses_merge_only[-1] < losses_merge_only[0], \ + assert losses_merge_only[-1] < losses_merge_only[0], ( f"Merge-only didn't converge: {losses_merge_only[0]:.4f} -> {losses_merge_only[-1]:.4f}" - assert losses_merge_reset[-1] < losses_merge_reset[0], \ + ) + assert losses_merge_reset[-1] < losses_merge_reset[0], ( f"Merge+reset didn't converge: {losses_merge_reset[0]:.4f} -> {losses_merge_reset[-1]:.4f}" + ) # Reset should not catastrophically hurt (within 2x of merge-only final loss) - assert losses_merge_reset[-1] < losses_merge_only[-1] * 2.0, \ + assert losses_merge_reset[-1] < losses_merge_only[-1] * 2.0, ( f"Reset hurt too much: {losses_merge_reset[-1]:.4f} vs merge-only {losses_merge_only[-1]:.4f}" + ) def test_maybe_merge_lora_with_optimizer_reset_integration(): """End-to-end test of maybe_merge_lora with reset_optimizer=True.""" model = _make_model() - inject_qlora_into_model(model, r=16, lora_alpha=16, - quant_format="block_fp8", quant_group_size=128) + inject_qlora_into_model(model, r=16, lora_alpha=16, quant_format="block_fp8", quant_group_size=128) _quantize_injected_model(model) model.train() trainable = [p for p in model.parameters() if p.requires_grad] @@ -718,9 +741,13 @@ def test_maybe_merge_lora_with_optimizer_reset_integration(): optimizer.step() losses.append(loss.item()) maybe_merge_lora( - model, enable_lora=False, enable_qlora=True, - merge_interval=10, global_step=step, - optimizer=optimizer, reset_optimizer=True, + model, + enable_lora=False, + enable_qlora=True, + merge_interval=10, + global_step=step, + optimizer=optimizer, + reset_optimizer=True, ) # Should still converge despite optimizer resets diff --git a/tests/qlora/test_quantize_error_reduction.py b/tests/qlora/test_quantize_error_reduction.py index ac6ea9e6..d1f9ced6 100644 --- a/tests/qlora/test_quantize_error_reduction.py +++ b/tests/qlora/test_quantize_error_reduction.py @@ -1,7 +1,6 @@ """Tests for QLoRA quantization error reduction techniques.""" import pytest -import torch -import torch.nn as nn + pytestmark = [pytest.mark.gpu] diff --git a/tests/server/api_server/test_api_server.py b/tests/server/api_server/test_api_server.py index a21ab21f..112117fd 100644 --- a/tests/server/api_server/test_api_server.py +++ b/tests/server/api_server/test_api_server.py @@ -7,18 +7,19 @@ import pytest + pytestmark = [pytest.mark.cpu, pytest.mark.server] -from xorl.server.api_server.server import APIServer from xorl.server.api_server.api_types import ( + AdamParams, Datum, DatumInput, ForwardBackwardRequest, + LoadWeightsRequest, OptimStepRequest, SaveWeightsRequest, - LoadWeightsRequest, - AdamParams, ) +from xorl.server.api_server.server import APIServer class TestAPIServerConfiguration: @@ -69,10 +70,12 @@ def test_request_creation_and_serialization(self): # Multiple samples with default model_id request = ForwardBackwardRequest( - forward_backward_input=DatumInput(data=[ - Datum(model_input={"input_ids": [1, 2]}, loss_fn_inputs={"labels": [2, 3]}), - Datum(model_input={"input_ids": [3, 4]}, loss_fn_inputs={"labels": [4, 5]}), - ]), + forward_backward_input=DatumInput( + data=[ + Datum(model_input={"input_ids": [1, 2]}, loss_fn_inputs={"labels": [2, 3]}), + Datum(model_input={"input_ids": [3, 4]}, loss_fn_inputs={"labels": [4, 5]}), + ] + ), ) assert len(request.forward_backward_input.data) == 2 assert request.model_id == "default" diff --git a/tests/server/api_server/test_api_types.py b/tests/server/api_server/test_api_types.py index 451aca10..b2f1276a 100644 --- a/tests/server/api_server/test_api_types.py +++ b/tests/server/api_server/test_api_types.py @@ -6,27 +6,28 @@ import pytest + pytestmark = [pytest.mark.cpu, pytest.mark.server] from pydantic import ValidationError from xorl.server.api_server.api_types import ( + AdamParams, Datum, DatumInput, - LossFnOutput, - AdamParams, + ErrorResponse, ForwardBackwardRequest, ForwardBackwardResponse, - OptimStepRequest, - OptimStepResponse, - SaveWeightsRequest, - SaveWeightsResponse, + HealthCheckResponse, LoadWeightsRequest, LoadWeightsResponse, + LossFnOutput, + OptimStepRequest, + OptimStepResponse, SaveWeightsForSamplerRequest, SaveWeightsForSamplerResponse, - HealthCheckResponse, - ErrorResponse, + SaveWeightsRequest, + SaveWeightsResponse, ) @@ -112,7 +113,8 @@ def test_datum_forward_backward_request_and_response(self): response = ForwardBackwardResponse( loss_fn_output_type="multi_loss", loss_fn_outputs=[LossFnOutput(loss=2.0), LossFnOutput(loss=3.0)], - metrics={}, info={}, + metrics={}, + info={}, ) assert len(response.loss_fn_outputs) == 2 @@ -138,7 +140,9 @@ def test_optim_step_types(self): assert defaults.eps == 1e-12 request = OptimStepRequest( - model_id="test-model", adam_params=AdamParams(learning_rate=1e-4), gradient_clip=1.0, + model_id="test-model", + adam_params=AdamParams(learning_rate=1e-4), + gradient_clip=1.0, ) assert request.model_id == "test-model" assert request.adam_params.learning_rate == 1e-4 @@ -150,7 +154,8 @@ def test_optim_step_types(self): assert request.gradient_clip is None response = OptimStepResponse( - metrics={"grad_norm": 1.234, "learning_rate": 1e-4, "step": 100}, info={}, + metrics={"grad_norm": 1.234, "learning_rate": 1e-4, "step": 100}, + info={}, ) assert response.metrics["grad_norm"] == 1.234 response = OptimStepResponse(metrics={}, info={}) @@ -201,12 +206,18 @@ def test_weights_health_error_and_serialization(self): # HealthCheckResponse response = HealthCheckResponse( - status="healthy", engine_running=True, active_requests=5, total_requests=100, + status="healthy", + engine_running=True, + active_requests=5, + total_requests=100, ) assert response.status == "healthy" assert response.engine_running is True response = HealthCheckResponse( - status="unhealthy", engine_running=False, active_requests=0, total_requests=0, + status="unhealthy", + engine_running=False, + active_requests=0, + total_requests=0, ) assert response.engine_running is False @@ -242,7 +253,10 @@ def test_weights_health_error_and_serialization(self): # HealthCheckResponse response = HealthCheckResponse( - status="healthy", engine_running=True, active_requests=0, total_requests=10, + status="healthy", + engine_running=True, + active_requests=0, + total_requests=10, ) data = response.model_dump() response2 = HealthCheckResponse(**data) diff --git a/tests/server/api_server/test_checkpoint_paths.py b/tests/server/api_server/test_checkpoint_paths.py index 86f1a35c..597ed2dc 100644 --- a/tests/server/api_server/test_checkpoint_paths.py +++ b/tests/server/api_server/test_checkpoint_paths.py @@ -13,13 +13,13 @@ import pytest -from xorl.server.api_server.server import APIServer, validate_model_id from xorl.server.api_server.api_types import ( - SaveWeightsRequest, - LoadWeightsRequest, - ListCheckpointsRequest, DeleteCheckpointRequest, + ListCheckpointsRequest, + LoadWeightsRequest, + SaveWeightsRequest, ) +from xorl.server.api_server.server import APIServer, validate_model_id class TestTomiUriAndPathConstruction: @@ -43,9 +43,7 @@ def test_xorl_uri_construction_and_parsing(self): assert uri == "xorl://user_123/weights/checkpoint-001" # Parse xorl:// URI - model_id, checkpoint_name, _ = self.server._from_xorl_uri( - "xorl://user_123/weights/checkpoint-001" - ) + model_id, checkpoint_name, _ = self.server._from_xorl_uri("xorl://user_123/weights/checkpoint-001") assert model_id == "user_123" assert checkpoint_name == "checkpoint-001" @@ -94,10 +92,14 @@ def test_save_load_all_formats_errors_and_auto_checkpoint(self): # --- Save creates correct structure --- model_id = "user_123" - mock_output.outputs = [{"checkpoint_path": f"{self.temp_dir}/weights/{model_id}/my_checkpoint", "success": True}] - with patch.object(self.server, '_wait_for_response', return_value=mock_output): + mock_output.outputs = [ + {"checkpoint_path": f"{self.temp_dir}/weights/{model_id}/my_checkpoint", "success": True} + ] + with patch.object(self.server, "_wait_for_response", return_value=mock_output): self.server.orchestrator_client.send_request = AsyncMock(return_value=AsyncMock()) - response = asyncio.run(self.server.save_weights(SaveWeightsRequest(model_id=model_id, path="my_checkpoint"))) + response = asyncio.run( + self.server.save_weights(SaveWeightsRequest(model_id=model_id, path="my_checkpoint")) + ) assert model_id in response.path and "my_checkpoint" in response.path # --- Save rejects existing checkpoint (409) --- @@ -115,37 +117,47 @@ def test_save_load_all_formats_errors_and_auto_checkpoint(self): # xorl:// URI os.makedirs(os.path.join(self.temp_dir, "weights", "user_abc", "ckpt-001"), exist_ok=True) - with patch.object(self.server, '_wait_for_response', return_value=mock_output): + with patch.object(self.server, "_wait_for_response", return_value=mock_output): self.server.orchestrator_client.send_request = AsyncMock() - response = asyncio.run(self.server.load_weights(LoadWeightsRequest(model_id="user_abc", path="xorl://user_abc/weights/ckpt-001"))) + response = asyncio.run( + self.server.load_weights( + LoadWeightsRequest(model_id="user_abc", path="xorl://user_abc/weights/ckpt-001") + ) + ) assert response.path == "xorl://user_abc/weights/ckpt-001" # weights/model_id/name (cross-model) os.makedirs(os.path.join(self.temp_dir, "weights", "default", "000000"), exist_ok=True) - with patch.object(self.server, '_wait_for_response', return_value=mock_output): + with patch.object(self.server, "_wait_for_response", return_value=mock_output): self.server.orchestrator_client.send_request = AsyncMock() - response = asyncio.run(self.server.load_weights(LoadWeightsRequest(model_id="run-2", path="weights/default/000000"))) + response = asyncio.run( + self.server.load_weights(LoadWeightsRequest(model_id="run-2", path="weights/default/000000")) + ) assert response.path == "weights/default/000000" # model_id/checkpoint format os.makedirs(os.path.join(self.temp_dir, "weights", "other_run", "step-100"), exist_ok=True) - with patch.object(self.server, '_wait_for_response', return_value=mock_output): + with patch.object(self.server, "_wait_for_response", return_value=mock_output): self.server.orchestrator_client.send_request = AsyncMock() - response = asyncio.run(self.server.load_weights(LoadWeightsRequest(model_id="my_run", path="other_run/step-100"))) + response = asyncio.run( + self.server.load_weights(LoadWeightsRequest(model_id="my_run", path="other_run/step-100")) + ) assert response.path == "other_run/step-100" # checkpoint name only os.makedirs(os.path.join(self.temp_dir, "weights", "my_model", "ckpt-002"), exist_ok=True) - with patch.object(self.server, '_wait_for_response', return_value=mock_output): + with patch.object(self.server, "_wait_for_response", return_value=mock_output): self.server.orchestrator_client.send_request = AsyncMock() response = asyncio.run(self.server.load_weights(LoadWeightsRequest(model_id="my_model", path="ckpt-002"))) assert response.path == "ckpt-002" # legacy weights/checkpoint format os.makedirs(os.path.join(self.temp_dir, "weights", "some_model", "000000"), exist_ok=True) - with patch.object(self.server, '_wait_for_response', return_value=mock_output): + with patch.object(self.server, "_wait_for_response", return_value=mock_output): self.server.orchestrator_client.send_request = AsyncMock() - response = asyncio.run(self.server.load_weights(LoadWeightsRequest(model_id="some_model", path="weights/000000"))) + response = asyncio.run( + self.server.load_weights(LoadWeightsRequest(model_id="some_model", path="weights/000000")) + ) assert response.path == "weights/000000" # --- Load errors: 404 and isolation --- @@ -155,9 +167,11 @@ def test_save_load_all_formats_errors_and_auto_checkpoint(self): # Explicit path uses path's model_id os.makedirs(os.path.join(self.temp_dir, "weights", "user_a", "step-50"), exist_ok=True) - with patch.object(self.server, '_wait_for_response', return_value=mock_output): + with patch.object(self.server, "_wait_for_response", return_value=mock_output): self.server.orchestrator_client.send_request = AsyncMock() - response = asyncio.run(self.server.load_weights(LoadWeightsRequest(model_id="user_b", path="weights/user_a/step-50"))) + response = asyncio.run( + self.server.load_weights(LoadWeightsRequest(model_id="user_b", path="weights/user_a/step-50")) + ) assert response.path == "weights/user_a/step-50" with pytest.raises(HTTPException) as exc_info: @@ -166,7 +180,7 @@ def test_save_load_all_formats_errors_and_auto_checkpoint(self): # --- Auto-save 000000 checkpoint --- mock_output.outputs = [{"checkpoint_path": f"{self.temp_dir}/weights/test_model/000000"}] - with patch.object(self.server, '_wait_for_response', return_value=mock_output): + with patch.object(self.server, "_wait_for_response", return_value=mock_output): self.server.orchestrator_client.send_request = AsyncMock() response = asyncio.run(self.server.save_weights(SaveWeightsRequest(model_id="test_model", path="000000"))) assert "000000" in response.path @@ -174,7 +188,7 @@ def test_save_load_all_formats_errors_and_auto_checkpoint(self): # Load 000000 after save os.makedirs(os.path.join(self.temp_dir, "weights", "test_model", "000000"), exist_ok=True) mock_output.outputs = [{"success": True}] - with patch.object(self.server, '_wait_for_response', return_value=mock_output): + with patch.object(self.server, "_wait_for_response", return_value=mock_output): self.server.orchestrator_client.send_request = AsyncMock() response = asyncio.run(self.server.load_weights(LoadWeightsRequest(model_id="test_model", path="000000"))) assert response.path == "000000" @@ -217,15 +231,21 @@ def test_list_delete_and_isolation(self): # --- Delete: success, reserved blocked, invalid format --- ckpt_path = os.path.join(self.temp_dir, "weights", "del_test", "to_delete") os.makedirs(ckpt_path, exist_ok=True) - response = asyncio.run(self.server.delete_checkpoint( - DeleteCheckpointRequest(model_id="del_test", checkpoint_id="weights/del_test/to_delete"))) + response = asyncio.run( + self.server.delete_checkpoint( + DeleteCheckpointRequest(model_id="del_test", checkpoint_id="weights/del_test/to_delete") + ) + ) assert response.success is True and not os.path.exists(ckpt_path) reserved = os.path.join(self.temp_dir, "weights", "res_test", "000000") os.makedirs(reserved, exist_ok=True) with pytest.raises(HTTPException) as exc_info: - asyncio.run(self.server.delete_checkpoint( - DeleteCheckpointRequest(model_id="res_test", checkpoint_id="weights/res_test/000000"))) + asyncio.run( + self.server.delete_checkpoint( + DeleteCheckpointRequest(model_id="res_test", checkpoint_id="weights/res_test/000000") + ) + ) assert exc_info.value.status_code == 403 and os.path.exists(reserved) with pytest.raises(HTTPException) as exc_info: @@ -267,8 +287,17 @@ def test_valid_invalid_defaults_and_traversal(self): assert validate_model_id(None) == "default" # Invalid - for mid in ["../etc/passwd", "user/name", "user\\name", "user name", "user@name", - "user.name", "_start", "-start", "a" * 129]: + for mid in [ + "../etc/passwd", + "user/name", + "user\\name", + "user name", + "user@name", + "user.name", + "_start", + "-start", + "a" * 129, + ]: with pytest.raises(HTTPException) as exc_info: validate_model_id(mid) assert exc_info.value.status_code == 400 @@ -315,13 +344,19 @@ def test_sampler_listing_deletion_and_adapter_tracking(self): # --- Sampler deletion --- sp = os.path.join(self.temp_dir, "sampler_weights", "to-delete") os.makedirs(sp, exist_ok=True) - response = asyncio.run(self.server.delete_checkpoint( - DeleteCheckpointRequest(model_id="any", checkpoint_id="sampler_weights/to-delete"))) + response = asyncio.run( + self.server.delete_checkpoint( + DeleteCheckpointRequest(model_id="any", checkpoint_id="sampler_weights/to-delete") + ) + ) assert response.success is True and not os.path.exists(sp) with pytest.raises(HTTPException) as exc_info: - asyncio.run(self.server.delete_checkpoint( - DeleteCheckpointRequest(model_id="any", checkpoint_id="sampler_weights/nonexistent"))) + asyncio.run( + self.server.delete_checkpoint( + DeleteCheckpointRequest(model_id="any", checkpoint_id="sampler_weights/nonexistent") + ) + ) assert exc_info.value.status_code == 404 # --- Path resolution --- diff --git a/tests/server/api_server/test_future_store.py b/tests/server/api_server/test_future_store.py index 2eecfa68..b2337859 100644 --- a/tests/server/api_server/test_future_store.py +++ b/tests/server/api_server/test_future_store.py @@ -3,17 +3,19 @@ """ import asyncio -import pytest import time +import pytest + + pytestmark = [pytest.mark.cpu, pytest.mark.server] from xorl.server.api_server.future_store import ( - FutureStore, - FutureStatus, FutureEntry, - make_try_again_response, + FutureStatus, + FutureStore, make_failed_response, + make_try_again_response, ) @@ -47,22 +49,27 @@ async def dummy_process(data): return {"result": "success"} request_id = await future_store.create( - model_id="model-1", request_type="test", - process_fn=dummy_process, request_data={"test": "data"}, + model_id="model-1", + request_type="test", + process_fn=dummy_process, + request_data={"test": "data"}, ) entry = future_store.get(request_id) assert entry is not None and entry.request_type == "test" and entry.model_id == "model-1" # --- Processing completes with results --- processed = asyncio.Event() + async def slow_process(data): await asyncio.sleep(0.1) processed.set() return {"result": data["value"] * 2} request_id = await future_store.create( - model_id="model-1", request_type="multiply", - process_fn=slow_process, request_data={"value": 5}, + model_id="model-1", + request_type="multiply", + process_fn=slow_process, + request_data={"value": 5}, ) await asyncio.wait_for(processed.wait(), timeout=5.0) await asyncio.sleep(0.1) @@ -74,8 +81,10 @@ async def failing_process(data): raise ValueError("Invalid input data") request_id = await future_store.create( - model_id="model-1", request_type="fail", - process_fn=failing_process, request_data={}, + model_id="model-1", + request_type="fail", + process_fn=failing_process, + request_data={}, ) await asyncio.sleep(0.2) entry = future_store.get(request_id) @@ -101,8 +110,10 @@ async def tracking_process(data): request_ids = [] for i in range(5): request_id = await future_store.create( - model_id=f"model-{i}", request_type="track", - process_fn=tracking_process, request_data={"index": i}, + model_id=f"model-{i}", + request_type="track", + process_fn=tracking_process, + request_data={"index": i}, ) request_ids.append(request_id) @@ -113,6 +124,7 @@ async def tracking_process(data): # --- Stats --- completed_event = asyncio.Event() + async def slow_process2(data): await asyncio.sleep(0.1) if data.get("index") == 2: @@ -121,8 +133,10 @@ async def slow_process2(data): for i in range(3): await future_store.create( - model_id=f"model-{i}", request_type="stats", - process_fn=slow_process2, request_data={"index": i}, + model_id=f"model-{i}", + request_type="stats", + process_fn=slow_process2, + request_data={"index": i}, ) await asyncio.wait_for(completed_event.wait(), timeout=5.0) await asyncio.sleep(0.1) @@ -137,15 +151,24 @@ async def test_model_ops_deletion_status_and_expiration(self, future_store): await future_store.start() try: + async def dummy_process(data): await asyncio.sleep(0.5) return {"done": True} # --- List and delete by model --- - ids_1 = [await future_store.create(model_id="model-1", request_type="test", - process_fn=dummy_process, request_data={"index": i}) for i in range(3)] - ids_2 = [await future_store.create(model_id="model-2", request_type="test", - process_fn=dummy_process, request_data={"index": i}) for i in range(2)] + ids_1 = [ + await future_store.create( + model_id="model-1", request_type="test", process_fn=dummy_process, request_data={"index": i} + ) + for i in range(3) + ] + ids_2 = [ + await future_store.create( + model_id="model-2", request_type="test", process_fn=dummy_process, request_data={"index": i} + ) + for i in range(2) + ] assert len(future_store.list_by_model("model-1")) == 3 assert len(future_store.list_by_model("model-2")) == 2 @@ -162,8 +185,9 @@ async def dummy_process(data): async def fast_process(data): return {"done": True} - rid = await future_store.create(model_id="model-1", request_type="test", - process_fn=fast_process, request_data={}) + rid = await future_store.create( + model_id="model-1", request_type="test", process_fn=fast_process, request_data={} + ) assert future_store.get(rid) is not None assert await future_store.delete(rid) is True assert future_store.get(rid) is None @@ -174,14 +198,17 @@ async def fast_process(data): # --- Status and result tracking --- processed = asyncio.Event() + async def slow_process(data): await asyncio.sleep(0.1) processed.set() return {"value": 42} request_id = await future_store.create( - model_id="model-1", request_type="test", - process_fn=slow_process, request_data={}, + model_id="model-1", + request_type="test", + process_fn=slow_process, + request_data={}, ) status = future_store.get_status(request_id) assert status in (FutureStatus.PENDING, FutureStatus.PROCESSING) @@ -198,8 +225,10 @@ async def failing_process(data): raise RuntimeError("Server crashed") request_id = await future_store.create( - model_id="model-1", request_type="test", - process_fn=failing_process, request_data={}, + model_id="model-1", + request_type="test", + process_fn=failing_process, + request_data={}, ) await asyncio.sleep(0.2) error_info = future_store.get_error(request_id) @@ -213,11 +242,15 @@ async def failing_process(data): short_ttl_store = FutureStore(default_ttl=0.1, max_concurrent=2, cleanup_interval=60.0) await short_ttl_store.start() try: + async def dummy(data): return {"done": True} + request_id = await short_ttl_store.create( - model_id="model-1", request_type="test", - process_fn=dummy, request_data={}, + model_id="model-1", + request_type="test", + process_fn=dummy, + request_data={}, ) await asyncio.sleep(0.15) assert short_ttl_store.get(request_id).status == FutureStatus.EXPIRED @@ -228,8 +261,11 @@ async def dummy(data): await future_store.start() try: request_id = await future_store.create( - model_id="model-1", request_type="test", - process_fn=dummy, request_data={}, ttl=3600.0, + model_id="model-1", + request_type="test", + process_fn=dummy, + request_data={}, + ttl=3600.0, ) assert future_store.get(request_id).expires_at > time.time() + 3500 finally: diff --git a/tests/server/orchestrator/__init__.py b/tests/server/orchestrator/__init__.py index a52fec30..a8ba749b 100644 --- a/tests/server/orchestrator/__init__.py +++ b/tests/server/orchestrator/__init__.py @@ -1,2 +1 @@ """Tests for the engine module.""" - diff --git a/tests/server/orchestrator/test_api_orchestrator_messages.py b/tests/server/orchestrator/test_api_orchestrator_messages.py index c657a9d8..5d149203 100644 --- a/tests/server/orchestrator/test_api_orchestrator_messages.py +++ b/tests/server/orchestrator/test_api_orchestrator_messages.py @@ -10,27 +10,28 @@ import pytest + pytestmark = [pytest.mark.cpu, pytest.mark.server] from xorl.server.protocol.api_orchestrator import ( - # Enums - RequestType, - OutputType, + OrchestratorOutputs, # Core Message Types OrchestratorRequest, - OrchestratorOutputs, + OutputType, + # Enums + RequestType, + create_error_output, # Response Builders create_forward_backward_output, + create_health_check_output, + create_load_state_output, create_optim_step_output, create_save_state_output, - create_load_state_output, - create_health_check_output, - create_error_output, - # Validation and Utilities - validate_request, - validate_output, get_operation_from_request, is_streaming_output, + validate_output, + # Validation and Utilities + validate_request, ) from xorl.server.protocol.operations import ( AbortData, @@ -72,7 +73,9 @@ def test_roundtrip_request_and_response_builders(self): request_id="roundtrip-output-test", output_type=OutputType.FORWARD_BACKWARD, outputs=[{"loss": 1.234, "valid_tokens": 512, "grads_norm": 0.987}], - finished=True, error=None, timestamp=3333.0, + finished=True, + error=None, + timestamp=3333.0, ) packed = original_out.to_msgpack() assert isinstance(packed, bytes) and len(packed) > 0 @@ -88,13 +91,15 @@ def test_roundtrip_request_and_response_builders(self): # Forward backward with custom ID data = [{"model_input": {"input_ids": [1, 2, 3]}, "loss_fn_inputs": {"labels": [2, 3, 4]}}] request = OrchestratorRequest( - operation="forward_backward", payload=ModelPassData(data=data, loss_fn="causallm_loss"), + operation="forward_backward", + payload=ModelPassData(data=data, loss_fn="causallm_loss"), ) assert request.request_type == RequestType.ADD assert request.operation == "forward_backward" assert request.payload.data == data custom = OrchestratorRequest( - request_id="custom-fb-id", operation="forward_backward", + request_id="custom-fb-id", + operation="forward_backward", payload=ModelPassData(data=[{"model_input": {"input_ids": [1]}}]), ) assert custom.request_id == "custom-fb-id" @@ -133,7 +138,8 @@ def test_roundtrip_request_and_response_builders(self): health = OrchestratorRequest(request_type=RequestType.UTILITY, operation="health_check") assert health.request_type == RequestType.UTILITY abort = OrchestratorRequest( - request_type=RequestType.ABORT, operation="abort", + request_type=RequestType.ABORT, + operation="abort", payload=AbortData(target_request_id="request-to-abort"), ) assert abort.payload.target_request_id == "request-to-abort" @@ -141,7 +147,10 @@ def test_roundtrip_request_and_response_builders(self): # --- All response builders --- # Forward backward (full, minimal, error, additional metrics) output = create_forward_backward_output( - request_id="fb-req-1", loss=2.345, valid_tokens=1024, grads_norm=1.23, + request_id="fb-req-1", + loss=2.345, + valid_tokens=1024, + grads_norm=1.23, ) assert output.output_type == OutputType.FORWARD_BACKWARD assert output.outputs[0]["loss"] == 2.345 @@ -154,13 +163,18 @@ def test_roundtrip_request_and_response_builders(self): assert with_error.error == "Forward pass failed" with_metrics = create_forward_backward_output( - request_id="fb-req-3", loss=2.0, additional_metrics={"perplexity": 10.5, "accuracy": 0.85}, + request_id="fb-req-3", + loss=2.0, + additional_metrics={"perplexity": 10.5, "accuracy": 0.85}, ) assert with_metrics.outputs[0]["perplexity"] == 10.5 # Optim step (full and minimal) output = create_optim_step_output( - request_id="optim-req-1", step=1000, learning_rate=0.0001, grad_norm=0.95, + request_id="optim-req-1", + step=1000, + learning_rate=0.0001, + grad_norm=0.95, ) assert output.outputs[0]["step"] == 1000 assert output.outputs[0]["lr"] == 0.0001 @@ -169,27 +183,38 @@ def test_roundtrip_request_and_response_builders(self): # Save state (success and error) output = create_save_state_output( - request_id="save-req-1", checkpoint_path="/tmp/checkpoint", success=True, + request_id="save-req-1", + checkpoint_path="/tmp/checkpoint", + success=True, ) assert output.outputs[0]["success"] is True error_out = create_save_state_output( - request_id="save-err", checkpoint_path="/tmp/failed", success=False, error="Disk full", + request_id="save-err", + checkpoint_path="/tmp/failed", + success=False, + error="Disk full", ) assert error_out.error == "Disk full" # Load state output = create_load_state_output( - request_id="load-req-1", checkpoint_path="/tmp/checkpoint", success=True, + request_id="load-req-1", + checkpoint_path="/tmp/checkpoint", + success=True, ) assert output.outputs[0]["success"] is True # Health check output = create_health_check_output( - request_id="health-req-1", status="healthy", active_requests=5, total_requests=1000, + request_id="health-req-1", + status="healthy", + active_requests=5, + total_requests=1000, ) assert output.outputs[0]["status"] == "healthy" with_info = create_health_check_output( - request_id="health-req-2", additional_info={"uptime": 3600, "memory_usage": 0.75}, + request_id="health-req-2", + additional_info={"uptime": 3600, "memory_usage": 0.75}, ) assert with_info.outputs[0]["uptime"] == 3600 @@ -200,7 +225,8 @@ def test_roundtrip_request_and_response_builders(self): assert output.outputs == [] custom_type = create_error_output( - request_id="error-req-2", error_message="Forward failed", + request_id="error-req-2", + error_message="Forward failed", operation_type=OutputType.FORWARD_BACKWARD, ) assert custom_type.output_type == OutputType.FORWARD_BACKWARD @@ -212,26 +238,51 @@ class TestValidationUtilitiesAndIntegration: def test_validation_and_utilities(self): """Test request/output validation, get_operation, and is_streaming_output.""" # Valid ADD - assert validate_request(OrchestratorRequest( - request_id="valid-req", request_type=RequestType.ADD, operation="forward_backward", - )) is True + assert ( + validate_request( + OrchestratorRequest( + request_id="valid-req", + request_type=RequestType.ADD, + operation="forward_backward", + ) + ) + is True + ) # Valid ABORT - assert validate_request(OrchestratorRequest( - request_id="abort-req", request_type=RequestType.ABORT, operation="abort", - payload=AbortData(target_request_id="req-to-abort"), - )) is True + assert ( + validate_request( + OrchestratorRequest( + request_id="abort-req", + request_type=RequestType.ABORT, + operation="abort", + payload=AbortData(target_request_id="req-to-abort"), + ) + ) + is True + ) # Valid UTILITY - assert validate_request(OrchestratorRequest( - request_id="utility-req", request_type=RequestType.UTILITY, operation="health_check", - )) is True + assert ( + validate_request( + OrchestratorRequest( + request_id="utility-req", + request_type=RequestType.UTILITY, + operation="health_check", + ) + ) + is True + ) # Missing request_id with pytest.raises(ValueError, match="must have request_id"): - validate_request(OrchestratorRequest( - request_id="", request_type=RequestType.ADD, operation="test", - )) + validate_request( + OrchestratorRequest( + request_id="", + request_type=RequestType.ADD, + operation="test", + ) + ) # ADD missing operation with pytest.raises(ValueError, match="must have 'operation'"): @@ -239,50 +290,95 @@ def test_validation_and_utilities(self): # ABORT missing target with pytest.raises(ValueError, match="must have 'target_request_id'"): - validate_request(OrchestratorRequest( - request_id="abort-no-target", request_type=RequestType.ABORT, operation="abort", - )) + validate_request( + OrchestratorRequest( + request_id="abort-no-target", + request_type=RequestType.ABORT, + operation="abort", + ) + ) # Valid streaming output - assert validate_output(OrchestratorOutputs( - request_id="streaming-output", output_type=OutputType.FORWARD_BACKWARD, - outputs=[{"partial": "data"}], finished=False, - )) is True + assert ( + validate_output( + OrchestratorOutputs( + request_id="streaming-output", + output_type=OutputType.FORWARD_BACKWARD, + outputs=[{"partial": "data"}], + finished=False, + ) + ) + is True + ) # Valid error output - assert validate_output(OrchestratorOutputs( - request_id="error-output", output_type=OutputType.ERROR, - outputs=[], finished=True, error="Test error", - )) is True + assert ( + validate_output( + OrchestratorOutputs( + request_id="error-output", + output_type=OutputType.ERROR, + outputs=[], + finished=True, + error="Test error", + ) + ) + is True + ) # Missing request_id with pytest.raises(ValueError, match="must have request_id"): - validate_output(OrchestratorOutputs( - request_id="", output_type=OutputType.FORWARD_BACKWARD, - )) + validate_output( + OrchestratorOutputs( + request_id="", + output_type=OutputType.FORWARD_BACKWARD, + ) + ) # get_operation_from_request assert get_operation_from_request(OrchestratorRequest(operation="forward_backward")) == "forward_backward" assert get_operation_from_request(OrchestratorRequest()) is None # is_streaming_output - assert is_streaming_output(OrchestratorOutputs( - request_id="stream-test", output_type=OutputType.FORWARD_BACKWARD, - finished=False, error=None, - )) is True - assert is_streaming_output(OrchestratorOutputs( - request_id="finished-test", output_type=OutputType.FORWARD_BACKWARD, finished=True, - )) is False - assert is_streaming_output(OrchestratorOutputs( - request_id="error-test", output_type=OutputType.ERROR, finished=False, error="Test error", - )) is False + assert ( + is_streaming_output( + OrchestratorOutputs( + request_id="stream-test", + output_type=OutputType.FORWARD_BACKWARD, + finished=False, + error=None, + ) + ) + is True + ) + assert ( + is_streaming_output( + OrchestratorOutputs( + request_id="finished-test", + output_type=OutputType.FORWARD_BACKWARD, + finished=True, + ) + ) + is False + ) + assert ( + is_streaming_output( + OrchestratorOutputs( + request_id="error-test", + output_type=OutputType.ERROR, + finished=False, + error="Test error", + ) + ) + is False + ) def test_complete_request_response_and_error_flow(self): """Test complete request-response flow and error handling for all operation types.""" # Forward backward flow data = [{"model_input": {"input_ids": [1, 2, 3, 4, 5]}, "loss_fn_inputs": {"labels": [2, 3, 4, 5, 6]}}] request = OrchestratorRequest( - operation="forward_backward", payload=ModelPassData(data=data, loss_fn="causallm_loss"), + operation="forward_backward", + payload=ModelPassData(data=data, loss_fn="causallm_loss"), ) assert validate_request(request) is True received_request = OrchestratorRequest.from_msgpack(request.to_msgpack()) @@ -296,11 +392,15 @@ def test_complete_request_response_and_error_flow(self): # Optim step flow optim_req = OrchestratorRequest( - operation="optim_step", payload=OptimStepData(lr=0.001, gradient_clip=1.0), + operation="optim_step", + payload=OptimStepData(lr=0.001, gradient_clip=1.0), ) received = OrchestratorRequest.from_msgpack(optim_req.to_msgpack()) optim_out = create_optim_step_output( - request_id=received.request_id, step=100, learning_rate=0.001, grad_norm=0.85, + request_id=received.request_id, + step=100, + learning_rate=0.001, + grad_norm=0.85, ) final = OrchestratorOutputs.from_msgpack(optim_out.to_msgpack()) assert final.outputs[0]["step"] == 100 and final.outputs[0]["grad_norm"] == 0.85 @@ -312,7 +412,9 @@ def test_complete_request_response_and_error_flow(self): ) received = OrchestratorRequest.from_msgpack(save_req.to_msgpack()) save_out = create_save_state_output( - request_id=received.request_id, checkpoint_path=received.payload.checkpoint_path, success=True, + request_id=received.request_id, + checkpoint_path=received.payload.checkpoint_path, + success=True, ) final = OrchestratorOutputs.from_msgpack(save_out.to_msgpack()) assert final.outputs[0]["success"] is True @@ -323,7 +425,8 @@ def test_complete_request_response_and_error_flow(self): payload=ModelPassData(data=[{"model_input": {"input_ids": [1, 2, 3]}}]), ) error_output = create_error_output( - request_id=request.request_id, error_message="CUDA out of memory", + request_id=request.request_id, + error_message="CUDA out of memory", operation_type=OutputType.FORWARD_BACKWARD, ) assert validate_output(error_output) is True diff --git a/tests/server/orchestrator/test_cu_seqlens_alignment.py b/tests/server/orchestrator/test_cu_seqlens_alignment.py index 5d111540..c0928555 100644 --- a/tests/server/orchestrator/test_cu_seqlens_alignment.py +++ b/tests/server/orchestrator/test_cu_seqlens_alignment.py @@ -15,17 +15,18 @@ """ import math +from unittest.mock import Mock, patch + import pytest import torch -from unittest.mock import Mock, patch from xorl.data.collators.packing_concat_collator import ( PackingConcatCollator, - add_flash_attention_kwargs_from_position_ids, ) -from xorl.server.orchestrator.packing import SequentialPacker, pack_samples +from xorl.server.orchestrator.packing import SequentialPacker from xorl.utils.seqlen_pos_transform_utils import prepare_fa_kwargs_from_position_ids + pytestmark = [pytest.mark.cpu] @@ -33,6 +34,7 @@ # Helpers # --------------------------------------------------------------------------- + def _server_path_cu_seqlens(datum_list, max_seq_len=4096, pad_to_multiple_of=1): """Simulate the server path: pack -> finalize -> convert_to_tensors.""" packer = SequentialPacker(enable_packing=True, log_stats=False, pad_to_multiple_of=pad_to_multiple_of) @@ -89,6 +91,7 @@ def _make_feature(input_ids, position_ids): # Tests # --------------------------------------------------------------------------- + class TestCuSeqlensAlignmentAndDtype: """Verify server/CLI cu_seq_lens match, dtype, and edge cases.""" @@ -103,10 +106,12 @@ def test_server_cli_match_dtype_and_edge_cases(self, mock_cli_ps, mock_server_ps # --- Two sequences match --- seq1, seq2 = [10, 20, 30], [40, 50, 60, 70] server_batch = _server_path_cu_seqlens([_make_datum(seq1), _make_datum(seq2)])[0] - cli_batch = _cli_path_cu_seqlens([ - _make_feature(seq1, list(range(len(seq1)))), - _make_feature(seq2, list(range(len(seq2)))), - ]) + cli_batch = _cli_path_cu_seqlens( + [ + _make_feature(seq1, list(range(len(seq1)))), + _make_feature(seq2, list(range(len(seq2)))), + ] + ) assert torch.equal(server_batch["cu_seq_lens_q"], cli_batch["cu_seq_lens_q"]) assert torch.equal(server_batch["cu_seq_lens_k"], cli_batch["cu_seq_lens_k"]) assert server_batch["max_length_q"] == cli_batch["max_length_q"] @@ -158,8 +163,10 @@ def test_sequence_shard_collator_preservation_and_stale_overwrite(self): from xorl.data.collators.sequence_shard_collator import TextSequenceShardCollator # Preservation (SP size 2, 6 tokens) - with patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") as mock_ps, \ - patch("xorl.data.collators.packing_concat_collator.get_parallel_state") as mock_ps2: + with ( + patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") as mock_ps, + patch("xorl.data.collators.packing_concat_collator.get_parallel_state") as mock_ps2, + ): mock = Mock(cp_size=2, cp_rank=0, cp_enabled=True, ringattn_size=1) mock_ps.return_value = mock mock_ps2.return_value = mock @@ -177,8 +184,10 @@ def test_sequence_shard_collator_preservation_and_stale_overwrite(self): assert result["position_ids"].shape[-1] >= 6 # Not padded (SP size 4, 5 tokens -> pad to 8, but _original keeps 5) - with patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") as mock_ps, \ - patch("xorl.data.collators.packing_concat_collator.get_parallel_state") as mock_ps2: + with ( + patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") as mock_ps, + patch("xorl.data.collators.packing_concat_collator.get_parallel_state") as mock_ps2, + ): mock = Mock(cp_size=4, cp_rank=0, cp_enabled=True, ringattn_size=1) mock_ps.return_value = mock mock_ps2.return_value = mock @@ -193,8 +202,10 @@ def test_sequence_shard_collator_preservation_and_stale_overwrite(self): assert result["position_ids"].shape[-1] == 8 # Stale cu_seq_lens overwritten - with patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") as mock_ps, \ - patch("xorl.data.collators.packing_concat_collator.get_parallel_state") as mock_ps2: + with ( + patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") as mock_ps, + patch("xorl.data.collators.packing_concat_collator.get_parallel_state") as mock_ps2, + ): mock = Mock(cp_size=4, cp_rank=0, cp_enabled=True, ringattn_size=1) mock_ps.return_value = mock mock_ps2.return_value = mock @@ -205,7 +216,8 @@ def test_sequence_shard_collator_preservation_and_stale_overwrite(self): "position_ids": torch.tensor([[0, 1, 2, 0, 1]], dtype=torch.long), "cu_seq_lens_q": torch.tensor([0, 3, 5], dtype=torch.int32), "cu_seq_lens_k": torch.tensor([0, 3, 5], dtype=torch.int32), - "max_length_q": 3, "max_length_k": 3, + "max_length_q": 3, + "max_length_k": 3, } result = collator(batch) assert result["cu_seq_lens_q"][-1].item() == 8 @@ -213,17 +225,22 @@ def test_sequence_shard_collator_preservation_and_stale_overwrite(self): # cu_seq_lens match with padding pad = 128 - with patch("xorl.server.orchestrator.packing.get_parallel_state") as mock_server_ps, \ - patch("xorl.data.collators.packing_concat_collator.get_parallel_state") as mock_cli_ps: + with ( + patch("xorl.server.orchestrator.packing.get_parallel_state") as mock_server_ps, + patch("xorl.data.collators.packing_concat_collator.get_parallel_state") as mock_cli_ps, + ): mock_ps = Mock(cp_enabled=False, cp_size=1, cp_rank=0, ringattn_size=1) mock_server_ps.return_value = mock_ps mock_cli_ps.return_value = mock_ps seq1, seq2 = [1, 2, 3], [4, 5, 6, 7, 8, 9, 10] server_batch = _server_path_cu_seqlens([_make_datum(seq1), _make_datum(seq2)], pad_to_multiple_of=pad)[0] - cli_batch = _cli_path_cu_seqlens([ - _make_feature(seq1, list(range(len(seq1)))), - _make_feature(seq2, list(range(len(seq2)))), - ], pad_to_multiple_of=pad) + cli_batch = _cli_path_cu_seqlens( + [ + _make_feature(seq1, list(range(len(seq1)))), + _make_feature(seq2, list(range(len(seq2)))), + ], + pad_to_multiple_of=pad, + ) assert server_batch["input_ids"].shape[-1] == pad assert torch.equal(server_batch["cu_seq_lens_q"], cli_batch["cu_seq_lens_q"]) assert server_batch["cu_seq_lens_q"][-1].item() == pad @@ -252,8 +269,9 @@ def test_lcm_padding_and_collator_divisibility(self, mock_ps): assert len(batches[0]["input_ids"][0]) == expected # RequestProcessor computes lcm correctly - from xorl.server.orchestrator.request_processor import RequestProcessor from xorl.server.backend import DummyBackend + from xorl.server.orchestrator.request_processor import RequestProcessor + backend = DummyBackend() assert RequestProcessor(backend=backend, pad_to_multiple_of=128, cp_size=3).pad_to_multiple_of == 384 assert RequestProcessor(backend=backend, pad_to_multiple_of=128, cp_size=8).pad_to_multiple_of == 128 @@ -261,25 +279,32 @@ def test_lcm_padding_and_collator_divisibility(self, mock_ps): # After TextSequenceShardCollator, lengths divisible by cp_size from xorl.data.collators.sequence_shard_collator import TextSequenceShardCollator + cp_size = 3 mock_ps.return_value = Mock(cp_enabled=True, cp_size=cp_size, cp_rank=0, ringattn_size=1) packer = SequentialPacker(enable_packing=True, log_stats=False, pad_to_multiple_of=math.lcm(128, cp_size)) - batches = packer.pack([_make_datum(list(range(50))), _make_datum(list(range(30)))], max_seq_len=4096, request_id="test") + batches = packer.pack( + [_make_datum(list(range(50))), _make_datum(list(range(30)))], max_seq_len=4096, request_id="test" + ) batch = batches[0] pre_len = len(batch["input_ids"][0]) assert pre_len % cp_size == 0 - with patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") as mock_sp, \ - patch("xorl.data.collators.packing_concat_collator.get_parallel_state") as mock_sp2: + with ( + patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") as mock_sp, + patch("xorl.data.collators.packing_concat_collator.get_parallel_state") as mock_sp2, + ): m = Mock(cp_size=cp_size, cp_rank=0, cp_enabled=True, ringattn_size=1) mock_sp.return_value = m mock_sp2.return_value = m collator = TextSequenceShardCollator() - result = collator({ - "input_ids": torch.tensor(batch["input_ids"], dtype=torch.long), - "labels": torch.tensor(batch["labels"], dtype=torch.long), - "position_ids": torch.tensor(batch["position_ids"], dtype=torch.long), - }) + result = collator( + { + "input_ids": torch.tensor(batch["input_ids"], dtype=torch.long), + "labels": torch.tensor(batch["labels"], dtype=torch.long), + "position_ids": torch.tensor(batch["position_ids"], dtype=torch.long), + } + ) assert result["position_ids"].shape[-1] == pre_len assert result["input_ids"].shape[-1] == pre_len // cp_size assert result["cu_seq_lens_q"][-1].item() == pre_len @@ -292,6 +317,7 @@ class TestUnpackingWithPadding: def test_padding_boundary_handling(self, mock_ps): """Padding creates extra boundary; no padding has correct count.""" from xorl.server.orchestrator.packing import unpack_per_token_outputs + mock_ps.return_value = Mock(cp_enabled=False, cp_size=1, cp_rank=0, ringattn_size=1) datum_list = [_make_datum([10, 20, 30]), _make_datum([40, 50, 60, 70])] @@ -302,7 +328,7 @@ def test_padding_boundary_handling(self, mock_ps): assert batch["num_samples"] == 2 and len(batch["position_ids"][0]) == 128 pos_ids = torch.tensor(batch["position_ids"][0], dtype=torch.long) unpacked = unpack_per_token_outputs(torch.randn(127), pos_ids) - assert len(unpacked) == 3 and len(unpacked[:batch["num_samples"]]) == 2 + assert len(unpacked) == 3 and len(unpacked[: batch["num_samples"]]) == 2 # Without padding: correct count packer = SequentialPacker(enable_packing=True, log_stats=False, pad_to_multiple_of=1) diff --git a/tests/server/orchestrator/test_orchestrator.py b/tests/server/orchestrator/test_orchestrator.py index 6690791f..122d7f2f 100644 --- a/tests/server/orchestrator/test_orchestrator.py +++ b/tests/server/orchestrator/test_orchestrator.py @@ -15,39 +15,38 @@ import pytest + pytestmark = [pytest.mark.cpu, pytest.mark.server] -import time -import queue import socket -import threading -from typing import Any, Dict, List, Optional +import time + +import zmq -from xorl.server.orchestrator.orchestrator import Orchestrator from xorl.server.backend import DummyBackend +from xorl.server.orchestrator.orchestrator import Orchestrator from xorl.server.protocol.api_orchestrator import ( - OrchestratorRequest, OrchestratorOutputs, - RequestType, + OrchestratorRequest, OutputType, + RequestType, ) from xorl.server.protocol.operations import ( - ModelPassData, - OptimStepData, AbortData, EmptyData, + ModelPassData, + OptimStepData, ) -import zmq - # ============================================================================ # Fixtures # ============================================================================ + def find_free_port(): """Find a free port.""" with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(('127.0.0.1', 0)) + s.bind(("127.0.0.1", 0)) s.listen(1) port = s.getsockname()[1] return port @@ -105,9 +104,11 @@ def output_socket(addresses): # Helper Functions # ============================================================================ + def receive_outputs(output_socket, timeout_ms=1000, max_outputs=10): """Receive outputs from output socket.""" import msgpack + outputs = [] while len(outputs) < max_outputs: @@ -125,6 +126,7 @@ def receive_outputs(output_socket, timeout_ms=1000, max_outputs=10): # Tests # ============================================================================ + def test_init_start_stop_and_all_operations(addresses, orchestrator, output_socket): """Test Orchestrator lifecycle and event loop processing of all operation types.""" # --- Init and start/stop lifecycle --- @@ -154,12 +156,15 @@ def test_init_start_stop_and_all_operations(addresses, orchestrator, output_sock # --- Forward backward --- fb_request = OrchestratorRequest( - request_id="req-fb-001", request_type=RequestType.ADD, + request_id="req-fb-001", + request_type=RequestType.ADD, operation="forward_backward", - payload=ModelPassData(data=[ - {"input_ids": [1, 2, 3, 4], "labels": [2, 3, 4, 5]}, - {"input_ids": [10, 20], "labels": [20, 30]}, - ]), + payload=ModelPassData( + data=[ + {"input_ids": [1, 2, 3, 4], "labels": [2, 3, 4, 5]}, + {"input_ids": [10, 20], "labels": [20, 30]}, + ] + ), ) orchestrator.input_queue.put(fb_request) time.sleep(0.5) @@ -171,8 +176,10 @@ def test_init_start_stop_and_all_operations(addresses, orchestrator, output_sock # --- Optim step --- opt_request = OrchestratorRequest( - request_id="req-opt-001", request_type=RequestType.ADD, - operation="optim_step", payload=OptimStepData(lr=0.001, gradient_clip=1.0), + request_id="req-opt-001", + request_type=RequestType.ADD, + operation="optim_step", + payload=OptimStepData(lr=0.001, gradient_clip=1.0), ) orchestrator.input_queue.put(opt_request) time.sleep(0.5) @@ -183,8 +190,10 @@ def test_init_start_stop_and_all_operations(addresses, orchestrator, output_sock # --- Health check --- health_request = OrchestratorRequest( - request_id="req-health-001", request_type=RequestType.UTILITY, - operation="health_check", payload=EmptyData(), + request_id="req-health-001", + request_type=RequestType.UTILITY, + operation="health_check", + payload=EmptyData(), ) orchestrator.input_queue.put(health_request) time.sleep(0.3) @@ -198,8 +207,10 @@ def test_errors_abort_e2e_and_concurrent(orchestrator, output_socket): """Test error handling, abort, end-to-end flow, and concurrent requests.""" # --- Invalid operation --- invalid_request = OrchestratorRequest( - request_id="req-error-001", request_type=RequestType.ADD, - operation="invalid_operation", payload=EmptyData(), + request_id="req-error-001", + request_type=RequestType.ADD, + operation="invalid_operation", + payload=EmptyData(), ) orchestrator.input_queue.put(invalid_request) time.sleep(0.5) @@ -210,8 +221,10 @@ def test_errors_abort_e2e_and_concurrent(orchestrator, output_socket): # --- Empty datum list --- empty_request = OrchestratorRequest( - request_id="req-empty-001", request_type=RequestType.ADD, - operation="forward_backward", payload=ModelPassData(data=[]), + request_id="req-empty-001", + request_type=RequestType.ADD, + operation="forward_backward", + payload=ModelPassData(data=[]), ) orchestrator.input_queue.put(empty_request) time.sleep(0.5) @@ -222,15 +235,18 @@ def test_errors_abort_e2e_and_concurrent(orchestrator, output_socket): # --- Abort --- fb_request = OrchestratorRequest( - request_id="req-abort-001", request_type=RequestType.ADD, + request_id="req-abort-001", + request_type=RequestType.ADD, operation="forward_backward", payload=ModelPassData(data=[{"input_ids": [1, 2], "labels": [2, 3]}]), ) orchestrator.input_queue.put(fb_request) time.sleep(0.1) abort_request = OrchestratorRequest( - request_id="abort-req", request_type=RequestType.ABORT, - operation="abort", payload=AbortData(target_request_id="req-abort-001"), + request_id="abort-req", + request_type=RequestType.ABORT, + operation="abort", + payload=AbortData(target_request_id="req-abort-001"), ) orchestrator.input_queue.put(abort_request) time.sleep(0.3) @@ -240,7 +256,8 @@ def test_errors_abort_e2e_and_concurrent(orchestrator, output_socket): # --- End-to-end --- request = OrchestratorRequest( - request_id="req-e2e-001", request_type=RequestType.ADD, + request_id="req-e2e-001", + request_type=RequestType.ADD, operation="forward_backward", payload=ModelPassData(data=[{"input_ids": [1, 2, 3, 4, 5], "labels": [2, 3, 4, 5, 6]}]), ) @@ -257,7 +274,8 @@ def test_errors_abort_e2e_and_concurrent(orchestrator, output_socket): requests = [] for i in range(5): req = OrchestratorRequest( - request_id=f"req-concurrent-{i:03d}", request_type=RequestType.ADD, + request_id=f"req-concurrent-{i:03d}", + request_type=RequestType.ADD, operation="forward_backward", payload=ModelPassData(data=[{"input_ids": list(range(10)), "labels": list(range(1, 11))}]), ) @@ -284,7 +302,8 @@ def test_stats_throughput_and_health_check(orchestrator): initial_stats = orchestrator.get_stats() for i in range(3): req = OrchestratorRequest( - request_id=f"req-throughput-{i:03d}", request_type=RequestType.ADD, + request_id=f"req-throughput-{i:03d}", + request_type=RequestType.ADD, operation="forward_backward", payload=ModelPassData(data=[{"input_ids": [1, 2], "labels": [2, 3]}]), ) @@ -301,20 +320,48 @@ def test_stats_throughput_and_health_check(orchestrator): ) # Health check identified correctly - assert bare._is_health_check_request(OrchestratorRequest( - request_id="test-health", request_type=RequestType.UTILITY, operation="health_check", - )) is True + assert ( + bare._is_health_check_request( + OrchestratorRequest( + request_id="test-health", + request_type=RequestType.UTILITY, + operation="health_check", + ) + ) + is True + ) # Non-health-check requests - assert bare._is_health_check_request(OrchestratorRequest( - request_id="test-fb", request_type=RequestType.ADD, operation="forward_backward", - )) is False - assert bare._is_health_check_request(OrchestratorRequest( - request_id="test-abort", request_type=RequestType.ABORT, operation="abort", - )) is False - assert bare._is_health_check_request(OrchestratorRequest( - request_id="test-adapter", request_type=RequestType.UTILITY, operation="get_adapter_info", - )) is False + assert ( + bare._is_health_check_request( + OrchestratorRequest( + request_id="test-fb", + request_type=RequestType.ADD, + operation="forward_backward", + ) + ) + is False + ) + assert ( + bare._is_health_check_request( + OrchestratorRequest( + request_id="test-abort", + request_type=RequestType.ABORT, + operation="abort", + ) + ) + is False + ) + assert ( + bare._is_health_check_request( + OrchestratorRequest( + request_id="test-adapter", + request_type=RequestType.UTILITY, + operation="get_adapter_info", + ) + ) + is False + ) # Methods exist assert callable(bare._is_health_check_request) diff --git a/tests/server/orchestrator/test_orchestrator_client_communication.py b/tests/server/orchestrator/test_orchestrator_client_communication.py index 968056da..62691ae4 100644 --- a/tests/server/orchestrator/test_orchestrator_client_communication.py +++ b/tests/server/orchestrator/test_orchestrator_client_communication.py @@ -9,21 +9,20 @@ import asyncio import logging + import pytest import pytest_asyncio import zmq import zmq.asyncio -from typing import Optional from xorl.server.api_server.orchestrator_client import OrchestratorClient from xorl.server.protocol.api_orchestrator import ( - OrchestratorRequest, OrchestratorOutputs, - RequestType, + OrchestratorRequest, OutputType, + RequestType, ) from xorl.server.protocol.operations import ( - EmptyData, ModelPassData, OptimStepData, ) @@ -36,6 +35,7 @@ # Mock Engine # ============================================================================ + class MockEngine: """ Mock Engine that simulates the INPUT/OUTPUT socket behavior of Orchestrator. @@ -116,33 +116,53 @@ async def _process_requests(self): async def _handle_request(self, request): op = request.operation if op == "forward_backward": - num_samples = len(request.payload.data) if hasattr(request.payload, 'data') else 0 - return OrchestratorOutputs(request_id=request.request_id, output_type=OutputType.FORWARD_BACKWARD, - outputs=[{"loss": 2.5, "num_samples": num_samples, "status": "success"}], finished=True) + num_samples = len(request.payload.data) if hasattr(request.payload, "data") else 0 + return OrchestratorOutputs( + request_id=request.request_id, + output_type=OutputType.FORWARD_BACKWARD, + outputs=[{"loss": 2.5, "num_samples": num_samples, "status": "success"}], + finished=True, + ) elif op == "optim_step": - lr = request.payload.lr if hasattr(request.payload, 'lr') else 0.0001 - return OrchestratorOutputs(request_id=request.request_id, output_type=OutputType.OPTIM_STEP, - outputs=[{"grad_norm": 1.23, "learning_rate": lr, "status": "success"}], finished=True) + lr = request.payload.lr if hasattr(request.payload, "lr") else 0.0001 + return OrchestratorOutputs( + request_id=request.request_id, + output_type=OutputType.OPTIM_STEP, + outputs=[{"grad_norm": 1.23, "learning_rate": lr, "status": "success"}], + finished=True, + ) elif op == "health_check": - return OrchestratorOutputs(request_id=request.request_id, output_type=OutputType.HEALTH_CHECK, - outputs=[{"status": "healthy", "workers_discovered": True, "active_workers": 8}], finished=True) + return OrchestratorOutputs( + request_id=request.request_id, + output_type=OutputType.HEALTH_CHECK, + outputs=[{"status": "healthy", "workers_discovered": True, "active_workers": 8}], + finished=True, + ) else: - return OrchestratorOutputs(request_id=request.request_id, output_type=OutputType.ERROR, - outputs=[], finished=True, error=f"Unknown operation: {op}") + return OrchestratorOutputs( + request_id=request.request_id, + output_type=OutputType.ERROR, + outputs=[], + finished=True, + error=f"Unknown operation: {op}", + ) # ============================================================================ # Fixtures # ============================================================================ + @pytest.fixture def zmq_addresses(): import socket as sock + def find_free_port(): with sock.socket(sock.AF_INET, sock.SOCK_STREAM) as s: - s.bind(('127.0.0.1', 0)) + s.bind(("127.0.0.1", 0)) s.listen(1) return s.getsockname()[1] + return {"input": f"tcp://127.0.0.1:{find_free_port()}", "output": f"tcp://127.0.0.1:{find_free_port()}"} @@ -166,12 +186,15 @@ async def orchestrator_client(zmq_addresses): # Helper to create requests # ============================================================================ + def _health_check_request(): return OrchestratorRequest(request_type=RequestType.UTILITY, operation="health_check") + def _forward_backward_request(data): return OrchestratorRequest(operation="forward_backward", payload=ModelPassData(data=data)) + def _optim_step_request(lr, gradient_clip=None): return OrchestratorRequest(operation="optim_step", payload=OptimStepData(lr=lr, gradient_clip=gradient_clip)) @@ -180,6 +203,7 @@ def _optim_step_request(lr, gradient_clip=None): # Tests # ============================================================================ + @pytest.mark.asyncio async def test_basic_communication_and_roundtrip(mock_engine, orchestrator_client): """Test connection, send, receive, and full roundtrip.""" @@ -205,10 +229,12 @@ async def test_forward_backward_and_optim_step(mock_engine, orchestrator_client) await asyncio.sleep(0.2) # Forward backward - fb_request = _forward_backward_request(data=[ - {"model_input": {"input_ids": [1, 2, 3]}, "loss_fn_inputs": {"labels": [2, 3, 4]}}, - {"model_input": {"input_ids": [5, 6, 7]}, "loss_fn_inputs": {"labels": [6, 7, 8]}}, - ]) + fb_request = _forward_backward_request( + data=[ + {"model_input": {"input_ids": [1, 2, 3]}, "loss_fn_inputs": {"labels": [2, 3, 4]}}, + {"model_input": {"input_ids": [5, 6, 7]}, "loss_fn_inputs": {"labels": [6, 7, 8]}}, + ] + ) future = await orchestrator_client.send_request(fb_request) output = await asyncio.wait_for(future, timeout=2.0) assert output.output_type == OutputType.FORWARD_BACKWARD @@ -269,8 +295,12 @@ async def test_edge_cases_and_lifecycle(zmq_addresses, mock_engine, orchestrator assert stats["running"] is True # Serialization roundtrip - data = [{"model_input": {"input_ids": list(range(100)), "attention_mask": [1]*100}, - "loss_fn_inputs": {"labels": list(range(100, 200))}}] + data = [ + { + "model_input": {"input_ids": list(range(100)), "attention_mask": [1] * 100}, + "loss_fn_inputs": {"labels": list(range(100, 200))}, + } + ] request = _forward_backward_request(data=data) await orchestrator_client.send_request(request) await asyncio.sleep(0.2) @@ -286,9 +316,10 @@ async def test_edge_cases_and_lifecycle(zmq_addresses, mock_engine, orchestrator # Client start/stop - use fresh ports to avoid address conflicts import socket as sock + def _find_free_port(): with sock.socket(sock.AF_INET, sock.SOCK_STREAM) as s: - s.bind(('127.0.0.1', 0)) + s.bind(("127.0.0.1", 0)) s.listen(1) return s.getsockname()[1] diff --git a/tests/server/orchestrator/test_orchestrator_runner_messages.py b/tests/server/orchestrator/test_orchestrator_runner_messages.py index dc9b8945..32004e5d 100644 --- a/tests/server/orchestrator/test_orchestrator_runner_messages.py +++ b/tests/server/orchestrator/test_orchestrator_runner_messages.py @@ -13,19 +13,9 @@ import pytest + pytestmark = [pytest.mark.cpu, pytest.mark.server] -from xorl.server.protocol.orchestrator_runner import ( - MessageType, - BaseMessage, - RunnerReady, - RunnerAck, - RunnerResponse, - RunnerDispatchCommand, - serialize_message, - deserialize_message, - create_ack_for_request, -) from xorl.server.protocol.operations import ( EmptyData, LoadStateData, @@ -33,6 +23,17 @@ OptimStepData, SaveStateData, ) +from xorl.server.protocol.orchestrator_runner import ( + BaseMessage, + MessageType, + RunnerAck, + RunnerDispatchCommand, + RunnerReady, + RunnerResponse, + create_ack_for_request, + deserialize_message, + serialize_message, +) def test_dispatch_command_all_operations(): @@ -51,7 +52,9 @@ def test_dispatch_command_all_operations(): # optim_step (with and without clip) msg = RunnerDispatchCommand.create( - operation="optim_step", payload=OptimStepData(lr=0.001, gradient_clip=1.0), request_id="opt-id", + operation="optim_step", + payload=OptimStepData(lr=0.001, gradient_clip=1.0), + request_id="opt-id", ) assert msg.payload.lr == 0.001 and msg.payload.gradient_clip == 1.0 msg_no_clip = RunnerDispatchCommand.create(operation="optim_step", payload=OptimStepData(lr=0.001)) @@ -59,13 +62,15 @@ def test_dispatch_command_all_operations(): # save_state msg = RunnerDispatchCommand.create( - operation="save_state", payload=SaveStateData(checkpoint_path="/tmp/ckpt.pt", save_optimizer=True), + operation="save_state", + payload=SaveStateData(checkpoint_path="/tmp/ckpt.pt", save_optimizer=True), ) assert msg.payload.checkpoint_path == "/tmp/ckpt.pt" and msg.payload.save_optimizer is True # load_state msg = RunnerDispatchCommand.create( - operation="load_state", payload=LoadStateData(checkpoint_path="/tmp/ckpt.pt", load_optimizer=False), + operation="load_state", + payload=LoadStateData(checkpoint_path="/tmp/ckpt.pt", load_optimizer=False), ) assert msg.payload.load_optimizer is False @@ -107,10 +112,16 @@ def test_json_conversion(): assert data["worker_rank"] == 0 and data["world_size"] == 4 # from_json - json_str = json.dumps({ - "message_type": "ready", "worker_rank": 2, "world_size": 8, - "device": "cuda:2", "message_id": "msg-123", "timestamp": time.time(), - }) + json_str = json.dumps( + { + "message_type": "ready", + "worker_rank": 2, + "world_size": 8, + "device": "cuda:2", + "message_id": "msg-123", + "timestamp": time.time(), + } + ) msg = BaseMessage.from_json(json_str) assert isinstance(msg, RunnerReady) assert msg.worker_rank == 2 and msg.world_size == 8 and msg.device == "cuda:2" @@ -119,12 +130,19 @@ def test_json_conversion(): def test_complex_data_and_edge_cases(): """Test large batches, complex results, None values, ID uniqueness, and timestamp accuracy.""" # Large batches - batches = [{"input_ids": [list(range(100)) for _ in range(10)], - "labels": [list(range(100, 200)) for _ in range(10)], - "position_ids": [list(range(100)) for _ in range(10)], - "request_id": "req-123", "batch_id": i} for i in range(5)] + batches = [ + { + "input_ids": [list(range(100)) for _ in range(10)], + "labels": [list(range(100, 200)) for _ in range(10)], + "position_ids": [list(range(100)) for _ in range(10)], + "request_id": "req-123", + "batch_id": i, + } + for i in range(5) + ] msg = RunnerDispatchCommand.create( - operation="forward_backward", payload=ModelPassData(batches=batches, loss_fn="causallm_loss"), + operation="forward_backward", + payload=ModelPassData(batches=batches, loss_fn="causallm_loss"), ) deserialized = deserialize_message(serialize_message(msg)) assert len(deserialized.payload.batches) == 5 @@ -153,7 +171,9 @@ def test_complex_data_and_edge_cases(): # ACK creation request = RunnerDispatchCommand.create( - operation="forward_backward", payload=ModelPassData(batches=[], loss_fn="test"), request_id="req-123", + operation="forward_backward", + payload=ModelPassData(batches=[], loss_fn="test"), + request_id="req-123", ) ack = create_ack_for_request(request) assert isinstance(ack, RunnerAck) diff --git a/tests/server/orchestrator/test_packing.py b/tests/server/orchestrator/test_packing.py index c72fe969..7fe6c9c2 100644 --- a/tests/server/orchestrator/test_packing.py +++ b/tests/server/orchestrator/test_packing.py @@ -4,14 +4,15 @@ import pytest import torch + pytestmark = [pytest.mark.cpu, pytest.mark.server] from xorl.server.orchestrator.packing import ( + Packer, SequentialPacker, pack_samples, - validate_micro_batches, unpack_per_token_outputs, - Packer, + validate_micro_batches, ) @@ -19,6 +20,7 @@ # Fixtures # ============================================================================ + @pytest.fixture def simple_data(): return [ @@ -43,6 +45,7 @@ def mixed_length_data(): # Core packing # ============================================================================ + def test_packing_enabled(simple_data): """Packing ON: samples concatenated into single sequence with correct shifting.""" packer = SequentialPacker(enable_packing=True, log_stats=False, pad_to_multiple_of=1) @@ -156,58 +159,96 @@ def test_mixed_length_and_capacity(mixed_length_data): assert len(batch["input_ids"][0]) <= 60 # Exact fit: 5+5=10 - batches = packer.pack( - [{"input_ids": [1] * 5}, {"input_ids": [2] * 5}], max_seq_len=10 - ) + batches = packer.pack([{"input_ids": [1] * 5}, {"input_ids": [2] * 5}], max_seq_len=10) assert len(batches) == 1 and batches[0]["num_samples"] == 2 # Off-by-one: 5+6=11 > 10 - batches = packer.pack( - [{"input_ids": [1] * 5}, {"input_ids": [2] * 6}], max_seq_len=10 - ) + batches = packer.pack([{"input_ids": [1] * 5}, {"input_ids": [2] * 6}], max_seq_len=10) assert len(batches) == 2 def test_validate_micro_batches(): """Validation: valid batches pass, missing field / empty / length mismatch fail.""" - valid = [{ - "input_ids": [[1, 2, 3], [4, 5]], - "labels": [[2, 3, 4], [5, 6]], - "position_ids": [[0, 1, 2], [0, 1]], - "request_id": "test", - "batch_id": 0, - }] + valid = [ + { + "input_ids": [[1, 2, 3], [4, 5]], + "labels": [[2, 3, 4], [5, 6]], + "position_ids": [[0, 1, 2], [0, 1]], + "request_id": "test", + "batch_id": 0, + } + ] assert validate_micro_batches(valid) is True # Missing request_id - assert validate_micro_batches([{ - "input_ids": [[1, 2, 3]], "labels": [[2, 3, 4]], - "position_ids": [[0, 1, 2]], "batch_id": 0, - }]) is False + assert ( + validate_micro_batches( + [ + { + "input_ids": [[1, 2, 3]], + "labels": [[2, 3, 4]], + "position_ids": [[0, 1, 2]], + "batch_id": 0, + } + ] + ) + is False + ) # Empty input_ids - assert validate_micro_batches([{ - "input_ids": [], "labels": [], "position_ids": [], - "request_id": "t", "batch_id": 0, - }]) is False + assert ( + validate_micro_batches( + [ + { + "input_ids": [], + "labels": [], + "position_ids": [], + "request_id": "t", + "batch_id": 0, + } + ] + ) + is False + ) # Length mismatch (labels vs input_ids) - assert validate_micro_batches([{ - "input_ids": [[1, 2, 3], [4, 5]], "labels": [[2, 3, 4]], - "position_ids": [[0, 1, 2], [0, 1]], "request_id": "t", "batch_id": 0, - }]) is False + assert ( + validate_micro_batches( + [ + { + "input_ids": [[1, 2, 3], [4, 5]], + "labels": [[2, 3, 4]], + "position_ids": [[0, 1, 2], [0, 1]], + "request_id": "t", + "batch_id": 0, + } + ] + ) + is False + ) # Position_ids length mismatch - assert validate_micro_batches([{ - "input_ids": [[1, 2, 3]], "labels": [[2, 3, 4]], - "position_ids": [[0, 1]], "request_id": "t", "batch_id": 0, - }]) is False + assert ( + validate_micro_batches( + [ + { + "input_ids": [[1, 2, 3]], + "labels": [[2, 3, 4]], + "position_ids": [[0, 1]], + "request_id": "t", + "batch_id": 0, + } + ] + ) + is False + ) def test_pack_samples_function(simple_data): """pack_samples convenience function: with/without packing, default params.""" - batches = pack_samples(simple_data, max_seq_len=10, enable_packing=True, - request_id="func-test", pad_to_multiple_of=1) + batches = pack_samples( + simple_data, max_seq_len=10, enable_packing=True, request_id="func-test", pad_to_multiple_of=1 + ) assert len(batches) == 1 and batches[0]["request_id"] == "func-test" batches_no_pack = pack_samples(simple_data, enable_packing=False, pad_to_multiple_of=1) @@ -221,10 +262,17 @@ def test_packer_abstract_and_custom(): class CustomPacker(Packer): def pack(self, datum_list, max_seq_len, request_id=""): - return [{"input_ids": [d["input_ids"]], "labels": [d.get("labels", [])], - "position_ids": [list(range(len(d["input_ids"])))], - "request_id": request_id, "batch_id": i} - for i, d in enumerate(datum_list) if "input_ids" in d] + return [ + { + "input_ids": [d["input_ids"]], + "labels": [d.get("labels", [])], + "position_ids": [list(range(len(d["input_ids"])))], + "request_id": request_id, + "batch_id": i, + } + for i, d in enumerate(datum_list) + if "input_ids" in d + ] cp = CustomPacker() assert cp.get_name() == "CustomPacker" @@ -284,6 +332,7 @@ def test_large_batch_and_numpy(): # Unpack per-token outputs # ============================================================================ + def test_unpack_per_token_outputs(): """Unpack: no-shift, shift, single/multi sample, 2D tensors, lists, min-length.""" # No-shift: output length == position_ids length @@ -333,6 +382,7 @@ def test_unpack_per_token_outputs(): # Full pipeline roundtrip # ============================================================================ + def test_full_pipeline_roundtrip(): """Pack -> simulate forward -> unpack: token counts preserved, single and multi-batch.""" data = [ diff --git a/tests/server/orchestrator/test_request_processor.py b/tests/server/orchestrator/test_request_processor.py index 23748405..a8e3f197 100644 --- a/tests/server/orchestrator/test_request_processor.py +++ b/tests/server/orchestrator/test_request_processor.py @@ -13,25 +13,23 @@ - Verify RequestProcessor correctly packs data and formats outputs """ -import asyncio import pytest import pytest_asyncio -from typing import List, Dict, Any, Optional -from xorl.server.orchestrator.request_processor import RequestProcessor from xorl.server.backend import DummyBackend +from xorl.server.orchestrator.request_processor import RequestProcessor from xorl.server.protocol.api_orchestrator import ( - OrchestratorRequest, OrchestratorOutputs, - RequestType, + OrchestratorRequest, OutputType, + RequestType, ) from xorl.server.protocol.operations import ( + EmptyData, + LoadStateData, ModelPassData, OptimStepData, SaveStateData, - LoadStateData, - EmptyData, ) @@ -39,6 +37,7 @@ # Fixtures # ============================================================================ + @pytest_asyncio.fixture async def processor(): """Create and start processor with DummyBackend.""" @@ -58,6 +57,7 @@ async def processor(): # Tests # ============================================================================ + @pytest.mark.asyncio async def test_lifecycle_and_ready_state(): """Test processor start, stop, and ready state.""" @@ -81,7 +81,8 @@ async def test_forward_backward_operations(processor): {"input_ids": [100, 200, 300], "labels": [200, 300, 400]}, ] request = OrchestratorRequest( - request_id="req-001", request_type=RequestType.ADD, + request_id="req-001", + request_type=RequestType.ADD, operation="forward_backward", payload=ModelPassData(data=datum_list, loss_fn="causallm_loss"), ) @@ -97,7 +98,8 @@ async def test_forward_backward_operations(processor): # Forward only (no gradients) request = OrchestratorRequest( - request_id="req-fwd", request_type=RequestType.ADD, + request_id="req-fwd", + request_type=RequestType.ADD, operation="forward", payload=ModelPassData(data=[{"input_ids": [1, 2, 3], "labels": [2, 3, 4]}]), ) @@ -111,8 +113,10 @@ async def test_optim_and_checkpoint_operations(processor): """Test optim_step, save_state, load_state, sleep, and wake_up.""" # Optim step request = OrchestratorRequest( - request_id="req-004", request_type=RequestType.ADD, - operation="optim_step", payload=OptimStepData(lr=0.001, gradient_clip=1.0), + request_id="req-004", + request_type=RequestType.ADD, + operation="optim_step", + payload=OptimStepData(lr=0.001, gradient_clip=1.0), ) output = await processor.execute_optim_step(request) assert output.output_type == OutputType.OPTIM_STEP @@ -121,8 +125,10 @@ async def test_optim_and_checkpoint_operations(processor): # Save state request = OrchestratorRequest( - request_id="req-save", request_type=RequestType.ADD, - operation="save_state", payload=SaveStateData(checkpoint_path="/tmp/ckpt"), + request_id="req-save", + request_type=RequestType.ADD, + operation="save_state", + payload=SaveStateData(checkpoint_path="/tmp/ckpt"), ) output = await processor.execute_save_state(request) assert output.output_type == OutputType.SAVE_STATE @@ -130,8 +136,10 @@ async def test_optim_and_checkpoint_operations(processor): # Load state request = OrchestratorRequest( - request_id="req-load", request_type=RequestType.ADD, - operation="load_state", payload=LoadStateData(checkpoint_path="/tmp/ckpt"), + request_id="req-load", + request_type=RequestType.ADD, + operation="load_state", + payload=LoadStateData(checkpoint_path="/tmp/ckpt"), ) output = await processor.execute_load_state(request) assert output.output_type == OutputType.LOAD_STATE @@ -139,16 +147,20 @@ async def test_optim_and_checkpoint_operations(processor): # Sleep request = OrchestratorRequest( - request_id="req-sleep", request_type=RequestType.ADD, - operation="sleep", payload=EmptyData(), + request_id="req-sleep", + request_type=RequestType.ADD, + operation="sleep", + payload=EmptyData(), ) output = await processor.execute_sleep(request) assert output.output_type == OutputType.SLEEP # Wake up request = OrchestratorRequest( - request_id="req-wake", request_type=RequestType.ADD, - operation="wake_up", payload=EmptyData(), + request_id="req-wake", + request_type=RequestType.ADD, + operation="wake_up", + payload=EmptyData(), ) output = await processor.execute_wake_up(request) assert output.output_type == OutputType.WAKE_UP @@ -160,7 +172,8 @@ async def test_statistics_tracking(processor): initial = processor.total_operations request = OrchestratorRequest( - request_id="req-stat", request_type=RequestType.ADD, + request_id="req-stat", + request_type=RequestType.ADD, operation="forward_backward", payload=ModelPassData(data=[{"input_ids": [1, 2], "labels": [2, 3]}]), ) @@ -178,15 +191,18 @@ async def test_error_handling(processor): """Test error handling for empty datum list, missing labels, and sequential ops.""" # Empty datum list request = OrchestratorRequest( - request_id="req-008", request_type=RequestType.ADD, - operation="forward_backward", payload=ModelPassData(data=[]), + request_id="req-008", + request_type=RequestType.ADD, + operation="forward_backward", + payload=ModelPassData(data=[]), ) output = await processor.execute_forward_backward(request) assert output.output_type == OutputType.ERROR # Without labels (no valid tokens) request = OrchestratorRequest( - request_id="req-009", request_type=RequestType.ADD, + request_id="req-009", + request_type=RequestType.ADD, operation="forward_backward", payload=ModelPassData(data=[{"input_ids": [1, 2, 3, 4, 5]}]), ) @@ -196,7 +212,8 @@ async def test_error_handling(processor): # Multiple sequential operations for i in range(5): request = OrchestratorRequest( - request_id=f"req-seq-{i}", request_type=RequestType.ADD, + request_id=f"req-seq-{i}", + request_type=RequestType.ADD, operation="forward_backward", payload=ModelPassData(data=[{"input_ids": list(range(10)), "labels": list(range(1, 11))}]), ) diff --git a/tests/server/orchestrator/test_scheduler.py b/tests/server/orchestrator/test_scheduler.py index 28ec5de6..69ed7cfd 100644 --- a/tests/server/orchestrator/test_scheduler.py +++ b/tests/server/orchestrator/test_scheduler.py @@ -11,14 +11,14 @@ import pytest + pytestmark = [pytest.mark.cpu, pytest.mark.server] import time from xorl.server.orchestrator.scheduler import ( - Scheduler, FIFOPolicy, ScheduledRequest, - SchedulingPolicy, + Scheduler, ) from xorl.server.protocol.api_orchestrator import ( OrchestratorRequest, @@ -202,17 +202,29 @@ def test_fifo_policy_order_remove_and_clear(): """Test FIFO ordering, remove, and clear operations.""" policy = FIFOPolicy() for i in range(5): - policy.add_request(ScheduledRequest(request=OrchestratorRequest( - request_id=f"req-{i}", request_type=RequestType.ADD, operation="test", - ))) + policy.add_request( + ScheduledRequest( + request=OrchestratorRequest( + request_id=f"req-{i}", + request_type=RequestType.ADD, + operation="test", + ) + ) + ) for i in range(5): assert policy.get_next_request().request_id == f"req-{i}" for i in range(5): - policy.add_request(ScheduledRequest(request=OrchestratorRequest( - request_id=f"r2-{i}", request_type=RequestType.ADD, operation="test", - ))) + policy.add_request( + ScheduledRequest( + request=OrchestratorRequest( + request_id=f"r2-{i}", + request_type=RequestType.ADD, + operation="test", + ) + ) + ) assert policy.remove_request("r2-0") is True and policy.size() == 4 assert policy.get_next_request().request_id == "r2-1" diff --git a/tests/server/runner/test_send_ready.py b/tests/server/runner/test_send_ready.py index 0635489b..0307a352 100644 --- a/tests/server/runner/test_send_ready.py +++ b/tests/server/runner/test_send_ready.py @@ -11,19 +11,19 @@ import pytest + pytestmark = [pytest.mark.cpu, pytest.mark.server] +from xorl.server.protocol.operations import ( + EmptyData, + ModelPassData, +) from xorl.server.protocol.orchestrator_runner import ( - RunnerDispatchCommand, - MessageType, RunnerAck, + RunnerDispatchCommand, RunnerReady, - serialize_message, deserialize_message, -) -from xorl.server.protocol.operations import ( - EmptyData, - ModelPassData, + serialize_message, ) from xorl.server.runner.utils.rank0_protocol import Rank0Protocol diff --git a/tests/server/weight_sync/test_pp_nccl_transfer.py b/tests/server/weight_sync/test_pp_nccl_transfer.py index db7854d3..1af81cb3 100644 --- a/tests/server/weight_sync/test_pp_nccl_transfer.py +++ b/tests/server/weight_sync/test_pp_nccl_transfer.py @@ -5,9 +5,9 @@ requiring a full distributed environment. """ -import pytest +from unittest.mock import MagicMock, patch + import torch -from unittest.mock import patch, MagicMock from xorl.server.weight_sync.handler import _prod @@ -38,6 +38,7 @@ def _make_handler(self, rank): handler.rank = rank # Bind the real method from xorl.server.weight_sync.handler import WeightSyncHandler + handler._pp_nccl_transfer_buffer = WeightSyncHandler._pp_nccl_transfer_buffer.__get__(handler) return handler @@ -80,6 +81,7 @@ def test_receiver_gets_empty_for_empty_send(self, mock_dist): # Simulate broadcast_object_list delivering empty metadata def fake_broadcast_object_list(obj, src, group): obj[0] = [] # empty metadata from sender + mock_dist.broadcast_object_list = fake_broadcast_object_list result = handler._pp_nccl_transfer_buffer(None, MagicMock(), src_global, "cuda:0") From 6d3611355dd6c87fd4c194ef05b7cc7cf93fb46e Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Thu, 2 Apr 2026 13:01:22 -0700 Subject: [PATCH 08/41] Improve dummy dataset generation and packing (#49) * Improve dummy dataset generation and packing - Update dummy data pipeline for better benchmark coverage - Improve dataset packing logic for local testing * Fix pre-commit formatting in packing.py (cherry picked from commit 96e918147fc6e2c23d13f96d601f42437b67beec) --- src/xorl/data/prepare/packing.py | 68 +++++++++++------------ src/xorl/data/prepare/prepare_datasets.py | 39 ++++--------- 2 files changed, 44 insertions(+), 63 deletions(-) diff --git a/src/xorl/data/prepare/packing.py b/src/xorl/data/prepare/packing.py index ac780339..8d1f1b6e 100644 --- a/src/xorl/data/prepare/packing.py +++ b/src/xorl/data/prepare/packing.py @@ -4,6 +4,7 @@ """ import os +import shutil import time from concurrent.futures import ProcessPoolExecutor from multiprocessing import cpu_count, get_context @@ -12,6 +13,7 @@ import numba import numpy as np +import torch.distributed as dist from datasets import Dataset as HFDataset from torch.utils.data import Dataset from transformers import PreTrainedTokenizer @@ -436,6 +438,8 @@ def _save_bins_cache(self, bins: List[List[int]]) -> None: # Ensure the directory exists cache_path.parent.mkdir(parents=True, exist_ok=True) + if cache_path.exists(): + shutil.rmtree(cache_path, ignore_errors=True) # Create a HF dataset from the bins bins_dataset = HFDataset.from_dict({"bins": bins}) @@ -457,8 +461,7 @@ def _load_or_compute_bins(self) -> List[List[int]]: """Load cached bins or compute new ones with distributed coordination.""" # Get rank from environment variable (set by torchrun/distributed launcher) rank = int(os.environ.get("RANK", "0")) - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - world_size = int(os.environ.get("WORLD_SIZE", "1")) + is_distributed = dist.is_available() and dist.is_initialized() # Try to load from cache first cached_bins = self._load_cached_bins() @@ -476,43 +479,40 @@ def _load_or_compute_bins(self) -> List[List[int]]: # Save to cache for future use if cache_path is not None: self._save_bins_cache(bins) + if is_distributed: + dist.barrier() else: # Other ranks wait for rank 0 to finish if cache_path is not None: LOG.info(f"Rank {rank} waiting for rank 0 to compute bins...") - # Wait for cache file to be created - max_wait_time = 3600 # 1 hour max wait - wait_interval = 5 # Check every 5 seconds - waited = 0 - - while not cache_path.exists() and waited < max_wait_time: - time.sleep(wait_interval) - waited += wait_interval - - if cache_path.exists(): - # Wait a bit more to ensure rank 0 finished writing - # (HF datasets creates the directory first, then writes files) - time.sleep(2) - - # Try to load with retries - bins = None - max_retries = 30 - for attempt in range(max_retries): - bins = self._load_cached_bins() - if bins is not None: - break - if attempt < max_retries - 1: - LOG.info(f"Rank {rank} retry {attempt + 1}/{max_retries} loading bins...") - time.sleep(5) - - if bins is None: - raise RuntimeError( - f"Rank {rank} failed to load bins cached by rank 0 after {max_retries} retries. " - f"Cache path: {cache_path}" - ) - LOG.info(f"Rank {rank} loaded {len(bins)} bins from cache") + if is_distributed: + dist.barrier() else: - raise RuntimeError(f"Rank {rank} timed out waiting for rank 0 to compute bins") + max_wait_time = 3600 + wait_interval = 5 + waited = 0 + while not cache_path.exists() and waited < max_wait_time: + time.sleep(wait_interval) + waited += wait_interval + if not cache_path.exists(): + raise RuntimeError(f"Rank {rank} timed out waiting for rank 0 to compute bins") + + bins = None + max_retries = 10 + for attempt in range(max_retries): + bins = self._load_cached_bins() + if bins is not None: + break + if attempt < max_retries - 1: + LOG.info(f"Rank {rank} retry {attempt + 1}/{max_retries} loading bins...") + time.sleep(2) + + if bins is None: + raise RuntimeError( + f"Rank {rank} failed to load bins cached by rank 0 after {max_retries} retries. " + f"Cache path: {cache_path}" + ) + LOG.info(f"Rank {rank} loaded {len(bins)} bins from cache") else: # No cache path, fall back to computing on all ranks LOG.warning(f"No cache path available, rank {rank} computing bins independently") diff --git a/src/xorl/data/prepare/prepare_datasets.py b/src/xorl/data/prepare/prepare_datasets.py index ecce8a2e..f1f5533c 100644 --- a/src/xorl/data/prepare/prepare_datasets.py +++ b/src/xorl/data/prepare/prepare_datasets.py @@ -47,48 +47,29 @@ def _create_dummy_dataset(seq_len: int, num_samples: int = 4096, seed: int = 42, step entirely. Token format: each sample is ``[1, 2, 3, ..., length-1, 0]`` where - 0 is the EOD marker. Tokens follow a global arithmetic sequence - ``(global_offset + i) % vocab_size`` so routing varies across samples. - Lengths are drawn uniformly from [1, seq_len] with seed=0 (independent - of the dataset seed so packing shape is always the same). + 0 is the EOD marker. All ``num_samples`` examples are intentionally + identical so local benchmarks can show clean loss convergence. """ import numpy as np import pyarrow as pa EOD = 0 VOCAB_SIZE = vocab_size - # Cap at seq_len so samples always fit in sample_packing_sequence_len. MAX_SAMPLE_LEN = max(seq_len, 1) - # Lengths: uniform [1, max_sample_len], fixed seed=0 so packing is - # reproducible regardless of the dataset seed argument. # Lengths are k*16+1 so that after ShiftTokensCollator drops 1 token, # effective length k*16 is divisible by 2*ringattn_size for zigzag ring - # attention. This is a dummy-data workaround; production data should use - # per-document alignment in the packing pipeline (TODO). + # attention. This keeps the repeated sample valid for the benchmark path. LENGTH_ALIGN = 16 - max_buckets = max(1, MAX_SAMPLE_LEN // LENGTH_ALIGN) - len_rng = np.random.RandomState(0) - lengths = len_rng.randint(1, max_buckets + 1, size=num_samples) - lengths = (lengths * LENGTH_ALIGN + 1).astype(np.int64) - - # Build all tokens as one global arithmetic sequence, then slice. - # Each sample: (global_offset, ..., global_offset+length-2, EOD) - total_tokens = int(lengths.sum()) - global_seq = np.arange(total_tokens, dtype=np.int64) % VOCAB_SIZE - - # Overwrite the last position of each sample with EOD (vectorized) - ends = np.cumsum(lengths) - 1 - global_seq[ends] = EOD - - flat_tokens = global_seq.astype(np.int32) - offsets = np.concatenate([[0], np.cumsum(lengths)]).astype(np.int64) + sample_len = max(1, ((MAX_SAMPLE_LEN - 1) // LENGTH_ALIGN) * LENGTH_ALIGN + 1) + lengths = np.full((num_samples,), sample_len, dtype=np.int64) + sample_tokens = (np.arange(sample_len, dtype=np.int64) + 1) % VOCAB_SIZE + sample_tokens[-1] = EOD + flat_tokens = np.tile(sample_tokens.astype(np.int32), num_samples) + offsets = np.arange(0, sample_len * (num_samples + 1), sample_len, dtype=np.int64) tokens = pa.ListArray.from_arrays(offsets, flat_tokens) - # position_ids: [0, 1, ..., length-1] per sample — vectorized via cumsum reset - pos_flat = np.arange(total_tokens, dtype=np.int32) - starts = np.concatenate([[0], np.cumsum(lengths[:-1])]).astype(np.int64) - pos_flat -= np.repeat(starts, lengths).astype(np.int32) + pos_flat = np.tile(np.arange(sample_len, dtype=np.int32), num_samples) position_ids = pa.ListArray.from_arrays(offsets, pos_flat) pa_lengths = pa.array(lengths.tolist(), type=pa.int64()) From 46eca4a914e45809eee320d7ba3c0582ae1b9ee2 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 2 Apr 2026 17:54:10 -0700 Subject: [PATCH 09/41] Fuse MoE expert gate and up projections (#64) * Fuse MoE expert gate and up projections * Fix EP MoE weight sync gathering * Fix weight sync readiness tracking (cherry picked from commit 3ade637fe432c005f1b613a6d56423e1bb7d17fd) --- .../server/password_memorization/README.md | 4 + .../run_password_test.py | 144 ++++++++++++++---- .../models/checkpoint_handlers/buffers.py | 11 +- src/xorl/models/layers/moe/common.py | 24 +++ src/xorl/models/layers/moe/experts.py | 53 ++++--- src/xorl/models/layers/moe/lora.py | 80 +++++++--- .../qwen3_5_moe/checkpoint_handler.py | 115 +++++++++++--- .../transformers/qwen3_5_moe/parallelize.py | 5 +- .../qwen3_moe/checkpoint_handler.py | 133 ++++++++++++++-- .../transformers/qwen3_moe/parallelize.py | 5 +- src/xorl/qlora/modules/moe_experts.py | 5 + src/xorl/qlora/utils.py | 24 ++- src/xorl/server/weight_sync/backends/base.py | 4 + .../weight_sync/backends/nccl_broadcast.py | 10 +- src/xorl/server/weight_sync/handler.py | 125 ++++++++++----- tests/models/test_moe_fused_gate_up_proj.py | 134 ++++++++++++++++ .../test_weight_version_forwarding.py | 52 +++++++ 17 files changed, 768 insertions(+), 160 deletions(-) create mode 100644 src/xorl/models/layers/moe/common.py create mode 100644 tests/models/test_moe_fused_gate_up_proj.py create mode 100644 tests/server/weight_sync/test_weight_version_forwarding.py diff --git a/examples/server/password_memorization/README.md b/examples/server/password_memorization/README.md index 6ffd3391..7c4fba31 100644 --- a/examples/server/password_memorization/README.md +++ b/examples/server/password_memorization/README.md @@ -60,6 +60,8 @@ python run_password_test.py --model Qwen/Qwen3-8B --steps 48 --lr 5e-5 --sync-qu | `--infer-url` | http://localhost:30000 | Inference endpoint URL(s), space-separated | | `--master-address` | localhost | Master address for NCCL weight sync | | `--log-interval` | 16 | Print loss every N steps | +| `--run-baseline` | false | Run a pre-training inference check before weight sync | +| `--sync-wait-timeout` | 120 | Seconds to wait for inference endpoints to report the new `weight_version` | --- @@ -148,3 +150,5 @@ FP8 re-quantization uses block-wise e4m3 with `weight_block_size=[128, 128]`, co - **Qwen3-235B**: use the instruct tokenizer for training — the base model chat template injects `` tags even with `enable_thinking=False`. - **30B MoE with EP**: cosine LR schedule is important — constant LR causes loss oscillation around ~1.0. - **NF4**: quantizes bf16 weights on-the-fly; no pre-quantized checkpoint needed. +- **Password e2e validation**: `run_password_test.py` now waits for the async `create_model` future and for each inference endpoint to report the requested `weight_version` before running post-sync recall checks. +- **Baseline queries**: the script skips baseline inference by default to avoid seeding radix-cache entries for the exact prompts later used in the no-flush post-sync validation. Pass `--run-baseline` if you explicitly want that pre-training check. diff --git a/examples/server/password_memorization/run_password_test.py b/examples/server/password_memorization/run_password_test.py index 4f6f9fb8..85977ffa 100644 --- a/examples/server/password_memorization/run_password_test.py +++ b/examples/server/password_memorization/run_password_test.py @@ -31,12 +31,15 @@ import argparse import math import time +import traceback from urllib.parse import urlparse import requests from transformers import AutoTokenizer +MODEL_ID = "default" + CODES = { "project_alpha": "SUNRISE-7742-DRAGON", "project_beta": "MOUNTAIN-3391-RIVER", @@ -86,10 +89,19 @@ def build_training_data(tokenizer): # --------------------------------------------------------------------------- +def _raise_on_failed_future(result, context): + if result.get("type") == "request_failed": + raise RuntimeError(f"{context} failed: {result.get('error', result)}") + if result.get("error"): + raise RuntimeError(f"{context} failed: {result['error']}") + return result + + def wait_for_future(train_url, request_id, timeout=600): deadline = time.time() + timeout while time.time() < deadline: resp = requests.post(f"{train_url}/api/v1/retrieve_future", json={"request_id": request_id}, timeout=120) + resp.raise_for_status() result = resp.json() if result.get("type") == "try_again": time.sleep(0.5) @@ -98,11 +110,13 @@ def wait_for_future(train_url, request_id, timeout=600): raise TimeoutError(f"Future {request_id} timed out after {timeout}s") -def wait_for_service(url, timeout=300): +def wait_for_training_service(train_url, timeout=300): deadline = time.time() + timeout while time.time() < deadline: try: - if requests.get(f"{url}/health", timeout=5).status_code == 200: + resp = requests.get(f"{train_url}/health", timeout=5) + resp.raise_for_status() + if resp.json().get("engine_running"): return True except Exception: pass @@ -110,10 +124,28 @@ def wait_for_service(url, timeout=300): return False +def wait_for_inference_service(infer_url, timeout=300): + deadline = time.time() + timeout + while time.time() < deadline: + for endpoint in ("/health", "/model_info", "/v1/models"): + try: + resp = requests.get(f"{infer_url}{endpoint}", timeout=5) + resp.raise_for_status() + return True + except Exception: + pass + time.sleep(3) + return False + + def create_model(train_url, model_name): - return requests.post( - f"{train_url}/api/v1/create_model", json={"model_id": "default", "base_model": model_name}, timeout=30 - ).json() + resp = requests.post( + f"{train_url}/api/v1/create_model", json={"model_id": MODEL_ID, "base_model": model_name}, timeout=30 + ) + resp.raise_for_status() + future = resp.json() + result = wait_for_future(train_url, future["request_id"]) + return _raise_on_failed_future(result, "create_model") def add_endpoints(train_url, infer_urls): @@ -123,6 +155,7 @@ def add_endpoints(train_url, infer_urls): resp = requests.post( f"{train_url}/add_inference_endpoint", json={"host": host, "port": port, "worker_port": port}, timeout=30 ) + resp.raise_for_status() result = resp.json() si = result.get("endpoint", {}).get("server_info", {}) if result else {} print( @@ -134,17 +167,19 @@ def add_endpoints(train_url, infer_urls): def train_step(train_url, data, lr): fb = requests.post( f"{train_url}/api/v1/forward_backward", - json={"model_id": "default", "forward_backward_input": {"data": data, "loss_fn": "causallm_loss"}}, + json={"model_id": MODEL_ID, "forward_backward_input": {"data": data, "loss_fn": "causallm_loss"}}, timeout=30, ) - fb_result = wait_for_future(train_url, fb.json()["request_id"]) + fb.raise_for_status() + fb_result = _raise_on_failed_future(wait_for_future(train_url, fb.json()["request_id"]), "forward_backward") loss = fb_result.get("metrics", {}).get("loss:mean", "N/A") opt = requests.post( f"{train_url}/api/v1/optim_step", - json={"model_id": "default", "adam_params": {"learning_rate": lr}, "gradient_clip": 1.0}, + json={"model_id": MODEL_ID, "adam_params": {"learning_rate": lr}, "gradient_clip": 1.0}, timeout=30, ) - opt_result = wait_for_future(train_url, opt.json()["request_id"]) + opt.raise_for_status() + opt_result = _raise_on_failed_future(wait_for_future(train_url, opt.json()["request_id"]), "optim_step") grad_norm = opt_result.get("metrics", {}).get("grad_norm", "N/A") return loss, grad_norm @@ -152,27 +187,57 @@ def train_step(train_url, data, lr): def set_sync_quantization(train_url): config = {"quant_method": "fp8", "fmt": "e4m3", "weight_block_size": [128, 128]} resp = requests.post(f"{train_url}/api/v1/set_sync_quantization", json={"quantization": config}, timeout=10) + resp.raise_for_status() print(f" Set sync quantization: {resp.json().get('message')}") -def sync_weights(train_url, master_address): +def sync_weights(train_url, master_address, weight_version): t0 = time.time() resp = requests.post( - f"{train_url}/api/v1/sync_inference_weights", json={"master_address": master_address}, timeout=600 + f"{train_url}/api/v1/sync_inference_weights", + json={"master_address": master_address, "weight_version": weight_version}, + timeout=600, ) + resp.raise_for_status() result = resp.json() print( f" Sync: success={result.get('success')}, {time.time() - t0:.1f}s, " - f"params={result.get('num_parameters', 'N/A')}" + f"params={result.get('num_parameters', 'N/A')}, weight_version={weight_version}" ) return result +def get_model_info(infer_url): + resp = requests.get(f"{infer_url}/model_info", timeout=10) + resp.raise_for_status() + return resp.json() + + +def wait_for_inference_weight_version(infer_urls, expected_version, timeout=120): + deadline = time.time() + timeout + last_seen = {} + while time.time() < deadline: + all_ready = True + for url in infer_urls: + try: + info = get_model_info(url) + version = info.get("weight_version") + last_seen[url] = version + if version != expected_version: + all_ready = False + except Exception: + all_ready = False + if all_ready: + return last_seen + time.sleep(1) + raise TimeoutError(f"Inference endpoints did not reach weight_version={expected_version}: {last_seen}") + + def query_inference(infer_url, prompt): resp = requests.post( f"{infer_url}/v1/chat/completions", json={ - "model": "default", + "model": MODEL_ID, "messages": [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt}], "max_tokens": 32, "temperature": 0, @@ -180,21 +245,27 @@ def query_inference(infer_url, prompt): }, timeout=30, ) - content = resp.json()["choices"][0]["message"]["content"] - return content.strip() if content else "(empty response)" + resp.raise_for_status() + payload = resp.json() + content = payload["choices"][0]["message"]["content"] + metadata = payload.get("metadata", {}) + return (content.strip() if content else "(empty response)"), metadata -def test_inference(infer_urls, label=""): +def test_inference(infer_urls, label="", expected_weight_version=None): print(f" Inference ({label}):") total_correct = 0 for url in infer_urls: port = url.split(":")[-1] correct = 0 for project, expected in CODES.items(): - answer = query_inference(url, f"What is the secret code for {project}?") - match = expected in answer + answer, metadata = query_inference(url, f"What is the secret code for {project}?") + version = metadata.get("weight_version") + version_ok = expected_weight_version is None or version == expected_weight_version + match = expected in answer and version_ok correct += match - print(f" [{'OK' if match else 'FAIL'}] :{port} {project}: '{answer}'") + version_suffix = f", version={version}" if version is not None else "" + print(f" [{'OK' if match else 'FAIL'}] :{port} {project}: '{answer}'{version_suffix}") total_correct += correct return total_correct @@ -236,18 +307,22 @@ def run_test(args, training_data): total_codes = len(CODES) * len(infer_urls) print(" Checking services...") - if not wait_for_service(args.train_url, timeout=10): + if not wait_for_training_service(args.train_url, timeout=10): print(" FAILED: Training server not running") return None for url in infer_urls: - if not wait_for_service(url, timeout=10): + if not wait_for_inference_service(url, timeout=10): print(f" FAILED: {url} not running") return None print(" All services ready.") - create_model(args.train_url, args.model) + create_result = create_model(args.train_url, args.model) + print(f" Model created: model_id={create_result.get('model_id', MODEL_ID)}") add_endpoints(args.train_url, infer_urls) - test_inference(infer_urls, "baseline") + if args.run_baseline: + test_inference(infer_urls, "baseline") + else: + print(" Skipping baseline inference to avoid seeding stale prompt cache before sync.") print(f"\n Training ({args.steps} steps, lr={args.lr}, schedule={args.lr_schedule})...") t0 = time.time() @@ -261,12 +336,20 @@ def run_test(args, training_data): if args.sync_quant == "fp8": set_sync_quantization(args.train_url) - sync_result = sync_weights(args.train_url, args.master_address) + sync_version = f"password-sync-{int(time.time())}" + sync_result = sync_weights(args.train_url, args.master_address, sync_version) if not sync_result.get("success"): print(f" SYNC FAILED: {sync_result}") return None + versions = wait_for_inference_weight_version( + infer_urls, + sync_version, + timeout=args.sync_wait_timeout, + ) + for url, version in versions.items(): + print(f" Endpoint ready: {url} weight_version={version}") - correct = test_inference(infer_urls, "after sync") + correct = test_inference(infer_urls, f"after sync ({sync_version})", expected_weight_version=sync_version) print(f" Score: {correct}/{total_codes}") return correct @@ -302,6 +385,15 @@ def main(): ) parser.add_argument("--master-address", type=str, default="localhost", help="Master address for NCCL weight sync") parser.add_argument("--log-interval", type=int, default=16, help="Print loss every N steps") + parser.add_argument( + "--run-baseline", action="store_true", help="Run a pre-training inference check before weight sync" + ) + parser.add_argument( + "--sync-wait-timeout", + type=int, + default=120, + help="Seconds to wait for inference endpoints to report the new weight_version", + ) args = parser.parse_args() tokenizer = AutoTokenizer.from_pretrained(args.model) @@ -315,8 +407,6 @@ def main(): correct = run_test(args, training_data) except Exception as e: print(f" ERROR: {e}") - import traceback - traceback.print_exc() correct = None diff --git a/src/xorl/models/checkpoint_handlers/buffers.py b/src/xorl/models/checkpoint_handlers/buffers.py index 76771c7b..e74e86d4 100644 --- a/src/xorl/models/checkpoint_handlers/buffers.py +++ b/src/xorl/models/checkpoint_handlers/buffers.py @@ -35,8 +35,8 @@ _EXPERT_ANY_KEY_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.(.+)$") # Model parameter names in fused expert format -# e.g., model.layers.0.mlp.experts.gate_proj -FUSED_EXPERT_PATTERN = re.compile(r"^model\.layers\.\d+\.mlp\.experts\.(gate|up|down)_proj$") +# e.g., model.layers.0.mlp.experts.gate_up_proj +FUSED_EXPERT_PATTERN = re.compile(r"^model\.layers\.\d+\.mlp\.experts\.(gate_up|gate|up|down)_proj$") # Dense MLP gate/up projection weight keys # e.g., model.layers.0.mlp.gate_proj.weight or model.layers.0.mlp.shared_expert.up_proj.weight @@ -304,7 +304,7 @@ def model_needs_expert_merging(parameter_names: Set[str]) -> bool: def model_has_gate_up_merged(parameter_names: Set[str]) -> bool: """Check if the model uses merged gate_up_proj layers.""" - return any(".gate_up_proj." in name for name in parameter_names) + return any(".gate_up_proj." in name or name.endswith(".gate_up_proj") for name in parameter_names) def checkpoint_has_separate_gate_up(checkpoint_keys: Set[str]) -> bool: @@ -416,6 +416,11 @@ def get_fused_name(layer_idx: int, proj: str) -> str: """Get the fused parameter name, e.g., 'model.layers.0.mlp.experts.gate_proj'.""" return f"model.layers.{layer_idx}.mlp.experts.{proj}_proj" + @staticmethod + def get_gate_up_name(layer_idx: int) -> str: + """Get the fused gate/up parameter name.""" + return f"model.layers.{layer_idx}.mlp.experts.gate_up_proj" + def count_skipped(self, layer_idx: int, proj: str) -> None: """Count an expert key as seen without buffering tensor data. diff --git a/src/xorl/models/layers/moe/common.py b/src/xorl/models/layers/moe/common.py new file mode 100644 index 00000000..491174bb --- /dev/null +++ b/src/xorl/models/layers/moe/common.py @@ -0,0 +1,24 @@ +"""Shared helpers for MoE expert weight layouts.""" + +import torch + + +def split_gate_up_proj( + gate_up_proj: torch.Tensor, + intermediate_size: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Split a fused ``gate_up_proj`` tensor into gate and up views.""" + if intermediate_size is None: + intermediate_size = gate_up_proj.shape[-1] // 2 + return ( + gate_up_proj[..., :intermediate_size], + gate_up_proj[..., intermediate_size:], + ) + + +def fuse_gate_up_proj( + gate_proj: torch.Tensor, + up_proj: torch.Tensor, +) -> torch.Tensor: + """Concatenate separate gate/up tensors into fused ``gate_up_proj`` format.""" + return torch.cat([gate_proj, up_proj], dim=-1) diff --git a/src/xorl/models/layers/moe/experts.py b/src/xorl/models/layers/moe/experts.py index 417687d5..ee3fad9f 100644 --- a/src/xorl/models/layers/moe/experts.py +++ b/src/xorl/models/layers/moe/experts.py @@ -7,6 +7,7 @@ from ..activations import ACT2FN from .backend import MOE_EXPERT_BACKENDS, MOE_EXPERT_BACKENDS_MOE_ACT +from .common import split_gate_up_proj def _flag_enabled(name: str) -> bool: @@ -26,9 +27,11 @@ class MoEExperts(nn.Module): Weights are stored in ``(G, K, N)`` format — ``[num_experts, in_features, out_features]``:: - gate_proj: [num_experts, hidden_dim, intermediate_size] - up_proj: [num_experts, hidden_dim, intermediate_size] - down_proj: [num_experts, intermediate_size, hidden_dim] + gate_up_proj: [num_experts, hidden_dim, 2 * intermediate_size] + down_proj: [num_experts, intermediate_size, hidden_dim] + + ``gate_proj`` and ``up_proj`` are exposed as views into ``gate_up_proj`` + for compatibility with existing backends and helpers. Args: num_experts: Total number of experts. @@ -52,12 +55,8 @@ def __init__( self.intermediate_size = intermediate_size self.moe_implementation = moe_implementation - self.gate_proj = nn.Parameter( - torch.empty(num_experts, hidden_dim, intermediate_size), - requires_grad=True, - ) - self.up_proj = nn.Parameter( - torch.empty(num_experts, hidden_dim, intermediate_size), + self.gate_up_proj = nn.Parameter( + torch.empty(num_experts, hidden_dim, 2 * intermediate_size), requires_grad=True, ) self.down_proj = nn.Parameter( @@ -75,6 +74,20 @@ def __init__( self.deepep_num_sms: int = 20 self.deepep_async_combine: bool = False + @property + def gate_proj(self) -> torch.Tensor: + gate_proj, _ = split_gate_up_proj(self.gate_up_proj, self.intermediate_size) + gate_proj.grad = ( + None if self.gate_up_proj.grad is None else self.gate_up_proj.grad[..., : self.intermediate_size] + ) + return gate_proj + + @property + def up_proj(self) -> torch.Tensor: + _, up_proj = split_gate_up_proj(self.gate_up_proj, self.intermediate_size) + up_proj.grad = None if self.gate_up_proj.grad is None else self.gate_up_proj.grad[..., self.intermediate_size :] + return up_proj + def forward( self, hidden_states: torch.Tensor, @@ -91,6 +104,8 @@ def forward( use the unified dispatch → compute → combine path via ``_ep_forward()``. """ _moe_act = self._moe_act + gate_proj = self.gate_proj.contiguous() + up_proj = self.up_proj.contiguous() if self.moe_implementation == "eager": fn = MOE_EXPERT_BACKENDS[self.moe_implementation] @@ -98,8 +113,8 @@ def forward( return fn( hidden_states, expert_idx, - self.gate_proj, - self.up_proj, + gate_proj, + up_proj, self.down_proj, self.act_fn, ) @@ -122,8 +137,8 @@ def forward( hidden_states, routing_weights, selected_experts, - self.gate_proj, - self.up_proj, + gate_proj, + up_proj, self.down_proj, num_experts=self.num_experts, ) @@ -165,6 +180,8 @@ def _ep_forward( compute_fn = EP_EXPERT_COMPUTE_MOE_ACT[self.moe_implementation] else: compute_fn = EP_EXPERT_COMPUTE[self.moe_implementation] + gate_proj = self.gate_proj.contiguous() + up_proj = self.up_proj.contiguous() # Step 1: Dispatch tokens to expert-owning ranks dispatch_kwargs = self._build_dispatch_kwargs(hidden_states, routing_weights, selected_experts, parallel_state) @@ -187,8 +204,8 @@ def _ep_forward( expert_output = compute_fn( permute_tokens, cumsum, - self.gate_proj, - self.up_proj, + gate_proj, + up_proj, self.down_proj, ) @@ -219,8 +236,8 @@ def _ep_forward_debug(self, dispatch_fn, combine_fn, compute_fn, dispatch_kwargs expert_output = compute_fn( permute_tokens, cumsum, - self.gate_proj, - self.up_proj, + self.gate_proj.contiguous(), + self.up_proj.contiguous(), self.down_proj, ) ev[3].record() @@ -273,7 +290,7 @@ def _build_dispatch_kwargs(self, hidden_states, routing_weights, selected_expert buffer_size_gb=self.deepep_buffer_size_gb, num_sms=self.deepep_num_sms, ) - kwargs["num_local_experts"] = self.gate_proj.shape[0] + kwargs["num_local_experts"] = self.gate_up_proj.shape[0] return kwargs def _build_combine_kwargs(self, expert_output, ctx, dispatch_kwargs, parallel_state): diff --git a/src/xorl/models/layers/moe/lora.py b/src/xorl/models/layers/moe/lora.py index d29fb86c..5405d2cf 100644 --- a/src/xorl/models/layers/moe/lora.py +++ b/src/xorl/models/layers/moe/lora.py @@ -26,6 +26,7 @@ from ....ops.group_gemm.kernel import compute_lora_scaling from ....utils import logging from ..activations import ACT2FN +from .common import split_gate_up_proj logger = logging.get_logger(__name__) @@ -53,7 +54,9 @@ class MoEExpertsLoRA(LoraModule, nn.Module): and native (torch._grouped_mm). Base weights are frozen; only LoRA weights are trainable. - All weights in (G, K, N) format — ``[num_experts, in_features, out_features]``. + Base weights use fused ``gate_up_proj`` storage in (G, K, N) format — + ``[num_experts, in_features, out_features]``. ``gate_proj`` and + ``up_proj`` remain available as views for compatibility. When Expert Parallelism is enabled, pass ``num_local_experts`` to create weights at the local (sharded) shape. @@ -79,12 +82,8 @@ def __init__( self.lora_alpha = self.lora_config.lora_alpha # Base weights (frozen) in (G, K, N) format - self.gate_proj = nn.Parameter( - torch.empty(self.num_experts, hidden_dim, intermediate_size), - requires_grad=False, - ) - self.up_proj = nn.Parameter( - torch.empty(self.num_experts, hidden_dim, intermediate_size), + self.gate_up_proj = nn.Parameter( + torch.empty(self.num_experts, hidden_dim, 2 * intermediate_size), requires_grad=False, ) self.down_proj = nn.Parameter( @@ -126,6 +125,20 @@ def __init__( self.deepep_num_sms: int = 20 self.deepep_async_combine: bool = False + @property + def gate_proj(self) -> torch.Tensor: + gate_proj, _ = split_gate_up_proj(self.gate_up_proj, self.intermediate_size) + gate_proj.grad = ( + None if self.gate_up_proj.grad is None else self.gate_up_proj.grad[..., : self.intermediate_size] + ) + return gate_proj + + @property + def up_proj(self) -> torch.Tensor: + _, up_proj = split_gate_up_proj(self.gate_up_proj, self.intermediate_size) + up_proj.grad = None if self.gate_up_proj.grad is None else self.gate_up_proj.grad[..., self.intermediate_size :] + return up_proj + def _create_lora_params( self, name: str, A_experts: int, B_experts: int, r: int, in_features: int, out_features: int ): @@ -170,7 +183,7 @@ def merge_weights(self) -> None: for proj_name in ("gate_proj", "up_proj", "down_proj"): if proj_name not in self.lora_config.target_modules: continue - base = getattr(self, proj_name) # nn.Parameter [E, K, N] + base = getattr(self, proj_name) delta = self._compute_proj_delta(proj_name).to(base.dtype) base.add_(delta) self.reset_lora_parameters() @@ -203,13 +216,15 @@ def forward( from .backend import MOE_EXPERT_BACKENDS_LORA fn = MOE_EXPERT_BACKENDS_LORA[self.moe_implementation] + gate_proj = self.gate_proj.contiguous() + up_proj = self.up_proj.contiguous() return fn( num_experts=self.num_experts, routing_weights=routing_weights, selected_experts=selected_experts, hidden_states=hidden_states, - gate_proj=self.gate_proj, - up_proj=self.up_proj, + gate_proj=gate_proj, + up_proj=up_proj, down_proj=self.down_proj, gate_proj_lora_A=self.gate_proj_lora_A, gate_proj_lora_B=self.gate_proj_lora_B, @@ -247,6 +262,8 @@ def _ep_forward( dispatch_fn = EP_DISPATCH[self.ep_dispatch] combine_fn = EP_COMBINE[self.ep_dispatch] compute_fn = EP_EXPERT_COMPUTE_LORA[self.moe_implementation] + gate_proj = self.gate_proj.contiguous() + up_proj = self.up_proj.contiguous() # Step 1: Dispatch tokens to expert-owning ranks dispatch_kwargs = self._build_dispatch_kwargs(hidden_states, routing_weights, selected_experts, parallel_state) @@ -256,8 +273,8 @@ def _ep_forward( expert_output = compute_fn( permute_tokens, cumsum, - self.gate_proj, - self.up_proj, + gate_proj, + up_proj, self.down_proj, self.gate_proj_lora_A, self.gate_proj_lora_B, @@ -290,7 +307,7 @@ def _build_dispatch_kwargs(self, hidden_states, routing_weights, selected_expert buffer_size_gb=self.deepep_buffer_size_gb, num_sms=self.deepep_num_sms, ) - kwargs["num_local_experts"] = self.gate_proj.shape[0] + kwargs["num_local_experts"] = self.gate_up_proj.shape[0] return kwargs def _build_combine_kwargs(self, expert_output, ctx, dispatch_kwargs, parallel_state): @@ -358,7 +375,8 @@ def from_module(cls, module: nn.Module, r: int, lora_alpha: int, **kwargs): hybrid_shared=hybrid_shared, ) - num_exp = module.gate_proj.shape[0] + base_gate_up = getattr(module, "gate_up_proj", None) + num_exp = base_gate_up.shape[0] if base_gate_up is not None else module.gate_proj.shape[0] hidden_dim = module.hidden_dim intermediate_size = module.intermediate_size moe_implementation = getattr(module, "moe_implementation", "triton") @@ -378,13 +396,17 @@ def from_module(cls, module: nn.Module, r: int, lora_alpha: int, **kwargs): lora_experts.deepep_num_sms = getattr(module, "deepep_num_sms", 20) lora_experts.deepep_async_combine = getattr(module, "deepep_async_combine", False) + base_weight = base_gate_up if base_gate_up is not None else module.gate_proj lora_experts = lora_experts.to( - device=module.gate_proj.device, - dtype=module.gate_proj.dtype, + device=base_weight.device, + dtype=base_weight.dtype, ) with torch.no_grad(): - lora_experts.gate_proj.copy_(module.gate_proj) - lora_experts.up_proj.copy_(module.up_proj) + if base_gate_up is not None: + lora_experts.gate_up_proj.copy_(base_gate_up) + else: + lora_experts.gate_proj.copy_(module.gate_proj) + lora_experts.up_proj.copy_(module.up_proj) lora_experts.down_proj.copy_(module.down_proj) return lora_experts @@ -408,7 +430,8 @@ def inject_lora_into_experts( hybrid_shared=hybrid_shared, ) - num_local_experts = block.experts.gate_proj.shape[0] + gate_up_proj = getattr(block.experts, "gate_up_proj", None) + num_local_experts = gate_up_proj.shape[0] if gate_up_proj is not None else block.experts.gate_proj.shape[0] hidden_dim = block.experts.hidden_dim intermediate_size = block.experts.intermediate_size moe_implementation = getattr(block.experts, "moe_implementation", "triton") @@ -428,14 +451,18 @@ def inject_lora_into_experts( lora_experts.deepep_num_sms = getattr(block.experts, "deepep_num_sms", 20) lora_experts.deepep_async_combine = getattr(block.experts, "deepep_async_combine", False) + base_weight = gate_up_proj if gate_up_proj is not None else block.experts.gate_proj lora_experts = lora_experts.to( - device=block.experts.gate_proj.device, - dtype=block.experts.gate_proj.dtype, + device=base_weight.device, + dtype=base_weight.dtype, ) with torch.no_grad(): - lora_experts.gate_proj.copy_(block.experts.gate_proj) - lora_experts.up_proj.copy_(block.experts.up_proj) + if gate_up_proj is not None: + lora_experts.gate_up_proj.copy_(gate_up_proj) + else: + lora_experts.gate_proj.copy_(block.experts.gate_proj) + lora_experts.up_proj.copy_(block.experts.up_proj) lora_experts.down_proj.copy_(block.experts.down_proj) block.experts = lora_experts @@ -451,8 +478,11 @@ def inject_lora_into_experts( def copy_weights_to_lora_experts(source_experts: nn.Module, target_experts: nn.Module): """Copy base weights from source experts to LoRA experts.""" with torch.no_grad(): - target_experts.gate_proj.copy_(source_experts.gate_proj) - target_experts.up_proj.copy_(source_experts.up_proj) + if hasattr(source_experts, "gate_up_proj") and hasattr(target_experts, "gate_up_proj"): + target_experts.gate_up_proj.copy_(source_experts.gate_up_proj) + else: + target_experts.gate_proj.copy_(source_experts.gate_proj) + target_experts.up_proj.copy_(source_experts.up_proj) target_experts.down_proj.copy_(source_experts.down_proj) diff --git a/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py b/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py index b8307ce7..505202b6 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py @@ -1,7 +1,7 @@ """Checkpoint handler for Qwen3_5 MoE models.""" import re -from typing import Callable, List, Optional, Set, Tuple +from typing import Callable, Dict, List, Optional, Set, Tuple import torch @@ -89,29 +89,70 @@ def __init__( self._local_num_experts = num_experts // ep_size self._expert_start = ep_rank * self._local_num_experts self._expert_end = self._expert_start + self._local_num_experts + self._stacked_gate_up_pending: Dict[int, Dict[str, torch.Tensor]] = {} - _FUSED_EXPERT_GATE_UP_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.gate_up_proj(?:\.weight)?$") - _FUSED_EXPERT_DOWN_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.down_proj(?:\.weight)?$") + _STACKED_EXPERT_SPLIT_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.(gate|up|down)_proj$") + _HF_FUSED_EXPERT_GATE_UP_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.gate_up_proj\.weight$") + _HF_FUSED_EXPERT_DOWN_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.down_proj\.weight$") + _INTERNAL_FUSED_EXPERT_GATE_UP_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.gate_up_proj$") + _INTERNAL_FUSED_EXPERT_DOWN_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.down_proj$") def _slice_expert_tensor_for_ep(self, tensor: torch.Tensor) -> torch.Tensor: if self._ep_size == 1: return tensor return tensor[self._expert_start : self._expert_end].contiguous() + def _maybe_finalize_per_expert_merge( + self, + layer_idx: int, + proj: str, + ) -> List[Tuple[str, torch.Tensor]]: + if self._expert_buffer is None: + return [] + + if proj in {"gate", "up"}: + if not ( + self._expert_buffer.is_complete(layer_idx, "gate") and self._expert_buffer.is_complete(layer_idx, "up") + ): + return [] + gate = self._expert_buffer.pop_stacked(layer_idx, "gate") + up = self._expert_buffer.pop_stacked(layer_idx, "up") + return [ + ( + ExpertWeightBuffer.get_gate_up_name(layer_idx), + torch.cat([gate, up], dim=2), + ) + ] + + if proj == "down" and self._expert_buffer.is_complete(layer_idx, "down"): + return [ + ( + ExpertWeightBuffer.get_fused_name(layer_idx, "down"), + self._expert_buffer.pop_stacked(layer_idx, "down"), + ) + ] + + return [] + def _handle_fused_expert_weights(self, key: str, tensor: torch.Tensor) -> Optional[List[Tuple[str, torch.Tensor]]]: - gate_up_match = self._FUSED_EXPERT_GATE_UP_PATTERN.match(key) + internal_gate_up_match = self._INTERNAL_FUSED_EXPERT_GATE_UP_PATTERN.match(key) + if internal_gate_up_match is not None: + return [(key, self._slice_expert_tensor_for_ep(tensor))] + + internal_down_match = self._INTERNAL_FUSED_EXPERT_DOWN_PATTERN.match(key) + if internal_down_match is not None: + return [(key, self._slice_expert_tensor_for_ep(tensor))] + + gate_up_match = self._HF_FUSED_EXPERT_GATE_UP_PATTERN.match(key) if gate_up_match is not None: layer_idx = int(gate_up_match.group(1)) tensor = self._slice_expert_tensor_for_ep(tensor) - half = tensor.shape[1] // 2 - gate = tensor[:, :half, :].transpose(1, 2).contiguous() - up = tensor[:, half:, :].transpose(1, 2).contiguous() + gate_up = tensor.transpose(1, 2).contiguous() return [ - (f"model.layers.{layer_idx}.mlp.experts.gate_proj", gate), - (f"model.layers.{layer_idx}.mlp.experts.up_proj", up), + (f"model.layers.{layer_idx}.mlp.experts.gate_up_proj", gate_up), ] - down_match = self._FUSED_EXPERT_DOWN_PATTERN.match(key) + down_match = self._HF_FUSED_EXPERT_DOWN_PATTERN.match(key) if down_match is not None: layer_idx = int(down_match.group(1)) tensor = self._slice_expert_tensor_for_ep(tensor) @@ -120,6 +161,28 @@ def _handle_fused_expert_weights(self, key: str, tensor: torch.Tensor) -> Option (f"model.layers.{layer_idx}.mlp.experts.down_proj", down), ] + split_match = self._STACKED_EXPERT_SPLIT_PATTERN.match(key) + if split_match is not None: + layer_idx = int(split_match.group(1)) + proj = split_match.group(2) + tensor = self._slice_expert_tensor_for_ep(tensor) + if proj == "down": + return [(f"model.layers.{layer_idx}.mlp.experts.down_proj", tensor)] + + pending = self._stacked_gate_up_pending.setdefault(layer_idx, {}) + pending[proj] = tensor + if "gate" in pending and "up" in pending: + gate = pending.pop("gate") + up = pending.pop("up") + del self._stacked_gate_up_pending[layer_idx] + return [ + ( + f"model.layers.{layer_idx}.mlp.experts.gate_up_proj", + torch.cat([gate, up], dim=2), + ) + ] + return [] + return None def _handle_linear_attention_weights( @@ -231,11 +294,7 @@ def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torc if parsed is not None: layer_idx, expert_idx, proj = parsed self._expert_buffer.add(layer_idx, expert_idx, proj, tensor) - if self._expert_buffer.is_complete(layer_idx, proj): - fused_name = ExpertWeightBuffer.get_fused_name(layer_idx, proj) - stacked = self._expert_buffer.pop_stacked(layer_idx, proj) - return [(fused_name, stacked)] - return [] + return self._maybe_finalize_per_expert_merge(layer_idx, proj) # 4. Check QKV merge (skipped when unfused for TP) if self._qkv_buffer is not None: @@ -275,10 +334,7 @@ def on_skip_weight(self, key: str) -> List[Tuple[str, torch.Tensor]]: if parsed is not None: layer_idx, _expert_idx, proj = parsed self._expert_buffer.count_skipped(layer_idx, proj) - if self._expert_buffer.is_complete(layer_idx, proj): - fused_name = ExpertWeightBuffer.get_fused_name(layer_idx, proj) - stacked = self._expert_buffer.pop_stacked(layer_idx, proj) - return [(fused_name, stacked)] + return self._maybe_finalize_per_expert_merge(layer_idx, proj) return [] def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: @@ -288,6 +344,8 @@ def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: pending = self._expert_buffer.get_pending_counts() if pending: warnings.warn(f"Incomplete expert weights after loading: {pending}") + if self._stacked_gate_up_pending: + warnings.warn(f"Incomplete stacked expert gate/up merges after loading: {self._stacked_gate_up_pending}") if self._gate_up_buffer is not None: pending_gu = self._gate_up_buffer.get_pending() if pending_gu: @@ -299,6 +357,25 @@ def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: return [] def on_save_weight(self, param_name: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: + # Split fused MoE experts into per-expert HF weights. + if param_name.endswith(".mlp.experts.gate_up_proj"): + prefix = param_name.rsplit(".gate_up_proj", 1)[0] + half = tensor.shape[2] // 2 + gate = tensor[:, :, :half].transpose(1, 2).contiguous() + up = tensor[:, :, half:].transpose(1, 2).contiguous() + result = [] + for expert_idx in range(tensor.shape[0]): + result.append((f"{prefix}.{expert_idx}.gate_proj.weight", gate[expert_idx])) + result.append((f"{prefix}.{expert_idx}.up_proj.weight", up[expert_idx])) + return result + + if param_name.endswith(".mlp.experts.down_proj"): + prefix = param_name.rsplit(".down_proj", 1)[0] + down = tensor.transpose(1, 2).contiguous() + return [ + (f"{prefix}.{expert_idx}.down_proj.weight", down[expert_idx]) for expert_idx in range(tensor.shape[0]) + ] + # Split gate_up_proj -> gate_proj + up_proj (shared expert / dense layers) if ".gate_up_proj." in param_name: prefix, suffix = param_name.rsplit(".gate_up_proj.", 1) diff --git a/src/xorl/models/transformers/qwen3_5_moe/parallelize.py b/src/xorl/models/transformers/qwen3_5_moe/parallelize.py index 0864f6d3..b1e480c5 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/parallelize.py +++ b/src/xorl/models/transformers/qwen3_5_moe/parallelize.py @@ -54,9 +54,8 @@ def get_ep_plan(): LoRA weights are initialized at GLOBAL shape and sharded here. """ ep_plan = { - # Expert weights (stacked [num_experts, H, I] format) - "model.layers.*.mlp.experts.gate_proj": Shard(0), - "model.layers.*.mlp.experts.up_proj": Shard(0), + # Expert weights (stacked [num_experts, H, 2I] format) + "model.layers.*.mlp.experts.gate_up_proj": Shard(0), "model.layers.*.mlp.experts.down_proj": Shard(0), # LoRA weights for experts "model.layers.*.mlp.experts.gate_proj_lora_A": Shard(0), diff --git a/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py b/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py index 57ea4c6c..52ab41cd 100644 --- a/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py @@ -1,6 +1,7 @@ """Checkpoint handler for Qwen3 MoE models.""" -from typing import Callable, List, Optional, Set, Tuple +import re +from typing import Callable, Dict, List, Optional, Set, Tuple import torch import torch.nn as nn @@ -50,6 +51,10 @@ class Qwen3MoeCheckpointHandler(CheckpointHandler): - Expert merging is always active (stacking per-expert HF weights). """ + _STACKED_EXPERT_SPLIT_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.(gate|up|down)_proj$") + _FUSED_EXPERT_GATE_UP_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.gate_up_proj$") + _FUSED_EXPERT_DOWN_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.down_proj$") + def __init__( self, num_experts: int, @@ -81,8 +86,14 @@ def __init__( self._gate_up_buffer: Optional[GateUpMergeBuffer] = None if not skip_gate_up_merge: self._gate_up_buffer = GateUpMergeBuffer() + self._stacked_gate_up_pending: Dict[int, Dict[str, torch.Tensor]] = {} self._q_dim = num_attention_heads * head_dim self._kv_dim = num_key_value_heads * head_dim + self._ep_rank = ep_rank + self._ep_size = ep_size + self._local_num_experts = num_experts // ep_size + self._expert_start = ep_rank * self._local_num_experts + self._expert_end = self._expert_start + self._local_num_experts self._is_prequantized = is_prequantized self._exclude_modules = exclude_modules or set() # Inline QLoRA weight loading for dense modules (attention, shared expert) @@ -99,6 +110,81 @@ def __init__( num_experts=num_experts, ) + def _slice_expert_tensor_for_ep(self, tensor: torch.Tensor) -> torch.Tensor: + if self._ep_size == 1: + return tensor + return tensor[self._expert_start : self._expert_end].contiguous() + + def _maybe_finalize_per_expert_merge( + self, + layer_idx: int, + proj: str, + ) -> List[Tuple[str, torch.Tensor]]: + if self._expert_buffer is None: + return [] + + if proj in {"gate", "up"}: + if not ( + self._expert_buffer.is_complete(layer_idx, "gate") and self._expert_buffer.is_complete(layer_idx, "up") + ): + return [] + gate = self._expert_buffer.pop_stacked(layer_idx, "gate") + up = self._expert_buffer.pop_stacked(layer_idx, "up") + return [ + ( + ExpertWeightBuffer.get_gate_up_name(layer_idx), + torch.cat([gate, up], dim=2), + ) + ] + + if proj == "down" and self._expert_buffer.is_complete(layer_idx, "down"): + return [ + ( + ExpertWeightBuffer.get_fused_name(layer_idx, "down"), + self._expert_buffer.pop_stacked(layer_idx, "down"), + ) + ] + + return [] + + def _handle_stacked_expert_weights( + self, + key: str, + tensor: torch.Tensor, + ) -> Optional[List[Tuple[str, torch.Tensor]]]: + gate_up_match = self._FUSED_EXPERT_GATE_UP_PATTERN.match(key) + if gate_up_match is not None: + return [(key, self._slice_expert_tensor_for_ep(tensor))] + + down_match = self._FUSED_EXPERT_DOWN_PATTERN.match(key) + if down_match is not None: + return [(key, self._slice_expert_tensor_for_ep(tensor))] + + split_match = self._STACKED_EXPERT_SPLIT_PATTERN.match(key) + if split_match is None: + return None + + layer_idx = int(split_match.group(1)) + proj = split_match.group(2) + tensor = self._slice_expert_tensor_for_ep(tensor) + + if proj == "down": + return [(f"model.layers.{layer_idx}.mlp.experts.down_proj", tensor)] + + pending = self._stacked_gate_up_pending.setdefault(layer_idx, {}) + pending[proj] = tensor + if "gate" in pending and "up" in pending: + gate = pending.pop("gate") + up = pending.pop("up") + del self._stacked_gate_up_pending[layer_idx] + return [ + ( + f"model.layers.{layer_idx}.mlp.experts.gate_up_proj", + torch.cat([gate, up], dim=2), + ) + ] + return [] + def get_skip_key_fn(self) -> Optional[Callable[[str], bool]]: """Return predicate to skip keys during loading. @@ -220,19 +306,20 @@ def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torc if OPROJ_WEIGHT_PATTERN.match(key) or DENSE_DOWN_PROJ_PATTERN.match(key): return [] - # 1. Check expert merge (disabled when prequantized — expert_buffer is None) + # 1. Handle internal stacked/fused expert tensors. + stacked_expert_results = self._handle_stacked_expert_weights(key, tensor) + if stacked_expert_results is not None: + return stacked_expert_results + + # 2. Check expert merge (disabled when prequantized — expert_buffer is None) if self._expert_buffer is not None: parsed = parse_expert_key(key) if parsed is not None: layer_idx, expert_idx, proj = parsed self._expert_buffer.add(layer_idx, expert_idx, proj, tensor) - if self._expert_buffer.is_complete(layer_idx, proj): - fused_name = ExpertWeightBuffer.get_fused_name(layer_idx, proj) - stacked = self._expert_buffer.pop_stacked(layer_idx, proj) - return [(fused_name, stacked)] - return [] + return self._maybe_finalize_per_expert_merge(layer_idx, proj) - # 2. Check QKV merge (skipped when unfused for TP) + # 3. Check QKV merge (skipped when unfused for TP) if self._qkv_buffer is not None: # When pre-quantized, QKV .weight keys are uint8 packed — # skip standard merging, they'll be loaded by QLoRALinear @@ -246,7 +333,7 @@ def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torc if self._qkv_buffer.is_qkv_key(key): return [] - # 3. Check gate/up merge (skipped when unfused for TP) + # 4. Check gate/up merge (skipped when unfused for TP) if self._gate_up_buffer is not None: # When pre-quantized, gate/up .weight keys are uint8 packed — # skip standard merging, they'll be loaded by QLoRALinear @@ -260,7 +347,7 @@ def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torc if self._gate_up_buffer.is_gate_up_key(key): return [] - # 4. Passthrough + # 5. Passthrough return [(key, tensor)] def on_skip_weight(self, key: str) -> List[Tuple[str, torch.Tensor]]: @@ -274,10 +361,7 @@ def on_skip_weight(self, key: str) -> List[Tuple[str, torch.Tensor]]: if parsed is not None: layer_idx, _expert_idx, proj = parsed self._expert_buffer.count_skipped(layer_idx, proj) - if self._expert_buffer.is_complete(layer_idx, proj): - fused_name = ExpertWeightBuffer.get_fused_name(layer_idx, proj) - stacked = self._expert_buffer.pop_stacked(layer_idx, proj) - return [(fused_name, stacked)] + return self._maybe_finalize_per_expert_merge(layer_idx, proj) return [] def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: @@ -287,6 +371,8 @@ def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: pending = self._expert_buffer.get_pending_counts() if pending: warnings.warn(f"Incomplete expert weights after loading: {pending}") + if self._stacked_gate_up_pending: + warnings.warn(f"Incomplete stacked expert gate/up merges after loading: {self._stacked_gate_up_pending}") if self._gate_up_buffer is not None: pending_gu = self._gate_up_buffer.get_pending() if pending_gu: @@ -308,6 +394,25 @@ def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: return [] def on_save_weight(self, param_name: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: + # Split fused MoE experts into per-expert HF weights. + if param_name.endswith(".mlp.experts.gate_up_proj"): + prefix = param_name.rsplit(".gate_up_proj", 1)[0] + half = tensor.shape[2] // 2 + gate = tensor[:, :, :half].transpose(1, 2).contiguous() + up = tensor[:, :, half:].transpose(1, 2).contiguous() + result = [] + for expert_idx in range(tensor.shape[0]): + result.append((f"{prefix}.{expert_idx}.gate_proj.weight", gate[expert_idx])) + result.append((f"{prefix}.{expert_idx}.up_proj.weight", up[expert_idx])) + return result + + if param_name.endswith(".mlp.experts.down_proj"): + prefix = param_name.rsplit(".down_proj", 1)[0] + down = tensor.transpose(1, 2).contiguous() + return [ + (f"{prefix}.{expert_idx}.down_proj.weight", down[expert_idx]) for expert_idx in range(tensor.shape[0]) + ] + # Split gate_up_proj -> gate_proj + up_proj (shared expert / dense layers) if ".gate_up_proj." in param_name: prefix, suffix = param_name.rsplit(".gate_up_proj.", 1) diff --git a/src/xorl/models/transformers/qwen3_moe/parallelize.py b/src/xorl/models/transformers/qwen3_moe/parallelize.py index 975295ca..1b82d9a7 100644 --- a/src/xorl/models/transformers/qwen3_moe/parallelize.py +++ b/src/xorl/models/transformers/qwen3_moe/parallelize.py @@ -53,9 +53,8 @@ def get_ep_plan(): LoRA weights are initialized at GLOBAL shape and sharded here. """ ep_plan = { - # Expert weights (stacked [num_experts, H, I] format) - "model.layers.*.mlp.experts.gate_proj": Shard(0), - "model.layers.*.mlp.experts.up_proj": Shard(0), + # Expert weights (stacked [num_experts, H, 2I] format) + "model.layers.*.mlp.experts.gate_up_proj": Shard(0), "model.layers.*.mlp.experts.down_proj": Shard(0), # LoRA weights for experts "model.layers.*.mlp.experts.gate_proj_lora_A": Shard(0), diff --git a/src/xorl/qlora/modules/moe_experts.py b/src/xorl/qlora/modules/moe_experts.py index 979470b3..e9cd6b28 100644 --- a/src/xorl/qlora/modules/moe_experts.py +++ b/src/xorl/qlora/modules/moe_experts.py @@ -420,6 +420,11 @@ def gate_proj(self) -> Tensor: """Dequantize all local gate_proj weights. [num_local_experts, hidden, intermediate]""" return self.dequantize_all_experts("gate", self.hidden_size, self.intermediate_size) + @property + def gate_up_proj(self) -> Tensor: + """Dequantize all local fused gate_up_proj weights. [num_local_experts, hidden, 2 * intermediate]""" + return torch.cat([self.gate_proj, self.up_proj], dim=2) + @property def up_proj(self) -> Tensor: """Dequantize all local up_proj weights. [num_local_experts, hidden, intermediate]""" diff --git a/src/xorl/qlora/utils.py b/src/xorl/qlora/utils.py index 3be7331a..aa8778a7 100644 --- a/src/xorl/qlora/utils.py +++ b/src/xorl/qlora/utils.py @@ -226,20 +226,30 @@ def inject_qlora_into_model( if not (hasattr(module, "gate") and hasattr(module, "experts")): continue experts = module.experts + fused_gate_up = getattr(experts, "gate_up_proj", None) has_3d_weights = ( - hasattr(experts, "gate_proj") - and isinstance(experts.gate_proj, nn.Parameter) - and hasattr(experts, "up_proj") - and isinstance(experts.up_proj, nn.Parameter) + isinstance(fused_gate_up, nn.Parameter) and hasattr(experts, "down_proj") and isinstance(experts.down_proj, nn.Parameter) - and experts.gate_proj.dim() == 3 + and fused_gate_up.dim() == 3 and not isinstance(experts, QLoRAMoeExperts) ) + if not has_3d_weights: + has_3d_weights = ( + hasattr(experts, "gate_proj") + and isinstance(experts.gate_proj, nn.Parameter) + and hasattr(experts, "up_proj") + and isinstance(experts.up_proj, nn.Parameter) + and hasattr(experts, "down_proj") + and isinstance(experts.down_proj, nn.Parameter) + and experts.gate_proj.dim() == 3 + and not isinstance(experts, QLoRAMoeExperts) + ) if not has_3d_weights: continue - num_experts = experts.gate_proj.shape[0] + base_gate = fused_gate_up if fused_gate_up is not None else experts.gate_proj + num_experts = base_gate.shape[0] num_local_experts = num_experts // ep_size expert_offset = ep_rank * num_local_experts @@ -257,7 +267,7 @@ def inject_qlora_into_model( experts_fqn = name + ".experts" if not name.endswith(".experts") else name if quant_format == "nf4": - is_meta = experts.gate_proj.device.type == "meta" + is_meta = base_gate.device.type == "meta" if is_meta: # Defer: load bf16 from checkpoint after FSDP, then quantize qlora_experts._source_fqn = experts_fqn diff --git a/src/xorl/server/weight_sync/backends/base.py b/src/xorl/server/weight_sync/backends/base.py index acddba7e..f348f600 100644 --- a/src/xorl/server/weight_sync/backends/base.py +++ b/src/xorl/server/weight_sync/backends/base.py @@ -125,6 +125,7 @@ def transfer_bucket( *, src_rank: int = 0, flush_cache: bool = False, + weight_version: Optional[str] = None, ) -> None: """Send a bucket of named tensors to inference. @@ -136,6 +137,9 @@ def transfer_bucket( flush_cache: If ``True``, tell the inference endpoint to flush its KV cache after loading this bucket (used for the final bucket of a sync). + weight_version: Optional version marker to apply on the inference + endpoint after this bucket is loaded. Handlers should only set + this on the final bucket of a sync. """ # ------------------------------------------------------------------ diff --git a/src/xorl/server/weight_sync/backends/nccl_broadcast.py b/src/xorl/server/weight_sync/backends/nccl_broadcast.py index 613315e3..f64471ae 100644 --- a/src/xorl/server/weight_sync/backends/nccl_broadcast.py +++ b/src/xorl/server/weight_sync/backends/nccl_broadcast.py @@ -558,6 +558,7 @@ def _transfer_single_bucket( self, bucket: List[Tuple[str, torch.Tensor]], flush_cache: bool = False, + weight_version: Optional[str] = None, ) -> List[Dict[str, Any]]: """ Transfer a single bucket of parameters. @@ -571,6 +572,7 @@ def _transfer_single_bucket( Args: bucket: List of (name, tensor) tuples flush_cache: Whether to flush cache after this bucket (only for last bucket) + weight_version: Optional weight version to apply with this bucket. Returns: List of results from each endpoint @@ -597,6 +599,7 @@ def call_single_endpoint(endpoint: EndpointInfo, endpoint_idx: int): "shapes": shapes, "group_name": self.group_name, "flush_cache": flush_cache, + "weight_version": weight_version, }, timeout=600, ) @@ -853,12 +856,17 @@ def transfer_bucket( *, src_rank: int = 0, flush_cache: bool = False, + weight_version: Optional[str] = None, ) -> None: if src_rank != 0: raise ValueError(f"NCCLBroadcastBackend only supports src_rank=0, got {src_rank}") if self._synchronizer is None: raise RuntimeError("Backend not initialized — call initialize() first") - self._synchronizer._transfer_single_bucket(bucket, flush_cache=flush_cache) + self._synchronizer._transfer_single_bucket( + bucket, + flush_cache=flush_cache, + weight_version=weight_version, + ) @property def sender_ranks(self) -> FrozenSet[int]: diff --git a/src/xorl/server/weight_sync/handler.py b/src/xorl/server/weight_sync/handler.py index 6990a9e8..071f68b5 100644 --- a/src/xorl/server/weight_sync/handler.py +++ b/src/xorl/server/weight_sync/handler.py @@ -83,12 +83,13 @@ async def handle_sync_inference_weights(self, command_dict: Dict[str, Any]) -> D sync_method = p.sync_method flush_cache = p.flush_cache pause_mode = p.pause_mode + weight_version = p.weight_version quantization = p.quantization logger.info( f"Rank {self.rank}: [WeightSync] sync_method={sync_method}, " f"endpoints={len(endpoints)}, flush_cache={flush_cache}, " - f"quantization={quantization}" + f"weight_version={weight_version}, quantization={quantization}" ) try: @@ -101,6 +102,7 @@ async def handle_sync_inference_weights(self, command_dict: Dict[str, Any]) -> D sync_method=sync_method, flush_cache=flush_cache, pause_mode=pause_mode, + weight_version=weight_version, quantization=quantization, ) @@ -124,6 +126,7 @@ def _sync_weights( sync_method: str, flush_cache: bool, pause_mode: str, + weight_version: Optional[str], quantization: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """ @@ -356,6 +359,7 @@ def _sync_weights( backend, current_buffer, flush_cache=(flush_cache and is_last_overall and not moe_contexts), + weight_version=weight_version if is_last_overall and not moe_contexts else None, ) total_bytes += b total_params += p @@ -370,6 +374,7 @@ def _sync_weights( backend, ctx, flush_cache=(flush_cache and is_last_overall), + weight_version=weight_version if is_last_overall else None, quantization=quantization, ps=_ps, ) @@ -382,6 +387,7 @@ def _sync_weights( backend, ctx, flush_cache=(flush_cache and is_last_overall), + weight_version=weight_version if is_last_overall else None, quantization=quantization, ) total_bytes += b @@ -423,6 +429,7 @@ def _sync_weights( backend, received, flush_cache=(flush_cache and is_last_overall), + weight_version=weight_version if is_last_overall else None, ) total_bytes += b total_params += p @@ -648,16 +655,19 @@ def _collect_ep_moe_data( continue # Get expert params — after unshard they may be plain tensors or DTensors - gate = getattr(mod, "gate_proj", None) - if gate is None or not isinstance(gate, torch.nn.Parameter): - continue + gate_up = getattr(mod, "gate_up_proj", None) + if isinstance(gate_up, torch.nn.Parameter): + E_local = gate_up.shape[0] + else: + gate = getattr(mod, "gate_proj", None) + if gate is None or not isinstance(gate, torch.nn.Parameter): + continue + E_local = gate.shape[0] if mname: full_prefix = f"{mod_name}.{mname}" if mod_name != "(root)" else mname else: full_prefix = mod_name - # gate_proj has shape [num_local_experts, K, N] — already local - E_local = gate.shape[0] # Clone local expert data for each projection. # With EP, each rank's module already holds only local experts [E_local, K, N]. @@ -675,8 +685,6 @@ def _collect_ep_moe_data( delta = mod._compute_proj_delta(proj_name) if isinstance(delta, DTensor): delta = delta.to_local() - elif delta.shape[0] > E_local: - delta = delta[start : start + E_local] local = local.to(torch.bfloat16) + delta.to(torch.bfloat16) else: local = local.to(torch.bfloat16) @@ -705,6 +713,7 @@ def _gather_and_broadcast_ep_moe_experts( backend, ctx: Dict[str, Any], flush_cache: bool = False, + weight_version: Optional[str] = None, bucket_size_bytes: int = _DEFAULT_MOE_BUCKET_BYTES, quantization: Optional[Dict[str, Any]] = None, ps=None, @@ -731,6 +740,20 @@ def _gather_and_broadcast_ep_moe_experts( device = f"cuda:{self.rank % torch.cuda.device_count()}" is_qlora = ctx.get("type") != "full_weight" + ep_fsdp_rank = 0 + if ps.ep_fsdp_device_mesh is not None: + ep_fsdp_rank = ps.ep_fsdp_device_mesh.get_local_rank("ep_fsdp") + + # EP-FSDP replicas hold identical local expert shards after unshard(). + # Only one replica column needs to gather and forward those experts to + # rank 0; the others would just duplicate the same weights. + if ep_fsdp_rank != 0: + if is_qlora: + ctx["lora_params"] = None + else: + ctx["local_experts"] = None + return 0, 0, 0 + if self.rank == 0: logger.info( f"Rank 0: [EP-Gather] prefix={full_prefix}, type={'qlora' if is_qlora else 'full_weight'}, " @@ -841,6 +864,7 @@ def _gather_and_broadcast_ep_moe_experts( backend, bucket, flush_cache=flush_cache, + weight_version=weight_version, ) total_bytes += b total_params += p @@ -863,6 +887,7 @@ def _broadcast_moe_experts_bucketed( backend, ctx: Dict[str, Any], flush_cache: bool = False, + weight_version: Optional[str] = None, bucket_size_bytes: int = _DEFAULT_MOE_BUCKET_BYTES, quantization: Optional[Dict[str, Any]] = None, ) -> Tuple[int, int, int]: @@ -957,6 +982,7 @@ def _broadcast_moe_experts_bucketed( backend, bucket, flush_cache=flush_cache, + weight_version=weight_version, ) total_bytes += b total_params += p @@ -1286,25 +1312,25 @@ def _extract_params_for_sync( buffer.append((full_name, cloned)) continue elif isinstance(lora_mod, MoEExpertsLoRA): - # MoE experts: fused [E, K, N] → per-expert HF format - # param_leaf is "gate_proj", "up_proj", or "down_proj" - if param_leaf in lora_mod.lora_config.target_modules: - delta = lora_mod._compute_proj_delta(param_leaf) - merged = param.data.to(dtype=torch.bfloat16) + delta.to(dtype=torch.bfloat16) - else: - merged = param.data.to(dtype=torch.bfloat16) - # Split fused [E, K, N] into per-expert [N, K] (HF format) - # full_name is like "model.layers.0.mlp.experts.gate_proj" - # HF expects "model.layers.0.mlp.experts.{i}.gate_proj.weight" - base = full_name # e.g. "model.layers.0.mlp.experts.gate_proj" - # Split "experts.gate_proj" → "experts.{i}.gate_proj.weight" - parts = base.rsplit(".", 1) # ["model.layers.0.mlp.experts", "gate_proj"] - for expert_idx in range(merged.shape[0]): - hf_name = f"{parts[0]}.{expert_idx}.{parts[1]}.weight" - # GKN [K, N] → HF [N, K] (transpose) - expert_weight = merged[expert_idx].t().contiguous() - buffer.append((hf_name, expert_weight)) - continue + if param_leaf == "gate_up_proj": + merged = param.data.to(dtype=torch.bfloat16).clone() + half = merged.shape[2] // 2 + if "gate_proj" in lora_mod.lora_config.target_modules: + gate_delta = lora_mod._compute_proj_delta("gate_proj") + merged[:, :, :half].add_(gate_delta.to(dtype=torch.bfloat16)) + if "up_proj" in lora_mod.lora_config.target_modules: + up_delta = lora_mod._compute_proj_delta("up_proj") + merged[:, :, half:].add_(up_delta.to(dtype=torch.bfloat16)) + buffer.append((full_name, merged)) + continue + if param_leaf == "down_proj": + if "down_proj" in lora_mod.lora_config.target_modules: + delta = lora_mod._compute_proj_delta("down_proj") + merged = param.data.to(dtype=torch.bfloat16) + delta.to(dtype=torch.bfloat16) + else: + merged = param.data.to(dtype=torch.bfloat16) + buffer.append((full_name, merged.clone())) + continue # Check if this is a non-LoRA MoEExperts fused tensor _is_moe_experts = False if parent_name: @@ -1315,19 +1341,13 @@ def _extract_params_for_sync( for p in parts_list: parent_mod = getattr(parent_mod, p) if isinstance(parent_mod, (MoEExperts, MoEExpertsLoRA)): - if param_leaf in ("gate_proj", "up_proj", "down_proj"): + if param_leaf in ("gate_up_proj", "down_proj"): _is_moe_experts = True except (AttributeError, TypeError): pass if _is_moe_experts: - # Split fused [E, K, N] into per-expert HF [N, K] - data = param.data.to(dtype=torch.bfloat16) - parts = full_name.rsplit(".", 1) - for expert_idx in range(data.shape[0]): - hf_name = f"{parts[0]}.{expert_idx}.{parts[1]}.weight" - expert_weight = data[expert_idx].t().contiguous() - buffer.append((hf_name, expert_weight)) + buffer.append((full_name, param.data.to(dtype=torch.bfloat16).clone())) else: cloned = param.data.to(dtype=torch.bfloat16).clone() buffer.append((full_name, cloned)) @@ -1359,10 +1379,17 @@ def _unfuse_for_inference( ) -> List[Tuple[str, torch.Tensor]]: """Split fused projections (qkv_proj, gate_up_proj) into HF-format names. - SGLang's load_weights expects HF names (q_proj, k_proj, v_proj, gate_proj, - up_proj) and has stacked_params_mapping that does string replacement. - If we send fused names like 'qkv_proj', the replacement mangles them - (e.g. 'v_proj' is a substring of 'qkv_proj' → 'qkqkv_proj'). + Handles: + - qkv_proj → q_proj + k_proj + v_proj (split fused attention) + - gate_up_proj → gate_proj + up_proj (split fused dense/shared MLP) + - MoE experts: gate_up_proj/down_proj → per-expert HF gate/up/down weights + - Qwen3.5 linear attention: remap split GatedDeltaNet params back to + HF fused names (q_proj/k_proj/v_proj → in_proj_qkv, etc.) + """ + from xorl.models.transformers.qwen3_5_shared import ( + has_linear_attention_layers, + remap_linear_attention_params_for_inference, + ) This splits fused weights back to individual projections before sending. """ @@ -1384,6 +1411,19 @@ def _unfuse_for_inference( result.append((f"{prefix}.q_proj.{suffix}", q)) result.append((f"{prefix}.k_proj.{suffix}", k)) result.append((f"{prefix}.v_proj.{suffix}", v)) + elif name.endswith(".mlp.experts.gate_up_proj"): + prefix = name.rsplit(".gate_up_proj", 1)[0] + half = tensor.shape[2] // 2 + gate = tensor[:, :, :half].transpose(1, 2).contiguous() + up = tensor[:, :, half:].transpose(1, 2).contiguous() + for expert_idx in range(tensor.shape[0]): + result.append((f"{prefix}.{expert_idx}.gate_proj.weight", gate[expert_idx])) + result.append((f"{prefix}.{expert_idx}.up_proj.weight", up[expert_idx])) + elif name.endswith(".mlp.experts.down_proj"): + prefix = name.rsplit(".down_proj", 1)[0] + down = tensor.transpose(1, 2).contiguous() + for expert_idx in range(tensor.shape[0]): + result.append((f"{prefix}.{expert_idx}.down_proj.weight", down[expert_idx])) elif ".gate_up_proj." in name: # Split [2*intermediate, hidden] → gate, up prefix, suffix = name.rsplit(".gate_up_proj.", 1) @@ -1401,6 +1441,7 @@ def _broadcast_buffer( backend, buffer: List[Tuple[str, torch.Tensor]], flush_cache: bool, + weight_version: Optional[str] = None, ) -> Tuple[int, int]: """ Broadcast a buffer of (name, tensor) pairs to inference endpoints. @@ -1417,7 +1458,11 @@ def _broadcast_buffer( bucket_bytes = sum(t.numel() * t.element_size() for _, t in buffer) logger.info(f"Rank {self.rank}: [WeightSync] Broadcasting {len(buffer)} params, {bucket_bytes / 1e6:.1f} MB") - backend.transfer_bucket(buffer, flush_cache=flush_cache) + backend.transfer_bucket( + buffer, + flush_cache=flush_cache, + weight_version=weight_version, + ) return bucket_bytes, len(buffer) @staticmethod diff --git a/tests/models/test_moe_fused_gate_up_proj.py b/tests/models/test_moe_fused_gate_up_proj.py new file mode 100644 index 00000000..9e0bd0a2 --- /dev/null +++ b/tests/models/test_moe_fused_gate_up_proj.py @@ -0,0 +1,134 @@ +import pytest +import torch + +from xorl.models.layers.moe import MoEExperts, MoEExpertsLoRA, MoELoRAConfig +from xorl.models.transformers.qwen3_5_moe.checkpoint_handler import Qwen3_5MoeCheckpointHandler +from xorl.models.transformers.qwen3_moe.checkpoint_handler import Qwen3MoeCheckpointHandler + + +pytestmark = [pytest.mark.cpu] + + +def test_moe_experts_register_only_fused_gate_up_proj(): + experts = MoEExperts(num_experts=3, hidden_dim=4, intermediate_size=5, moe_implementation="eager") + + named_params = dict(experts.named_parameters()) + assert set(named_params) == {"gate_up_proj", "down_proj"} + assert named_params["gate_up_proj"].shape == (3, 4, 10) + assert experts.gate_proj.shape == (3, 4, 5) + assert experts.up_proj.shape == (3, 4, 5) + + with torch.no_grad(): + experts.gate_up_proj.zero_() + experts.gate_proj.fill_(1.0) + experts.up_proj.fill_(2.0) + + torch.testing.assert_close(experts.gate_up_proj[..., :5], torch.ones(3, 4, 5)) + torch.testing.assert_close(experts.gate_up_proj[..., 5:], torch.full((3, 4, 5), 2.0)) + + +def test_moe_experts_lora_registers_fused_base_weight(): + experts = MoEExpertsLoRA( + num_experts=3, + hidden_dim=4, + intermediate_size=5, + moe_implementation="eager", + lora_config=MoELoRAConfig(r=2, lora_alpha=4), + ) + + named_params = dict(experts.named_parameters()) + assert "gate_up_proj" in named_params + assert "gate_proj" not in named_params + assert "up_proj" not in named_params + assert named_params["gate_up_proj"].shape == (3, 4, 10) + assert experts.gate_proj.shape == (3, 4, 5) + assert experts.up_proj.shape == (3, 4, 5) + + +def test_qwen3_moe_checkpoint_handler_round_trips_fused_experts(): + hidden_size = 4 + intermediate_size = 3 + handler = Qwen3MoeCheckpointHandler( + num_experts=2, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=2, + ) + + gate_0 = torch.arange(0, 12, dtype=torch.float32).view(intermediate_size, hidden_size) + gate_1 = torch.arange(12, 24, dtype=torch.float32).view(intermediate_size, hidden_size) + up_0 = torch.arange(24, 36, dtype=torch.float32).view(intermediate_size, hidden_size) + up_1 = torch.arange(36, 48, dtype=torch.float32).view(intermediate_size, hidden_size) + down_0 = torch.arange(48, 60, dtype=torch.float32).view(hidden_size, intermediate_size) + down_1 = torch.arange(60, 72, dtype=torch.float32).view(hidden_size, intermediate_size) + + results = [] + for key, tensor in [ + ("model.layers.0.mlp.experts.0.gate_proj.weight", gate_0), + ("model.layers.0.mlp.experts.1.gate_proj.weight", gate_1), + ("model.layers.0.mlp.experts.0.up_proj.weight", up_0), + ("model.layers.0.mlp.experts.1.up_proj.weight", up_1), + ("model.layers.0.mlp.experts.0.down_proj.weight", down_0), + ("model.layers.0.mlp.experts.1.down_proj.weight", down_1), + ]: + results.extend(handler.on_load_weight(key, tensor)) + + loaded = dict(results) + expected_gate = torch.stack([gate_0.t(), gate_1.t()], dim=0) + expected_up = torch.stack([up_0.t(), up_1.t()], dim=0) + expected_gate_up = torch.cat([expected_gate, expected_up], dim=2) + expected_down = torch.stack([down_0.t(), down_1.t()], dim=0) + + assert set(loaded) == { + "model.layers.0.mlp.experts.gate_up_proj", + "model.layers.0.mlp.experts.down_proj", + } + torch.testing.assert_close(loaded["model.layers.0.mlp.experts.gate_up_proj"], expected_gate_up) + torch.testing.assert_close(loaded["model.layers.0.mlp.experts.down_proj"], expected_down) + + saved = dict(handler.on_save_weight("model.layers.0.mlp.experts.gate_up_proj", expected_gate_up)) + saved.update(handler.on_save_weight("model.layers.0.mlp.experts.down_proj", expected_down)) + + torch.testing.assert_close(saved["model.layers.0.mlp.experts.0.gate_proj.weight"], gate_0) + torch.testing.assert_close(saved["model.layers.0.mlp.experts.1.gate_proj.weight"], gate_1) + torch.testing.assert_close(saved["model.layers.0.mlp.experts.0.up_proj.weight"], up_0) + torch.testing.assert_close(saved["model.layers.0.mlp.experts.1.up_proj.weight"], up_1) + torch.testing.assert_close(saved["model.layers.0.mlp.experts.0.down_proj.weight"], down_0) + torch.testing.assert_close(saved["model.layers.0.mlp.experts.1.down_proj.weight"], down_1) + + +def test_qwen3_5_moe_checkpoint_handler_round_trips_fused_experts(): + hidden_size = 4 + intermediate_size = 3 + gate_up_weight = torch.arange(0, 48, dtype=torch.float32).view(2, 2 * intermediate_size, hidden_size) + down_weight = torch.arange(48, 72, dtype=torch.float32).view(2, hidden_size, intermediate_size) + + handler = Qwen3_5MoeCheckpointHandler( + num_experts=2, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=2, + linear_key_dim=2, + linear_value_dim=2, + ) + + loaded = dict( + handler.on_load_weight("model.layers.0.mlp.experts.gate_up_proj.weight", gate_up_weight) + + handler.on_load_weight("model.layers.0.mlp.experts.down_proj.weight", down_weight) + ) + + expected_gate_up = gate_up_weight.transpose(1, 2).contiguous() + expected_down = down_weight.transpose(1, 2).contiguous() + torch.testing.assert_close(loaded["model.layers.0.mlp.experts.gate_up_proj"], expected_gate_up) + torch.testing.assert_close(loaded["model.layers.0.mlp.experts.down_proj"], expected_down) + + saved = dict(handler.on_save_weight("model.layers.0.mlp.experts.gate_up_proj", expected_gate_up)) + saved.update(handler.on_save_weight("model.layers.0.mlp.experts.down_proj", expected_down)) + + gate, up = gate_up_weight.split(intermediate_size, dim=1) + torch.testing.assert_close(saved["model.layers.0.mlp.experts.0.gate_proj.weight"], gate[0]) + torch.testing.assert_close(saved["model.layers.0.mlp.experts.1.gate_proj.weight"], gate[1]) + torch.testing.assert_close(saved["model.layers.0.mlp.experts.0.up_proj.weight"], up[0]) + torch.testing.assert_close(saved["model.layers.0.mlp.experts.1.up_proj.weight"], up[1]) + torch.testing.assert_close(saved["model.layers.0.mlp.experts.0.down_proj.weight"], down_weight[0]) + torch.testing.assert_close(saved["model.layers.0.mlp.experts.1.down_proj.weight"], down_weight[1]) diff --git a/tests/server/weight_sync/test_weight_version_forwarding.py b/tests/server/weight_sync/test_weight_version_forwarding.py new file mode 100644 index 00000000..1287056b --- /dev/null +++ b/tests/server/weight_sync/test_weight_version_forwarding.py @@ -0,0 +1,52 @@ +"""Regression tests for weight_version forwarding in weight sync.""" + +from unittest.mock import MagicMock + +import torch + +from xorl.server.weight_sync.backends.base import EndpointConfig, TransportConfig +from xorl.server.weight_sync.backends.nccl_broadcast import NCCLBroadcastBackend +from xorl.server.weight_sync.handler import WeightSyncHandler + + +class TestWeightVersionForwarding: + def test_handler_broadcast_buffer_forwards_weight_version(self): + handler = MagicMock() + handler.rank = 0 + handler._broadcast_buffer = WeightSyncHandler._broadcast_buffer.__get__(handler) + + backend = MagicMock() + buffer = [("layer.weight", torch.ones(2, 3, dtype=torch.bfloat16))] + + bucket_bytes, num_params = handler._broadcast_buffer( + backend, + buffer, + flush_cache=False, + weight_version="sync-v1", + ) + + assert bucket_bytes == buffer[0][1].numel() * buffer[0][1].element_size() + assert num_params == 1 + backend.transfer_bucket.assert_called_once_with( + buffer, + flush_cache=False, + weight_version="sync-v1", + ) + + def test_nccl_backend_transfer_bucket_forwards_weight_version(self): + backend = NCCLBroadcastBackend(TransportConfig(endpoints=[EndpointConfig(host="127.0.0.1", port=30000)])) + backend._synchronizer = MagicMock() + + buffer = [("layer.weight", torch.ones(1, dtype=torch.bfloat16))] + + backend.transfer_bucket( + buffer, + flush_cache=False, + weight_version="sync-v2", + ) + + backend._synchronizer._transfer_single_bucket.assert_called_once_with( + buffer, + flush_cache=False, + weight_version="sync-v2", + ) From c1fb95eb0a576576059fc51bc5f8b624c2e3e076 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 3 Apr 2026 15:28:04 -0700 Subject: [PATCH 10/41] Improve RMSNorm modes and fused residual path (#50) * Improve RMSNorm modes and fused residual path (#37) * Improve RMSNorm modes and fused residual path * Support native and compile RMSNorm for Qwen3.5 * Improve RMSNorm benchmark and Qwen3.5 fusion * Remove quack RMSNorm mode support * Update server docs for rmsnorm_mode * Remove stale quack benchmark tags * Fix pre-commit lint and formatting issues * Fix nested config flattening to propagate all fields The load_server_arguments() nested-config path previously cherry-picked a hardcoded subset of model/train keys, silently dropping fields like rmsnorm_mode, ep_dispatch, router_fp32, and all lora config. Replace with generic flattening that copies all keys from each section, with explicit remapping only for worker.* prefixes and lora.exclude_modules -> qlora_exclude_modules. New fields added to ServerArguments + to_config_dict() are now picked up automatically. * Add missing fields to to_config_dict() for complete round-trip model_name, pp_variable_seq_lengths, log_level, and sync_inference_method were not emitted by to_config_dict(), causing them to be silently lost when round-tripping through nested YAML configs via load_server_arguments(). (cherry picked from commit a77acd4802f9f2121a55dedd9d578997e2834ab3) --- .../content/docs/config-reference/server.md | 2 +- .../content/docs/server-training/sglang.mdx | 4 +- .../full/qwen3_235b_a22b_8node_ep64.yaml | 2 +- src/xorl/arguments.py | 8 ++ src/xorl/cli/direct_train.py | 1 + src/xorl/models/auto.py | 6 +- src/xorl/models/layers/__init__.py | 4 +- src/xorl/models/layers/normalization.py | 100 +++++++++++++++++- .../transformers/qwen3/modeling_qwen3.py | 9 +- .../transformers/qwen3_5/modeling_qwen3_5.py | 41 +++++-- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 41 +++++-- .../qwen3_moe/modeling_qwen3_moe.py | 9 +- src/xorl/server/launcher.py | 93 +++++++--------- src/xorl/server/runner/model_runner.py | 2 +- src/xorl/server/server_arguments.py | 15 ++- src/xorl/trainers/model_builder.py | 4 +- src/xorl/trainers/trainer.py | 1 + 17 files changed, 242 insertions(+), 100 deletions(-) diff --git a/docs/src/content/docs/config-reference/server.md b/docs/src/content/docs/config-reference/server.md index e2faf5d8..d24655b7 100644 --- a/docs/src/content/docs/config-reference/server.md +++ b/docs/src/content/docs/config-reference/server.md @@ -47,7 +47,7 @@ These flags align the training model's numerics with the inference engine (SGLan |---|---|---| | `router_fp32` | `true` | Upcast MoE router gate logits to float32 for numerical stability. | | `lm_head_fp32` | `true` | Upcast LM head logits to float32. | -| `rmsnorm_native` | `false` | Use unfused PyTorch RMSNorm instead of Triton kernel. | +| `rmsnorm_mode` | `native` | RMSNorm implementation: `eager`, `native`, or `compile`. | | `activation_native` | `false` | Use unfused SiLU instead of fused Triton kernel. | | `rope_native` | `false` | Use unfused RoPE instead of flash_attn kernel. | | `attention_cast_bf16` | `false` | Explicitly cast Q/K to BF16 after RoPE. | diff --git a/docs/src/content/docs/server-training/sglang.mdx b/docs/src/content/docs/server-training/sglang.mdx index 9f5697b8..2992bdbd 100644 --- a/docs/src/content/docs/server-training/sglang.mdx +++ b/docs/src/content/docs/server-training/sglang.mdx @@ -16,7 +16,7 @@ The core RL training loop requires capabilities that upstream SGLang does not pr 3. **MoE routing data export (R3)** — for MoE models, the inference server's routing decisions must be exported so training can replay them exactly, ensuring gradient consistency. Upstream SGLang does not export expert routing indices. -4. **Numerical alignment flags** — training and inference must produce identical logits for the same input. xorl-sglang supports the same numerical alignment flags as the xorl training server (`router_fp32`, `lm_head_fp32`, `rmsnorm_native`, `attention_cast_bf16`, etc.). +4. **Numerical alignment flags** — training and inference must produce identical logits for the same input. xorl-sglang supports the same numerical alignment flags as the xorl training server (`router_fp32`, `lm_head_fp32`, `rmsnorm_mode`, `attention_cast_bf16`, etc.). --- @@ -78,7 +78,7 @@ Configure both sides identically: # xorl server config (server_config.yaml) router_fp32: true lm_head_fp32: true -rmsnorm_native: false +rmsnorm_mode: compile activation_native: false rope_native: false attention_cast_bf16: false diff --git a/examples/server/configs/full/qwen3_235b_a22b_8node_ep64.yaml b/examples/server/configs/full/qwen3_235b_a22b_8node_ep64.yaml index 0982ee5a..bdaec631 100644 --- a/examples/server/configs/full/qwen3_235b_a22b_8node_ep64.yaml +++ b/examples/server/configs/full/qwen3_235b_a22b_8node_ep64.yaml @@ -11,7 +11,7 @@ moe_implementation: triton ep_dispatch: alltoall router_fp32: true lm_head_fp32: true -rmsnorm_native: true +rmsnorm_mode: native activation_native: true rope_native: true attention_cast_bf16: true diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 445df854..b98aafd0 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -472,6 +472,14 @@ class ModelArguments: "(e.g., tensor parallelism, independent LoRA per projection)." }, ) + rmsnorm_mode: Literal["eager", "native", "compile"] = field( + default="native", + metadata={ + "help": "RMSNorm implementation mode. 'native' uses torch.nn.functional.rms_norm " + "and is the default. 'compile' runs that native path through torch.compile. " + "'eager' uses the plain eager implementation." + }, + ) def __post_init__(self): if self.config_path is None and self.model_path is None: diff --git a/src/xorl/cli/direct_train.py b/src/xorl/cli/direct_train.py index 5f62c25c..53f2b560 100644 --- a/src/xorl/cli/direct_train.py +++ b/src/xorl/cli/direct_train.py @@ -176,6 +176,7 @@ def main(): deepep_buffer_size_gb=args.model.deepep_buffer_size_gb, deepep_num_sms=args.model.deepep_num_sms, deepep_async_combine=args.model.deepep_async_combine, + rmsnorm_mode=args.model.rmsnorm_mode, init_device=args.train.init_device, ) model_config = model.config diff --git a/src/xorl/models/auto.py b/src/xorl/models/auto.py index cf3affdd..af1d2d01 100644 --- a/src/xorl/models/auto.py +++ b/src/xorl/models/auto.py @@ -13,6 +13,7 @@ from ..distributed.parallel_state import get_parallel_state from ..utils import logging from .layers.attention import ATTENTION_FUNCTIONS +from .layers.normalization import set_rmsnorm_mode from .loader import ModelLoader, get_loader from .transformers.qwen3_5_shared import ( LINEAR_ATTENTION_RING_UNSUPPORTED_MESSAGE, @@ -110,7 +111,7 @@ def build_foundation_model( deepep_async_combine: bool = False, router_fp32: bool = True, lm_head_fp32: bool = True, - rmsnorm_native: bool = False, + rmsnorm_mode: Literal["eager", "native", "compile"] = "native", activation_native: bool = False, rope_native: bool = False, attention_cast_bf16: bool = False, @@ -144,7 +145,8 @@ def build_foundation_model( config._deepep_async_combine = deepep_async_combine config._router_fp32 = router_fp32 config._lm_head_fp32 = lm_head_fp32 - config._rmsnorm_native = rmsnorm_native + set_rmsnorm_mode(rmsnorm_mode) + config._rmsnorm_mode = rmsnorm_mode config._activation_native = activation_native config._rope_native = rope_native config._attention_cast_bf16 = attention_cast_bf16 diff --git a/src/xorl/models/layers/__init__.py b/src/xorl/models/layers/__init__.py index 1304bf36..202714cd 100644 --- a/src/xorl/models/layers/__init__.py +++ b/src/xorl/models/layers/__init__.py @@ -17,7 +17,7 @@ MoELoRAConfig, TopKRouter, ) -from .normalization import RMSNorm +from .normalization import RMSNorm, get_rmsnorm_mode, set_rmsnorm_mode from .rope import ( ROPE_INIT_FUNCTIONS, RotaryEmbedding, @@ -41,6 +41,7 @@ "MoELoRAConfig", "TopKRouter", "RMSNorm", + "get_rmsnorm_mode", "ROPE_INIT_FUNCTIONS", "RotaryEmbedding", "apply_rotary_pos_emb", @@ -50,5 +51,6 @@ "repeat_kv", "rope_config_validation", "rotate_half", + "set_rmsnorm_mode", "update_causal_mask", ] diff --git a/src/xorl/models/layers/normalization.py b/src/xorl/models/layers/normalization.py index 38fa7d61..7a0552cd 100644 --- a/src/xorl/models/layers/normalization.py +++ b/src/xorl/models/layers/normalization.py @@ -1,8 +1,28 @@ +from typing import Callable, Literal, Optional, Union + import torch +import torch.nn.functional as F from torch import nn -def rms_norm(hidden_states: torch.Tensor, weight: torch.Tensor, variance_epsilon: float) -> torch.Tensor: +RMSNormMode = Literal["eager", "native", "compile"] +_RMSNORM_MODE: RMSNormMode = "native" +_COMPILED_NATIVE_RMS_NORM: Optional[Callable[[torch.Tensor, torch.Tensor, float], torch.Tensor]] = None +_COMPILED_ZERO_CENTERED_RMS_NORM: Optional[Callable[[torch.Tensor, torch.Tensor, float], torch.Tensor]] = None + + +def set_rmsnorm_mode(mode: RMSNormMode) -> None: + global _RMSNORM_MODE + if mode not in {"eager", "native", "compile"}: + raise ValueError(f"Unsupported rmsnorm_mode: {mode}") + _RMSNORM_MODE = mode + + +def get_rmsnorm_mode() -> RMSNormMode: + return _RMSNORM_MODE + + +def eager_rms_norm(hidden_states: torch.Tensor, weight: torch.Tensor, variance_epsilon: float) -> torch.Tensor: """Eager implementation of RMSNorm.""" input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) @@ -11,16 +31,86 @@ def rms_norm(hidden_states: torch.Tensor, weight: torch.Tensor, variance_epsilon return weight * hidden_states.to(input_dtype) +def native_rms_norm(hidden_states: torch.Tensor, weight: torch.Tensor, variance_epsilon: float) -> torch.Tensor: + return F.rms_norm(hidden_states, (weight.shape[0],), weight, eps=variance_epsilon) + + +def compiled_rms_norm(hidden_states: torch.Tensor, weight: torch.Tensor, variance_epsilon: float) -> torch.Tensor: + global _COMPILED_NATIVE_RMS_NORM + if _COMPILED_NATIVE_RMS_NORM is None: + _COMPILED_NATIVE_RMS_NORM = torch.compile(native_rms_norm) + return _COMPILED_NATIVE_RMS_NORM(hidden_states, weight, variance_epsilon) + + +def eager_zero_centered_rms_norm( + hidden_states: torch.Tensor, + weight: torch.Tensor, + variance_epsilon: float, +) -> torch.Tensor: + output = hidden_states.float() + output = output * torch.rsqrt(output.pow(2).mean(-1, keepdim=True) + variance_epsilon) + output = output * (1.0 + weight.float()) + return output.type_as(hidden_states) + + +def native_zero_centered_rms_norm( + hidden_states: torch.Tensor, + weight: torch.Tensor, + variance_epsilon: float, +) -> torch.Tensor: + output = F.rms_norm( + hidden_states.float(), + (weight.shape[0],), + 1.0 + weight.float(), + eps=variance_epsilon, + ) + return output.type_as(hidden_states) + + +def compiled_zero_centered_rms_norm( + hidden_states: torch.Tensor, + weight: torch.Tensor, + variance_epsilon: float, +) -> torch.Tensor: + global _COMPILED_ZERO_CENTERED_RMS_NORM + if _COMPILED_ZERO_CENTERED_RMS_NORM is None: + _COMPILED_ZERO_CENTERED_RMS_NORM = torch.compile(native_zero_centered_rms_norm) + return _COMPILED_ZERO_CENTERED_RMS_NORM(hidden_states, weight, variance_epsilon) + + class RMSNorm(nn.Module): """Root Mean Square Layer Normalization.""" - def __init__(self, hidden_size: int, eps: float = 1e-6): + def __init__(self, hidden_size: int, eps: float = 1e-6, mode: Optional[RMSNormMode] = None): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps + self.mode: RMSNormMode = mode or get_rmsnorm_mode() + + def forward( + self, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor] = None, + prenorm: bool = False, + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + residual_out: Optional[torch.Tensor] = None + norm_input = hidden_states + if residual is not None: + residual_out = hidden_states + residual + norm_input = residual_out + + if self.mode == "eager": + out = eager_rms_norm(norm_input, self.weight, self.variance_epsilon) + elif self.mode == "native": + out = native_rms_norm(norm_input, self.weight, self.variance_epsilon) + elif self.mode == "compile": + out = compiled_rms_norm(norm_input, self.weight, self.variance_epsilon) + else: + raise ValueError(f"Unsupported rmsnorm_mode: {self.mode}") - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return rms_norm(hidden_states, self.weight, self.variance_epsilon) + if residual_out is not None and prenorm: + return out, residual_out + return out def extra_repr(self): - return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}, mode={self.mode}" diff --git a/src/xorl/models/transformers/qwen3/modeling_qwen3.py b/src/xorl/models/transformers/qwen3/modeling_qwen3.py index 4b5b441a..f977e37f 100644 --- a/src/xorl/models/transformers/qwen3/modeling_qwen3.py +++ b/src/xorl/models/transformers/qwen3/modeling_qwen3.py @@ -110,11 +110,12 @@ def forward( position_embeddings=position_embeddings, **kwargs, ) - hidden_states = residual + hidden_states - # Fully Connected - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states, residual = self.post_attention_layernorm( + hidden_states, + residual=residual, + prenorm=True, + ) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states diff --git a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py index 7bf43347..938db2d1 100644 --- a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py +++ b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py @@ -19,6 +19,12 @@ ) from xorl.models.layers.attention.backend import ATTENTION_FUNCTIONS from xorl.models.layers.attention.backend.eager import eager_attention_forward +from xorl.models.layers.normalization import ( + compiled_zero_centered_rms_norm, + eager_zero_centered_rms_norm, + get_rmsnorm_mode, + native_zero_centered_rms_norm, +) from xorl.models.module_utils import GradientCheckpointingLayer from xorl.models.outputs import BaseModelOutput, CausalLMOutput from xorl.models.transformers.qwen3_5 import parallelize @@ -91,14 +97,32 @@ def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.zeros(dim)) + self.mode = get_rmsnorm_mode() - def _norm(self, x): - return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + def forward( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + prenorm: bool = False, + ): + residual_out: Optional[torch.Tensor] = None + norm_input = x + if residual is not None: + residual_out = x + residual + norm_input = residual_out + + if self.mode == "eager": + out = eager_zero_centered_rms_norm(norm_input, self.weight, self.eps) + elif self.mode == "native": + out = native_zero_centered_rms_norm(norm_input, self.weight, self.eps) + elif self.mode == "compile": + out = compiled_zero_centered_rms_norm(norm_input, self.weight, self.eps) + else: + raise NotImplementedError(f"Unsupported rmsnorm_mode for Qwen3.5 RMSNorm: {self.mode}") - def forward(self, x): - output = self._norm(x.float()) - output = output * (1.0 + self.weight.float()) - return output.type_as(x) + if residual_out is not None and prenorm: + return out, residual_out + return out class Qwen3_5Attention(nn.Module): @@ -265,10 +289,7 @@ def forward( position_embeddings=position_embeddings, **kwargs, ) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual=residual, prenorm=True) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states diff --git a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py index c37b3555..a6d2b90f 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -18,6 +18,12 @@ from xorl.models.layers.attention.backend import ATTENTION_FUNCTIONS from xorl.models.layers.attention.backend.eager import eager_attention_forward from xorl.models.layers.moe import MoEBlock +from xorl.models.layers.normalization import ( + compiled_zero_centered_rms_norm, + eager_zero_centered_rms_norm, + get_rmsnorm_mode, + native_zero_centered_rms_norm, +) from xorl.models.outputs import MoeCausalLMOutput, MoeModelOutput from xorl.models.transformers.qwen3_5_moe import parallelize from xorl.models.transformers.qwen3_5_moe.checkpoint_handler import Qwen3_5MoeCheckpointHandler @@ -88,14 +94,32 @@ def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.zeros(dim)) + self.mode = get_rmsnorm_mode() - def _norm(self, x): - return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + def forward( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + prenorm: bool = False, + ): + residual_out: Optional[torch.Tensor] = None + norm_input = x + if residual is not None: + residual_out = x + residual + norm_input = residual_out + + if self.mode == "eager": + out = eager_zero_centered_rms_norm(norm_input, self.weight, self.eps) + elif self.mode == "native": + out = native_zero_centered_rms_norm(norm_input, self.weight, self.eps) + elif self.mode == "compile": + out = compiled_zero_centered_rms_norm(norm_input, self.weight, self.eps) + else: + raise NotImplementedError(f"Unsupported rmsnorm_mode for Qwen3.5 MoE RMSNorm: {self.mode}") - def forward(self, x): - output = self._norm(x.float()) - output = output * (1.0 + self.weight.float()) - return output.type_as(x) + if residual_out is not None and prenorm: + return out, residual_out + return out class Qwen3_5MoeAttention(nn.Module): @@ -306,10 +330,7 @@ def forward( position_embeddings=position_embeddings, **kwargs, ) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual=residual, prenorm=True) hidden_states = self.mlp(hidden_states) if isinstance(hidden_states, tuple): hidden_states, router_logits = hidden_states diff --git a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py index 6f4e7fad..84cf861f 100644 --- a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py +++ b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py @@ -267,11 +267,12 @@ def forward( position_embeddings=position_embeddings, **kwargs, ) - hidden_states = residual + hidden_states - # Fully Connected - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states, residual = self.post_attention_layernorm( + hidden_states, + residual=residual, + prenorm=True, + ) if _selective and "mlp" in self._recompute_modules: hidden_states = self._gradient_checkpointing_func( diff --git a/src/xorl/server/launcher.py b/src/xorl/server/launcher.py index 5eaf5c58..ca8fc2f7 100644 --- a/src/xorl/server/launcher.py +++ b/src/xorl/server/launcher.py @@ -370,66 +370,51 @@ def load_server_arguments(config_path: str, overrides: Optional[Dict[str, any]] # Check if this is a nested config (has model/train/worker sections) if "model" in config and isinstance(config["model"], dict): - # Nested config - flatten it + # Nested config - flatten all keys from each section into one dict. + # Keys whose names already match a ServerArguments field are copied + # directly; a few nested keys need explicit remapping (worker.*, + # lora.exclude_modules). This is forward-compatible: new fields + # added to ServerArguments + to_config_dict() are picked up + # automatically without touching this flattening logic. flat_config = {} - # Model section - model_config = config.get("model", {}) - flat_config["model_path"] = model_config.get("model_path") - flat_config["model_name"] = model_config.get("model_name") - flat_config["config_path"] = model_config.get("config_path") - flat_config["tokenizer_path"] = model_config.get("tokenizer_path") - flat_config["attn_implementation"] = model_config.get("attn_implementation", "flash_attention_3") - flat_config["moe_implementation"] = model_config.get("moe_implementation") - # Train section - train_config = config.get("train", {}) - flat_config["data_parallel_mode"] = train_config.get("data_parallel_mode", "fsdp2") - flat_config["data_parallel_shard_size"] = train_config.get("data_parallel_shard_size", 1) - flat_config["data_parallel_replicate_size"] = train_config.get("data_parallel_replicate_size", 1) - flat_config["ulysses_parallel_size"] = train_config.get("ulysses_parallel_size", 1) - flat_config["expert_parallel_size"] = train_config.get("expert_parallel_size", 1) - flat_config["enable_mixed_precision"] = train_config.get("enable_mixed_precision", True) - flat_config["enable_gradient_checkpointing"] = train_config.get("enable_gradient_checkpointing", True) - flat_config["enable_full_shard"] = train_config.get("enable_full_shard", True) - flat_config["enable_activation_offload"] = train_config.get("enable_activation_offload", False) - flat_config["init_device"] = train_config.get("init_device", "meta") - flat_config["load_checkpoint_path"] = train_config.get("load_checkpoint_path", "") - flat_config["ckpt_manager"] = train_config.get("ckpt_manager", "dcp") - flat_config["log_level"] = train_config.get("log_level", "INFO") - - # Data processing section (can be in train or data section) - data_config = config.get("data", {}) - flat_config["sample_packing_sequence_len"] = train_config.get("sample_packing_sequence_len") or data_config.get( - "sample_packing_sequence_len", 32000 - ) - flat_config["enable_packing"] = train_config.get("enable_packing", data_config.get("enable_packing", True)) - - # Output directory (can be in train section or top-level) - used for checkpoints, sampler weights, logs - # Note: This uses output_dir from config, which should be on shared filesystem for multi-node - flat_config["output_dir"] = train_config.get("output_dir", config.get("output_dir", "outputs")) + # model.* and train.* keys map 1:1 to ServerArguments fields + for section in ("model", "train"): + for k, v in config.get(section, {}).items(): + flat_config[k] = v - # Storage limit (can be in train section or top-level) - limits disk usage for output_dir - flat_config["storage_limit"] = train_config.get("storage_limit", config.get("storage_limit")) - - # Idle session timeout (can be in train section or top-level) - sessions inactive for this duration are cleaned up - flat_config["idle_session_timeout"] = train_config.get( - "idle_session_timeout", config.get("idle_session_timeout", 7200.0) - ) + # lora.* keys also map 1:1 except exclude_modules → qlora_exclude_modules + for k, v in config.get("lora", {}).items(): + if k == "exclude_modules": + flat_config["qlora_exclude_modules"] = v + else: + flat_config[k] = v - # Training flags - flat_config["skip_initial_checkpoint"] = train_config.get("skip_initial_checkpoint", False) - flat_config["log_gradient_norms"] = train_config.get("log_gradient_norms", True) - flat_config["log_router_stats"] = train_config.get("log_router_stats", True) - flat_config["freeze_router"] = train_config.get("freeze_router", True) + # data.* — only a few fields are relevant for the server + data_config = config.get("data", {}) + if "sample_packing_sequence_len" not in flat_config: + flat_config["sample_packing_sequence_len"] = data_config.get("sample_packing_sequence_len", 32000) + if "enable_packing" not in flat_config: + flat_config["enable_packing"] = data_config.get("enable_packing", True) - # Worker section + # worker.* keys are prefixed with worker_ in ServerArguments worker_config = config.get("worker", {}) - flat_config["worker_bind_address"] = worker_config.get("bind_address", "auto") - flat_config["worker_bind_host"] = worker_config.get("bind_host", "0.0.0.0") - flat_config["worker_bind_port"] = worker_config.get("bind_port", 5556) - flat_config["engine_connect_host"] = worker_config.get("engine_connect_host") - flat_config["worker_connection_timeout"] = worker_config.get("connection_timeout", 120.0) - flat_config["worker_max_retries"] = worker_config.get("max_retries", 3) + _worker_key_map = { + "bind_address": "worker_bind_address", + "bind_host": "worker_bind_host", + "bind_port": "worker_bind_port", + "engine_connect_host": "engine_connect_host", + "connection_timeout": "worker_connection_timeout", + "max_retries": "worker_max_retries", + } + for nested_key, flat_key in _worker_key_map.items(): + if nested_key in worker_config: + flat_config[flat_key] = worker_config[nested_key] + + # Top-level keys that can appear outside any section + for k in ("output_dir", "storage_limit", "idle_session_timeout"): + if k not in flat_config and k in config: + flat_config[k] = config[k] filtered_config = {k: v for k, v in flat_config.items() if k in valid_fields and v is not None} else: diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index 7e63817a..de7f67a2 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -494,7 +494,7 @@ def _initialize_model(self): freeze_router=self.train_config.get("freeze_router", False), router_fp32=self.model_config.get("router_fp32", True), lm_head_fp32=self.model_config.get("lm_head_fp32", True), - rmsnorm_native=self.model_config.get("rmsnorm_native", False), + rmsnorm_mode=self.model_config.get("rmsnorm_mode", "native"), activation_native=self.model_config.get("activation_native", False), rope_native=self.model_config.get("rope_native", False), attention_cast_bf16=self.model_config.get("attention_cast_bf16", False), diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index 887f2eac..9a7c0e04 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -99,8 +99,13 @@ class ServerArguments: default=True, metadata={"help": "Upcast LM head logits computation to float32 for numerical stability."} ) - rmsnorm_native: bool = field( - default=False, metadata={"help": "Use native RMSNorm (no fused kernels) for SGLang alignment."} + rmsnorm_mode: Literal["eager", "native", "compile"] = field( + default="native", + metadata={ + "help": "RMSNorm implementation mode. 'native' uses torch.nn.functional.rms_norm " + "and is the default. 'compile' runs that native path through torch.compile. " + "'eager' uses the plain eager implementation." + }, ) activation_native: bool = field( @@ -454,6 +459,7 @@ def to_config_dict(self) -> Dict[str, Any]: config = { "model": { "model_path": self.model_path, + "model_name": self.model_name, "config_path": self.config_path, "tokenizer_path": self.tokenizer_path, "attn_implementation": self.attn_implementation, @@ -468,7 +474,7 @@ def to_config_dict(self) -> Dict[str, Any]: "merge_qkv": self.merge_qkv, "router_fp32": self.router_fp32, "lm_head_fp32": self.lm_head_fp32, - "rmsnorm_native": self.rmsnorm_native, + "rmsnorm_mode": self.rmsnorm_mode, "activation_native": self.activation_native, "rope_native": self.rope_native, "attention_cast_bf16": self.attention_cast_bf16, @@ -511,6 +517,9 @@ def to_config_dict(self) -> Dict[str, Any]: "freeze_router": self.freeze_router, "pipeline_parallel_size": self.pipeline_parallel_size, "pipeline_parallel_schedule": self.pipeline_parallel_schedule, + "pp_variable_seq_lengths": self.pp_variable_seq_lengths, + "log_level": self.log_level, + "sync_inference_method": self.sync_inference_method, }, "data": { # Empty data section - data comes from client at runtime diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index e5395c25..92ac391b 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -76,7 +76,7 @@ def build_training_model( # --- SGLang numerical alignment --- router_fp32: bool = True, lm_head_fp32: bool = True, - rmsnorm_native: bool = False, + rmsnorm_mode: str = "native", activation_native: bool = False, rope_native: bool = False, attention_cast_bf16: bool = False, @@ -117,7 +117,7 @@ def build_training_model( deepep_async_combine=deepep_async_combine, router_fp32=router_fp32, lm_head_fp32=lm_head_fp32, - rmsnorm_native=rmsnorm_native, + rmsnorm_mode=rmsnorm_mode, activation_native=activation_native, rope_native=rope_native, attention_cast_bf16=attention_cast_bf16, diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index 526dd777..ac9ac395 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -359,6 +359,7 @@ def _build_model(self) -> None: deepep_buffer_size_gb=args.model.deepep_buffer_size_gb, deepep_num_sms=args.model.deepep_num_sms, deepep_async_combine=args.model.deepep_async_combine, + rmsnorm_mode=args.model.rmsnorm_mode, init_device=args.train.init_device, ) self.model_config = self.model.config From 3c52a2ed0e66f530f3f956a0cb56dd98a7dc9429 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Sun, 5 Apr 2026 21:03:58 -0700 Subject: [PATCH 11/41] Use rank-local node cache dirs for Triton, TorchInductor, and Quack (#82) * Use rank-local node cache dirs for Triton and Quack * Share rank-local cache setup across entrypoints (cherry picked from commit 49af7821117750a1d84b8daf59c1bd2b9822d9af) --- src/xorl/cli/direct_train.py | 5 +++++ src/xorl/cli/train.py | 5 +++++ src/xorl/server/runner/setup.py | 19 ++----------------- src/xorl/utils/compile_cache.py | 25 +++++++++++++++++++++++++ 4 files changed, 37 insertions(+), 17 deletions(-) create mode 100644 src/xorl/utils/compile_cache.py diff --git a/src/xorl/cli/direct_train.py b/src/xorl/cli/direct_train.py index 53f2b560..adb869a2 100644 --- a/src/xorl/cli/direct_train.py +++ b/src/xorl/cli/direct_train.py @@ -1,9 +1,14 @@ +# ruff: noqa: E402 + import os +from xorl.utils.compile_cache import configure_rank_local_compile_caches + # Must be set before importing torch / initializing CUDA so the # allocator picks up the setting on first use. os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") +configure_rank_local_compile_caches() import json import socket diff --git a/src/xorl/cli/train.py b/src/xorl/cli/train.py index 2351c126..80d91d94 100644 --- a/src/xorl/cli/train.py +++ b/src/xorl/cli/train.py @@ -1,9 +1,14 @@ +# ruff: noqa: E402 + import os +from xorl.utils.compile_cache import configure_rank_local_compile_caches + # Must be set before importing torch / initializing CUDA so the # allocator picks up the setting on first use. os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") +configure_rank_local_compile_caches() from xorl.arguments import Arguments, parse_args from xorl.trainers import Trainer diff --git a/src/xorl/server/runner/setup.py b/src/xorl/server/runner/setup.py index ede90fe4..8a36d191 100644 --- a/src/xorl/server/runner/setup.py +++ b/src/xorl/server/runner/setup.py @@ -21,6 +21,7 @@ from xorl.server.runner.model_runner import ModelRunner from xorl.server.runner.runner_dispatcher import RunnerDispatcher from xorl.server.server_arguments import ServerArguments, parse_server_args +from xorl.utils.compile_cache import configure_rank_local_compile_caches from xorl.utils.device import get_nccl_backend @@ -266,23 +267,7 @@ def main(): os.environ.setdefault("NCCL_NVLS_ENABLE", "0") os.environ.setdefault("TORCH_NCCL_BLOCKING_WAIT", "1") - # Set unique Triton cache directory per rank to avoid race conditions - # during parallel kernel compilation - local_rank = os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0")) - triton_base = os.environ.get("TRITON_CACHE_DIR", os.path.expanduser("~/.triton")) - triton_cache_dir = os.path.join(triton_base, f"cache_rank{local_rank}") - os.environ["TRITON_CACHE_DIR"] = triton_cache_dir - os.makedirs(triton_cache_dir, exist_ok=True) - - # Set per-rank TorchInductor cache to prevent cross-rank cubin path conflicts. - # TorchInductor's FxGraphCache at /tmp/torchinductor_/ is node-local and shared - # by all local ranks. It stores absolute cubin paths from the compiling rank's - # TRITON_CACHE_DIR, which breaks when a different local rank gets a cache hit. - inductor_local_rank = os.environ.get("LOCAL_RANK", "0") - inductor_base = os.environ.get("TORCHINDUCTOR_CACHE_DIR", f"/tmp/torchinductor_{os.environ.get('USER', 'unknown')}") - inductor_cache_dir = os.path.join(inductor_base, f"rank{inductor_local_rank}") - os.environ["TORCHINDUCTOR_CACHE_DIR"] = inductor_cache_dir - os.makedirs(inductor_cache_dir, exist_ok=True) + configure_rank_local_compile_caches() # Find config file path config_path = None diff --git a/src/xorl/utils/compile_cache.py b/src/xorl/utils/compile_cache.py new file mode 100644 index 00000000..eb54aa4d --- /dev/null +++ b/src/xorl/utils/compile_cache.py @@ -0,0 +1,25 @@ +import os + + +def configure_rank_local_compile_caches() -> None: + """Put compile/autotune caches on rank-local paths under a user-scoped root.""" + cache_user = os.environ.get("USER", "unknown") + + triton_rank = os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0")) + triton_base = os.environ.get("TRITON_CACHE_DIR", f"/tmp/triton_cache_{cache_user}") + triton_cache_dir = os.path.join(triton_base, f"cache_rank{triton_rank}") + os.environ["TRITON_CACHE_DIR"] = triton_cache_dir + os.makedirs(triton_cache_dir, exist_ok=True) + + inductor_rank = os.environ.get("LOCAL_RANK", "0") + inductor_base = os.environ.get("TORCHINDUCTOR_CACHE_DIR", f"/tmp/torchinductor_{cache_user}") + inductor_cache_dir = os.path.join(inductor_base, f"rank{inductor_rank}") + os.environ["TORCHINDUCTOR_CACHE_DIR"] = inductor_cache_dir + os.makedirs(inductor_cache_dir, exist_ok=True) + + # Quack keeps its own autotune cache separate from Triton's cache. + quack_rank = os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0")) + quack_base = os.environ.get("QUACK_CACHE_DIR", f"/tmp/quack_cache_{cache_user}") + quack_cache_dir = os.path.join(quack_base, f"cache_rank{quack_rank}") + os.environ["QUACK_CACHE_DIR"] = quack_cache_dir + os.makedirs(quack_cache_dir, exist_ok=True) From dc543e11f4e060d92a7213c2cfd3297abb7e1a62 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Sun, 5 Apr 2026 21:04:28 -0700 Subject: [PATCH 12/41] Default train_router=False for MoE (#53) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * DeepEP improvements: train_router fix, NFS retry, benchmark infra (#41) * Add local benchmark configs/jobs and improve DeepEP async combine Benchmark infra: - Add 1-node and 2-node EP16 benchmark configs for alltoall and deepep - Fix 2-node k8s jobs: POD_NAME downward API for correct node_rank, hostIPC, privileged, convert StatefulSet to two coordinated Jobs (backoffLimit=0) - Set rmsnorm_mode=compile in all benchmark configs DeepEP improvements: - Default async_combine=True to overlap combine comm with shared expert compute - scatter_add_ → index_add_ in unpermute step - Remove redundant .contiguous() before combine * Revert DeepEP async_combine changes (slower on H100) * Add 4-node EP32 Qwen3-235B-A22B benchmark configs (deepep + alltoall) * Default train_router=False for MoE, fix packing NFS retry, add benchmark scripts - Set train_router default to False so routing weights are detached before expert computation. Router is trained only via auxiliary losses, ensuring consistent convergence between DeepEP and AllToAll dispatch backends. DeepEP's custom autograd cannot propagate gradients through routing weights (buffer communication breaks the autograd graph), causing worse convergence when train_router=True. - Increase packing bins cache retry from 10x2s to 30x5s for multi-node NFS visibility delays. - Add local benchmark configs and slurm scripts for 30B 2-node and 235B 4-node benchmarks (deepep, alltoall, split_broadcast, A/B test). - Add NFS parallel read throughput benchmark (1-rank and 32-rank variants). - Add DeepEP vs AllToAll precision comparison test. * Assert train_router=False in MoE forward DeepEP cannot propagate gradients through routing weights, so train_router=True would silently produce different convergence between dispatch backends. Fail fast instead. * Remove unnecessary benchmark files and test configs Clean up AB test configs, split_broadcast variants, NFS bench scripts, and precision test that were used during development. * Remove benchmark results from repo * Move slurm scripts to experiments/local_benchmark/slurm/ * Remove packing.py NFS retry change (moved to separate PR) * Remove experiments/ (moved to local-benchmark PR) * Restore experiments/ to match main * Make train_router explicit for DeepEP * Align train_router defaults with PR intent --------- Co-authored-by: kiddyboots216 (cherry picked from commit 956da30f6542e4f85d700241d1ca871c4e4a1f0b) --- .../configs/full/qwen3_30b_a3b_muon.yaml | 1 + .../configs/full/qwen3_coder_30b_a3b.yaml | 1 + .../full/qwen3_235b_a22b_8node_ep64.yaml | 5 +- .../configs/full/qwen3_5_35b_a3b_full.yaml | 44 +++++++++++ src/xorl/arguments.py | 7 ++ src/xorl/cli/direct_train.py | 1 + src/xorl/models/auto.py | 8 ++ src/xorl/models/layers/moe/moe_block.py | 14 ++-- src/xorl/server/runner/model_runner.py | 2 +- src/xorl/server/server_arguments.py | 9 +++ src/xorl/trainers/model_builder.py | 2 + src/xorl/trainers/trainer.py | 1 + .../models/test_moe_train_router_dispatch.py | 79 +++++++++++++++++++ 13 files changed, 167 insertions(+), 7 deletions(-) create mode 100644 examples/server/configs/full/qwen3_5_35b_a3b_full.yaml create mode 100644 tests/models/test_moe_train_router_dispatch.py diff --git a/examples/local/dummy/configs/full/qwen3_30b_a3b_muon.yaml b/examples/local/dummy/configs/full/qwen3_30b_a3b_muon.yaml index 5abe89fb..0b063c82 100644 --- a/examples/local/dummy/configs/full/qwen3_30b_a3b_muon.yaml +++ b/examples/local/dummy/configs/full/qwen3_30b_a3b_muon.yaml @@ -3,6 +3,7 @@ model: attn_implementation: flash_attention_3 moe_implementation: quack ep_dispatch: deepep + train_router: false deepep_buffer_size_gb: 2.0 deepep_num_sms: 48 diff --git a/examples/local/dummy/configs/full/qwen3_coder_30b_a3b.yaml b/examples/local/dummy/configs/full/qwen3_coder_30b_a3b.yaml index cdbb3044..e4164390 100644 --- a/examples/local/dummy/configs/full/qwen3_coder_30b_a3b.yaml +++ b/examples/local/dummy/configs/full/qwen3_coder_30b_a3b.yaml @@ -3,6 +3,7 @@ model: attn_implementation: flash_attention_3 moe_implementation: quack ep_dispatch: deepep + train_router: false deepep_buffer_size_gb: 2.0 deepep_num_sms: 48 diff --git a/examples/server/configs/full/qwen3_235b_a22b_8node_ep64.yaml b/examples/server/configs/full/qwen3_235b_a22b_8node_ep64.yaml index bdaec631..1ef5a128 100644 --- a/examples/server/configs/full/qwen3_235b_a22b_8node_ep64.yaml +++ b/examples/server/configs/full/qwen3_235b_a22b_8node_ep64.yaml @@ -8,7 +8,10 @@ model_path: Qwen/Qwen3-235B-A22B-Instruct-2507 tokenizer_path: Qwen/Qwen3-235B-A22B-Instruct-2507 attn_implementation: flash_attention_3 moe_implementation: triton -ep_dispatch: alltoall +ep_dispatch: deepep +train_router: false +deepep_buffer_size_gb: 2.0 +deepep_num_sms: 48 router_fp32: true lm_head_fp32: true rmsnorm_mode: native diff --git a/examples/server/configs/full/qwen3_5_35b_a3b_full.yaml b/examples/server/configs/full/qwen3_5_35b_a3b_full.yaml new file mode 100644 index 00000000..d58083c9 --- /dev/null +++ b/examples/server/configs/full/qwen3_5_35b_a3b_full.yaml @@ -0,0 +1,44 @@ +# Server-side configuration for XORL Training Server (Qwen3.5-35B-A3B full weights, bf16) +# +# Full weight RL fine-tuning of the 35B MoE model. +# 8 GPUs: Ulysses SP=8, EP=8, FSDP shard=1. Muon optimizer (1 state) fits in GPU memory. + +model_path: Qwen/Qwen3.5-35B-A3B +tokenizer_path: Qwen/Qwen3.5-35B-A3B +attn_implementation: flash_attention_3 +moe_implementation: quack +ep_dispatch: deepep +train_router: false + +data_parallel_mode: fsdp2 +expert_parallel_size: 8 +ulysses_parallel_size: 8 +data_parallel_replicate_size: 1 +data_parallel_shard_size: 1 + +enable_mixed_precision: true +enable_gradient_checkpointing: true +enable_full_shard: true +enable_activation_offload: false +init_device: meta + +output_dir: outputs/Qwen3.5-35B-A3B-server-full-rl +load_checkpoint_path: "" +ckpt_manager: dcp + +log_level: INFO + +worker_connection_timeout: 60.0 +worker_max_retries: 3 + +sample_packing_sequence_len: 128000 +enable_packing: true + +skip_initial_checkpoint: true +ce_mode: compiled + +optimizer: muon +optimizer_dtype: bf16 +muon_lr: 0.02 +muon_momentum: 0.95 +muon_adjust_lr_fn: match_rms_adamw diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index b98aafd0..149cd2b8 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -446,6 +446,13 @@ class ModelArguments: default="alltoall", metadata={"help": "EP dispatch strategy: 'alltoall' (default) or 'deepep' (NVLink-optimized)."}, ) + train_router: bool = field( + default=False, + metadata={ + "help": "Whether expert computation gradients should flow through routing weights. " + "Disabled by default and must remain False when ep_dispatch='deepep'." + }, + ) deepep_buffer_size_gb: float = field( default=2.0, metadata={"help": "DeepEP buffer size in GB (effective when ep_dispatch='deepep')."}, diff --git a/src/xorl/cli/direct_train.py b/src/xorl/cli/direct_train.py index adb869a2..d348d74d 100644 --- a/src/xorl/cli/direct_train.py +++ b/src/xorl/cli/direct_train.py @@ -178,6 +178,7 @@ def main(): attn_implementation=args.model.attn_implementation, moe_implementation=args.model.moe_implementation, ep_dispatch=args.model.ep_dispatch, + train_router=args.model.train_router, deepep_buffer_size_gb=args.model.deepep_buffer_size_gb, deepep_num_sms=args.model.deepep_num_sms, deepep_async_combine=args.model.deepep_async_combine, diff --git a/src/xorl/models/auto.py b/src/xorl/models/auto.py index af1d2d01..c6dbdc8f 100644 --- a/src/xorl/models/auto.py +++ b/src/xorl/models/auto.py @@ -106,6 +106,7 @@ def build_foundation_model( ] = "flash_attention_3", moe_implementation: Optional[Literal["eager", "triton", "native", "quack"]] = None, ep_dispatch: str = "alltoall", + train_router: bool = False, deepep_buffer_size_gb: float = 2.0, deepep_num_sms: int = 20, deepep_async_combine: bool = False, @@ -139,7 +140,14 @@ def build_foundation_model( config._moe_implementation = moe_implementation logger.info_rank0(f"Moe implementation: {moe_implementation}") + if ep_dispatch == "deepep" and train_router: + raise ValueError( + "train_router=True is not supported with ep_dispatch='deepep'. " + "Set train_router=False or use ep_dispatch='alltoall'." + ) + config._ep_dispatch = ep_dispatch + config.train_router = train_router config._deepep_buffer_size_gb = deepep_buffer_size_gb config._deepep_num_sms = deepep_num_sms config._deepep_async_combine = deepep_async_combine diff --git a/src/xorl/models/layers/moe/moe_block.py b/src/xorl/models/layers/moe/moe_block.py index 7a83ab0a..85022102 100644 --- a/src/xorl/models/layers/moe/moe_block.py +++ b/src/xorl/models/layers/moe/moe_block.py @@ -35,7 +35,7 @@ class MoEBlock(nn.Module): norm_topk_prob: Whether to renormalize top-k routing weights. moe_implementation: Backend name — ``"eager"``, ``"triton"``, ``"native"``, or ``"quack"``. train_router: If True (default), gradients flow from expert - computation through routing weights back to the gate. If False, + computation through routing weights back to the gate. If False, routing weights are detached before expert computation so the gate is only trained via auxiliary losses on ``router_logits``. """ @@ -190,9 +190,13 @@ def forward(self, hidden_states: torch.Tensor): # No replay active: use standard router routing_weights, selected_experts = self.router(router_logits, hidden_states.dtype) - # When train_router is False, detach routing_weights so the gate is - # only trained via auxiliary losses on router_logits, not through - # expert computation. + ep_dispatch = getattr(self.experts, "ep_dispatch", "alltoall") + if self.train_router and ep_dispatch == "deepep": + raise AssertionError( + "train_router=True is not supported with ep_dispatch='deepep'. " + "DeepEP cannot propagate gradients through routing weights. " + "Set train_router=False or switch to ep_dispatch='alltoall'." + ) if not self.train_router: routing_weights = routing_weights.detach() @@ -239,5 +243,5 @@ def from_config(cls, config, moe_implementation: str = "triton"): hidden_act=config.hidden_act, norm_topk_prob=config.norm_topk_prob, moe_implementation=moe_implementation, - train_router=getattr(config, "train_router", True), + train_router=getattr(config, "train_router", False), ) diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index de7f67a2..710d8390 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -466,6 +466,7 @@ def _initialize_model(self): attn_implementation=self.model_config.get("attn_implementation", "sdpa"), moe_implementation=self.model_config.get("moe_implementation"), ep_dispatch=self.model_config.get("ep_dispatch", "alltoall"), + train_router=self.model_config.get("train_router", False), deepep_buffer_size_gb=self.model_config.get("deepep_buffer_size_gb", 2.0), deepep_num_sms=self.model_config.get("deepep_num_sms", 20), deepep_async_combine=self.model_config.get("deepep_async_combine", False), @@ -967,7 +968,6 @@ def _forward_loop( ): """Core forward (+ optional backward) loop shared between forward and forward_backward.""" params = loss_fn_params or {} - return_per_token = params.get("return_per_token", True) # Count valid tokens globally global_valid_tokens = self._count_global_valid_tokens(micro_batches) diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index 9a7c0e04..f6ab2ec0 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -78,6 +78,14 @@ class ServerArguments: metadata={"help": "EP dispatch strategy: 'alltoall' (default) or 'deepep' (NVLink-optimized)."}, ) + train_router: bool = field( + default=False, + metadata={ + "help": "Whether expert computation gradients should flow through routing weights. " + "Disabled by default and must remain False when ep_dispatch='deepep'." + }, + ) + deepep_buffer_size_gb: float = field( default=2.0, metadata={"help": "DeepEP buffer size in GB (effective when ep_dispatch='deepep')."} ) @@ -465,6 +473,7 @@ def to_config_dict(self) -> Dict[str, Any]: "attn_implementation": self.attn_implementation, "moe_implementation": self.moe_implementation, "ep_dispatch": self.ep_dispatch, + "train_router": self.train_router, "deepep_buffer_size_gb": self.deepep_buffer_size_gb, "deepep_num_sms": self.deepep_num_sms, "deepep_async_combine": self.deepep_async_combine, diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index 92ac391b..e236184c 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -43,6 +43,7 @@ def build_training_model( attn_implementation: str = "flash_attention_3", moe_implementation: Optional[str] = None, ep_dispatch: str = "alltoall", + train_router: bool = False, deepep_buffer_size_gb: float = 2.0, deepep_num_sms: int = 20, deepep_async_combine: bool = False, @@ -112,6 +113,7 @@ def build_training_model( attn_implementation=attn_implementation, moe_implementation=moe_implementation, ep_dispatch=ep_dispatch, + train_router=train_router, deepep_buffer_size_gb=deepep_buffer_size_gb, deepep_num_sms=deepep_num_sms, deepep_async_combine=deepep_async_combine, diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index ac9ac395..1450490a 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -356,6 +356,7 @@ def _build_model(self) -> None: attn_implementation=args.model.attn_implementation, moe_implementation=args.model.moe_implementation, ep_dispatch=args.model.ep_dispatch, + train_router=args.model.train_router, deepep_buffer_size_gb=args.model.deepep_buffer_size_gb, deepep_num_sms=args.model.deepep_num_sms, deepep_async_combine=args.model.deepep_async_combine, diff --git a/tests/models/test_moe_train_router_dispatch.py b/tests/models/test_moe_train_router_dispatch.py new file mode 100644 index 00000000..b62377fc --- /dev/null +++ b/tests/models/test_moe_train_router_dispatch.py @@ -0,0 +1,79 @@ +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from xorl.arguments import ModelArguments +from xorl.models.layers.moe.moe_block import MoEBlock +from xorl.server.server_arguments import ServerArguments + + +pytestmark = [pytest.mark.cpu] + + +def test_train_router_true_allowed_with_alltoall(): + moe = MoEBlock( + hidden_size=16, + num_experts=4, + top_k=2, + intermediate_size=32, + moe_implementation="eager", + train_router=True, + ) + moe.experts.ep_dispatch = "alltoall" + moe.train() + for p in moe.parameters(): + nn.init.normal_(p, std=0.01) + + x = torch.randn(1, 4, 16, requires_grad=True) + out, _ = moe(x) + out.sum().backward() + + assert moe.gate.weight.grad is not None + assert torch.isfinite(moe.gate.weight.grad).all() + assert moe.gate.weight.grad.abs().sum() > 0 + + +def test_train_router_true_rejected_with_deepep(): + moe = MoEBlock( + hidden_size=16, + num_experts=4, + top_k=2, + intermediate_size=32, + moe_implementation="eager", + train_router=True, + ) + moe.experts.ep_dispatch = "deepep" + + x = torch.randn(1, 4, 16) + with pytest.raises(AssertionError, match="ep_dispatch='deepep'"): + moe(x) + + +def test_model_arguments_default_train_router_false(): + args = ModelArguments(config_path="Qwen/Qwen3-8B") + + assert args.train_router is False + + +def test_server_arguments_default_train_router_false(): + args = ServerArguments(model_path="Qwen/Qwen3-8B") + + assert args.train_router is False + assert args.to_config_dict()["model"]["train_router"] is False + + +def test_from_config_defaults_train_router_false(): + config = SimpleNamespace( + hidden_size=16, + num_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=32, + hidden_act="silu", + norm_topk_prob=True, + ) + + moe = MoEBlock.from_config(config, moe_implementation="eager") + + assert moe.train_router is False From 58af94368d0e71a7fad6866e518f8821d7e69b1b Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Mon, 6 Apr 2026 09:52:14 -0700 Subject: [PATCH 13/41] Fix LoRA mixed-precision and 2-node 397B all-ranks loading paths (#72) (cherry picked from commit e3f720dcb52f87f997be264941c20a1b46c4f1c5) --- src/xorl/cli/direct_train.py | 39 ++- src/xorl/distributed/torch_parallelize.py | 6 +- src/xorl/models/layers/rope.py | 12 +- src/xorl/models/module_utils.py | 258 +++++++++++++++--- .../qwen3_5_moe/checkpoint_handler.py | 14 +- .../qwen3_moe/checkpoint_handler.py | 14 +- src/xorl/server/runner/model_runner.py | 17 +- src/xorl/trainers/model_builder.py | 61 ++++- src/xorl/trainers/trainer.py | 39 ++- .../test_lora_mixed_precision_dtype_policy.py | 110 ++++++++ 10 files changed, 501 insertions(+), 69 deletions(-) create mode 100644 tests/trainers/test_lora_mixed_precision_dtype_policy.py diff --git a/src/xorl/cli/direct_train.py b/src/xorl/cli/direct_train.py index d348d74d..e29ca888 100644 --- a/src/xorl/cli/direct_train.py +++ b/src/xorl/cli/direct_train.py @@ -14,7 +14,7 @@ import socket import time from dataclasses import asdict -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional import torch.distributed as dist from tqdm import trange @@ -33,6 +33,11 @@ from xorl.models.layers.moe.aux_loss import global_load_balancing_loss_func from xorl.models.module_utils import compute_loss from xorl.optim import build_lr_scheduler, build_optimizer +from xorl.trainers.model_builder import ( + maybe_upcast_trainable_adapter_params, + resolve_training_model_dtype, + should_skip_generic_param_upcast, +) from xorl.trainers.training_utils import ( clip_gradients, count_valid_tokens, @@ -50,6 +55,17 @@ logger = helper.create_logger(__name__) +_direct_train_cpu_group: Optional[dist.ProcessGroup] = None + + +def _get_direct_train_cpu_group() -> Optional[dist.ProcessGroup]: + """Return a cached CPU/Gloo group for bootstrap object collectives.""" + global _direct_train_cpu_group + if not dist.is_available() or not dist.is_initialized() or dist.get_world_size() <= 1: + return None + if _direct_train_cpu_group is None: + _direct_train_cpu_group = dist.new_group(backend="gloo") + return _direct_train_cpu_group def main(): @@ -85,7 +101,7 @@ def main(): "hostname": socket.gethostname(), } gathered_hosts = [None] * args.train.world_size - dist.all_gather_object(gathered_hosts, host_payload) + dist.all_gather_object(gathered_hosts, host_payload, group=_get_direct_train_cpu_group()) if args.train.global_rank == 0: unique_hostnames = sorted({item["hostname"] for item in gathered_hosts if item is not None}) rank_to_hostname = {str(item["global_rank"]): item["hostname"] for item in gathered_hosts if item is not None} @@ -171,10 +187,15 @@ def main(): logger.info_rank0(f"Save every {args.train.save_epochs} epoch(s) = every {save_epoch_steps} steps") logger.info_rank0("Prepare model") + model_dtype = resolve_training_model_dtype( + enable_lora=args.lora.enable_lora, + enable_qlora=args.lora.enable_qlora, + enable_mixed_precision=args.train.enable_mixed_precision, + ) model = build_foundation_model( config_path=args.model.config_path, weights_path=args.model.model_path, - torch_dtype="float32" if args.train.enable_mixed_precision else "bfloat16", + torch_dtype=model_dtype, attn_implementation=args.model.attn_implementation, moe_implementation=args.model.moe_implementation, ep_dispatch=args.model.ep_dispatch, @@ -261,6 +282,13 @@ def main(): ) helper.print_device_mem_info("VRAM usage after LoRA injection") + maybe_upcast_trainable_adapter_params( + model, + enable_lora=args.lora.enable_lora, + enable_qlora=args.lora.enable_qlora, + enable_mixed_precision=args.train.enable_mixed_precision, + ) + get_optimizer_pre_hook = getattr(model, "get_optimizer_pre_hook", None) build_result = build_parallelize_model( model, @@ -278,7 +306,10 @@ def main(): load_weights_mode=args.train.load_weights_mode, pp_schedule=args.train.pipeline_parallel_schedule if args.train.pipeline_parallel_size > 1 else None, reshard_after_forward=args.train.reshard_after_forward, - skip_param_upcast=args.lora.enable_qlora, + skip_param_upcast=should_skip_generic_param_upcast( + enable_lora=args.lora.enable_lora, + enable_qlora=args.lora.enable_qlora, + ), ) # PP returns dict with stages + model_parts; otherwise returns model directly diff --git a/src/xorl/distributed/torch_parallelize.py b/src/xorl/distributed/torch_parallelize.py index b7c434db..5fc1d4c7 100644 --- a/src/xorl/distributed/torch_parallelize.py +++ b/src/xorl/distributed/torch_parallelize.py @@ -369,8 +369,12 @@ def build_parallelize_model( if not parallel_state.fsdp_enabled: if kwargs.get("init_device") not in ["cuda", "npu"]: raise ValueError("Only FSDP training supports `init_device=cpu` or `init_device=meta`.") - if enable_mixed_precision and not kwargs.pop("skip_param_upcast", False): + skip_param_upcast = kwargs.pop("skip_param_upcast", False) + if enable_mixed_precision and not skip_param_upcast: + logger.info_rank0("Applying generic fp32 param upcast before parallelization") model = model.float() + elif enable_mixed_precision: + logger.info_rank0("Skipping generic fp32 param upcast before parallelization") if pp_schedule is not None: # ---- Pipeline Parallelism path ---- diff --git a/src/xorl/models/layers/rope.py b/src/xorl/models/layers/rope.py index 5f5861d4..94eafe71 100644 --- a/src/xorl/models/layers/rope.py +++ b/src/xorl/models/layers/rope.py @@ -538,7 +538,17 @@ def rope_config_validation(config, ignore_keys=None): FutureWarning, ) config.standardize_rope_params() - config.validate_rope(ignore_keys=ignore_keys) + if ignore_keys is None: + config.validate_rope() + return + + try: + config.validate_rope(ignore_keys=ignore_keys) + except TypeError as exc: + if "unexpected keyword argument 'ignore_keys'" not in str(exc): + raise + # transformers>=5.5 removed the ignore_keys kwarg from validate_rope(). + config.validate_rope() __all__ = [ diff --git a/src/xorl/models/module_utils.py b/src/xorl/models/module_utils.py index a7c11234..ad472b01 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -5,6 +5,7 @@ from collections import OrderedDict from contextlib import contextmanager from dataclasses import dataclass +from datetime import timedelta from functools import partial from typing import TYPE_CHECKING, Any, Callable, Dict, Generator, List, Literal, Optional, Sequence, Set, Tuple, Union @@ -19,9 +20,14 @@ from xorl.distributed.parallel_state import get_parallel_state from xorl.lora.modules.linear import LoraLinear +from xorl.models.checkpoint_handlers.buffers import ( # noqa: F401 + ExpertWeightBuffer, + checkpoint_has_per_expert_weights, + parse_expert_key, +) from xorl.ops.loss import get_loss_function from xorl.utils import logging -from xorl.utils.device import synchronize +from xorl.utils.device import get_device_id, get_device_type, synchronize from xorl.utils.helper import empty_cache, get_dtype_size @@ -36,13 +42,55 @@ logger = logging.get_logger(__name__) +_cpu_group = None +_weight_load_group = None -# Re-export checkpoint handler utilities for backward compatibility (used by tests) -from xorl.models.checkpoint_handlers.buffers import ( # noqa: F401 - ExpertWeightBuffer, - checkpoint_has_per_expert_weights, - parse_expert_key, -) + +def _get_cpu_group(): + """Get or create a Gloo process group for CPU-based broadcasts. + + Using Gloo instead of the default NCCL backend avoids deadlocks when + broadcast_object_list is the first collective operation (NCCL lazy init + can hang if not all ranks are ready simultaneously). + """ + global _cpu_group + if _cpu_group is None and dist.is_initialized(): + _cpu_group = dist.new_group(backend="gloo") + return _cpu_group + + +def _get_weight_load_group(): + """Get or create a dedicated process group for weight loading. + + Weight loading can spend many minutes in rank-0 shard I/O between collectives. + Use a separate process group with a larger timeout so NCCL does not trip the + default watchdog while other ranks are waiting for the next broadcast. + """ + global _weight_load_group + if _weight_load_group is None and dist.is_initialized(): + timeout_sec = int(os.getenv("XORL_WEIGHT_LOAD_TIMEOUT_SEC", "7200")) + _weight_load_group = dist.new_group(backend=dist.get_backend(), timeout=timedelta(seconds=timeout_sec)) + return _weight_load_group + + +def _get_weight_load_object_device() -> Optional[torch.device]: + """Return the device to use for NCCL object broadcasts in the load path.""" + group = _get_weight_load_group() + if group is None: + return None + if dist.get_backend(group) == "nccl": + return torch.device(f"{get_device_type()}:{get_device_id()}") + return None + + +def _broadcast_object_list_weight_load(obj_list: List[Any], src: int) -> None: + """Broadcast Python metadata for weight loading on the dedicated load group.""" + group = _get_weight_load_group() + device = _get_weight_load_object_device() + kwargs = {"src": src, "group": group} + if device is not None: + kwargs["device"] = device + dist.broadcast_object_list(obj_list, **kwargs) def _get_checkpoint_keys(weights_path: str) -> Optional[Set[str]]: @@ -326,6 +374,25 @@ def _try_load_state_dict(weights_path: str, **kwargs): logger.error(f"Failed to get shard files after {max_retries} attempts") raise + _broadcast_object_list_weight_load(resolved_paths, src=0) + shard_files = resolved_paths[0] + if shard_files is None: + return None + return [StateDictIterator(f) for f in shard_files] + + +def _try_load_state_dict_local(weights_path: str, **kwargs): + """Resolve shard file paths locally (no broadcast). Used by rank 0 or single-rank.""" + cache_kwargs = {"_raise_exceptions_for_missing_entries": False, **kwargs} + resolved_weight_file = cached_file(weights_path, SAFE_WEIGHTS_NAME, **cache_kwargs) + if resolved_weight_file: + return [StateDictIterator(resolved_weight_file)] + + resolved_weight_file = cached_file(weights_path, SAFE_WEIGHTS_INDEX_NAME, **cache_kwargs) + if resolved_weight_file: + shard_files, _ = get_checkpoint_shard_files(weights_path, resolved_weight_file, **kwargs) + return [StateDictIterator(shard_file) for shard_file in shard_files] + resolved_weight_file = cached_file(weights_path, WEIGHTS_NAME, **cache_kwargs) if resolved_weight_file: return [StateDictIterator(resolved_weight_file)] @@ -392,6 +459,64 @@ def _build_compiled_key_map( return mapping +def _get_expert_scatter_target_shape( + model: Union["nn.Module", "PreTrainedModel"], + parameter_name: str, + tensor: "torch.Tensor", + parallel_plan: Optional["ParallelPlan"], + parallel_state, +) -> Optional[Tuple[int, ...]]: + """Return the EP-local parameter shape if this tensor can be rank0-scattered. + + We only use this optimized path when the current rank topology has a simple + 2D EP mesh (ep, ep_fsdp). More complex layouts, such as PP+EP, fall back to + the existing full-tensor broadcast path. + """ + if ( + parallel_plan is None + or not parallel_state.ep_enabled + or parallel_state.ep_fsdp_device_mesh is None + or parallel_state.ep_fsdp_device_mesh.mesh.ndim != 2 + or not parallel_plan._is_expert_parameter(parameter_name) + ): + return None + + module, local_name = _find_submodule(model, parameter_name) + param = module._parameters.get(local_name) + if param is None or param.ndim == 0 or tensor.ndim == 0: + return None + + target_shape = tuple(param.shape) + if not target_shape or tensor.shape[0] <= target_shape[0] or tensor.shape[0] % target_shape[0] != 0: + return None + + return target_shape + + +def _build_expert_scatter_list( + tensor: "torch.Tensor", + target_shape: Tuple[int, ...], + parallel_state, + torch_device: "torch.device", +) -> Tuple["torch.Tensor", List["torch.Tensor"]]: + """Prepare rank0-side per-rank expert views for NCCL scatter.""" + if tensor.device.type == "cpu" and torch_device.type == "cuda": + full_tensor = tensor.pin_memory().to(torch_device, non_blocking=True) + else: + full_tensor = tensor.to(torch_device, non_blocking=True) + + ep_mesh = parallel_state.ep_fsdp_device_mesh.mesh.cpu() + local_experts = target_shape[0] + scatter_list: List[torch.Tensor] = [None] * int(ep_mesh.numel()) + for ep_rank in range(ep_mesh.shape[0]): + local_view = full_tensor.narrow(0, ep_rank * local_experts, local_experts) + for ep_fsdp_rank in range(ep_mesh.shape[1]): + global_rank = int(ep_mesh[ep_rank, ep_fsdp_rank]) + scatter_list[global_rank] = local_view + + return full_tensor, scatter_list + + class _MultiStreamDMA: """Manages multi-stream H2D DMA transfers for overlapping copy engine usage. @@ -998,7 +1123,7 @@ def rank0_load_and_broadcast_weights( dtype=torch.int64, device=torch_device if torch_device.type != "cpu" else torch.device("cpu"), ) - dist.broadcast(shard_count_tensor, src=0) + dist.broadcast(shard_count_tensor, src=0, group=_get_weight_load_group()) shard_count = int(shard_count_tensor.item()) # Rank 0: create prefetching generator that bulk-loads upcoming shards in @@ -1039,32 +1164,61 @@ def rank0_load_and_broadcast_weights( # broadcast tensors one by one. This replaces N metadata broadcasts # with 1, eliminating per-tensor pickle + NCCL launch overhead. if global_rank == 0: - batch_meta = [(n, t.shape, t.dtype) for n, t in merged_queue] + batch_meta = [] + for name, tensor in merged_queue: + model_name = _compiled_key_map.get(name, name) + target_shape = None + if model_name in parameter_names_to_load: + target_shape = _get_expert_scatter_target_shape(model, model_name, tensor, parallel_plan, _ps) + if target_shape is not None: + batch_meta.append((name, torch.Size(target_shape), tensor.dtype, "expert_scatter")) + else: + batch_meta.append((name, tensor.shape, tensor.dtype, "broadcast")) else: batch_meta = None batch_meta = [batch_meta] - dist.broadcast_object_list(batch_meta, src=0) + _broadcast_object_list_weight_load(batch_meta, src=0) batch_meta = batch_meta[0] - for name, shape, dtype in batch_meta: - # Broadcast tensor from rank 0 to all ranks. - # For expert tensors, this broadcasts the full [num_experts, ...] tensor; - # _dispatch_parameter -> shard_tensor handles EP slicing locally per rank. - if global_rank != 0: - tensor = torch.empty(shape, dtype=dtype, device=torch_device) + for name, shape, dtype, transfer_mode in batch_meta: + if global_rank == 0: + _, source_tensor = merged_queue.pop(0) else: - _, broadcast_tensor = merged_queue.pop(0) - if broadcast_tensor.device.type == "cpu" and torch_device.type == "cuda": - tensor = broadcast_tensor.pin_memory().to(torch_device, non_blocking=True) - else: - tensor = broadcast_tensor.to(torch_device, non_blocking=True) + source_tensor = None start_time = time.perf_counter() - dist.broadcast(tensor, src=0) - logger.info_rank0( - f"{name=}, {shape=}, {dtype=}, broadcast time (ms): {1000 * (time.perf_counter() - start_time)}" - ) + if transfer_mode == "expert_scatter": + if global_rank == 0: + full_tensor, scatter_list = _build_expert_scatter_list( + source_tensor, tuple(shape), _ps, torch_device + ) + tensor = scatter_list[global_rank].clone() + full_shape = tuple(source_tensor.shape) + else: + full_tensor = None + scatter_list = None + tensor = torch.empty(shape, dtype=dtype, device=torch_device) + full_shape = shape + + dist.scatter(tensor, scatter_list=scatter_list, src=0, group=_get_weight_load_group()) + logger.info_rank0( + f"{name=}, full_shape={full_shape}, local_shape={shape}, {dtype=}, " + f"scatter time (ms): {1000 * (time.perf_counter() - start_time)}" + ) + else: + if global_rank != 0: + tensor = torch.empty(shape, dtype=dtype, device=torch_device) + else: + if source_tensor.device.type == "cpu" and torch_device.type == "cuda": + tensor = source_tensor.pin_memory().to(torch_device, non_blocking=True) + else: + tensor = source_tensor.to(torch_device, non_blocking=True) + + dist.broadcast(tensor, src=0, group=_get_weight_load_group()) + logger.info_rank0( + f"{name=}, {shape=}, {dtype=}, broadcast time (ms): {1000 * (time.perf_counter() - start_time)}" + ) # Resolve _orig_mod prefix from torch.compile model_name = _compiled_key_map.get(name, name) @@ -1079,7 +1233,10 @@ def rank0_load_and_broadcast_weights( del tensor if global_rank == 0: - del broadcast_tensor + del source_tensor + if transfer_mode == "expert_scatter": + del full_tensor + del scatter_list empty_cache() @@ -1089,25 +1246,47 @@ def rank0_load_and_broadcast_weights( # Broadcast remaining items from handler flush (same batched pattern) if global_rank == 0: - flush_meta = [(n, t.shape, t.dtype) for n, t in merged_queue] + flush_meta = [] + for name, tensor in merged_queue: + model_name = _compiled_key_map.get(name, name) + target_shape = None + if model_name in parameter_names_to_load: + target_shape = _get_expert_scatter_target_shape(model, model_name, tensor, parallel_plan, _ps) + if target_shape is not None: + flush_meta.append((name, torch.Size(target_shape), tensor.dtype, "expert_scatter")) + else: + flush_meta.append((name, tensor.shape, tensor.dtype, "broadcast")) else: flush_meta = None flush_meta = [flush_meta] - dist.broadcast_object_list(flush_meta, src=0) + _broadcast_object_list_weight_load(flush_meta, src=0) flush_meta = flush_meta[0] - for name, shape, dtype in flush_meta: - if global_rank != 0: - tensor = torch.empty(shape, dtype=dtype, device=torch_device) + for name, shape, dtype, transfer_mode in flush_meta: + if global_rank == 0: + _, source_tensor = merged_queue.pop(0) else: - _, broadcast_tensor = merged_queue.pop(0) - if broadcast_tensor.device.type == "cpu" and torch_device.type == "cuda": - tensor = broadcast_tensor.pin_memory().to(torch_device, non_blocking=True) - else: - tensor = broadcast_tensor.to(torch_device, non_blocking=True) + source_tensor = None - dist.broadcast(tensor, src=0) + if transfer_mode == "expert_scatter": + if global_rank == 0: + full_tensor, scatter_list = _build_expert_scatter_list(source_tensor, tuple(shape), _ps, torch_device) + tensor = scatter_list[global_rank].clone() + else: + full_tensor = None + scatter_list = None + tensor = torch.empty(shape, dtype=dtype, device=torch_device) + dist.scatter(tensor, scatter_list=scatter_list, src=0, group=_get_weight_load_group()) + else: + if global_rank != 0: + tensor = torch.empty(shape, dtype=dtype, device=torch_device) + else: + if source_tensor.device.type == "cpu" and torch_device.type == "cuda": + tensor = source_tensor.pin_memory().to(torch_device, non_blocking=True) + else: + tensor = source_tensor.to(torch_device, non_blocking=True) + dist.broadcast(tensor, src=0, group=_get_weight_load_group()) # Resolve _orig_mod prefix from torch.compile model_name = _compiled_key_map.get(name, name) @@ -1117,7 +1296,10 @@ def rank0_load_and_broadcast_weights( del tensor if global_rank == 0: - del broadcast_tensor + del source_tensor + if transfer_mode == "expert_scatter": + del full_tensor + del scatter_list post_process_after_weight_loading(model, buffer_dict, parameter_names_to_load, dtensor_factory) diff --git a/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py b/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py index 505202b6..0ee076f6 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py @@ -137,11 +137,17 @@ def _maybe_finalize_per_expert_merge( def _handle_fused_expert_weights(self, key: str, tensor: torch.Tensor) -> Optional[List[Tuple[str, torch.Tensor]]]: internal_gate_up_match = self._INTERNAL_FUSED_EXPERT_GATE_UP_PATTERN.match(key) if internal_gate_up_match is not None: - return [(key, self._slice_expert_tensor_for_ep(tensor))] + tensor = self._slice_expert_tensor_for_ep(tensor) + if tensor.ndim == 3 and tensor.shape[1] < tensor.shape[2]: + tensor = tensor.transpose(1, 2).contiguous() + return [(key, tensor)] internal_down_match = self._INTERNAL_FUSED_EXPERT_DOWN_PATTERN.match(key) if internal_down_match is not None: - return [(key, self._slice_expert_tensor_for_ep(tensor))] + tensor = self._slice_expert_tensor_for_ep(tensor) + if tensor.ndim == 3 and tensor.shape[1] > tensor.shape[2]: + tensor = tensor.transpose(1, 2).contiguous() + return [(key, tensor)] gate_up_match = self._HF_FUSED_EXPERT_GATE_UP_PATTERN.match(key) if gate_up_match is not None: @@ -167,7 +173,7 @@ def _handle_fused_expert_weights(self, key: str, tensor: torch.Tensor) -> Option proj = split_match.group(2) tensor = self._slice_expert_tensor_for_ep(tensor) if proj == "down": - return [(f"model.layers.{layer_idx}.mlp.experts.down_proj", tensor)] + return [(f"model.layers.{layer_idx}.mlp.experts.down_proj", tensor.transpose(1, 2).contiguous())] pending = self._stacked_gate_up_pending.setdefault(layer_idx, {}) pending[proj] = tensor @@ -178,7 +184,7 @@ def _handle_fused_expert_weights(self, key: str, tensor: torch.Tensor) -> Option return [ ( f"model.layers.{layer_idx}.mlp.experts.gate_up_proj", - torch.cat([gate, up], dim=2), + torch.cat([gate.transpose(1, 2), up.transpose(1, 2)], dim=2).contiguous(), ) ] return [] diff --git a/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py b/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py index 52ab41cd..a48eb537 100644 --- a/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py @@ -154,11 +154,17 @@ def _handle_stacked_expert_weights( ) -> Optional[List[Tuple[str, torch.Tensor]]]: gate_up_match = self._FUSED_EXPERT_GATE_UP_PATTERN.match(key) if gate_up_match is not None: - return [(key, self._slice_expert_tensor_for_ep(tensor))] + tensor = self._slice_expert_tensor_for_ep(tensor) + if tensor.ndim == 3 and tensor.shape[1] < tensor.shape[2]: + tensor = tensor.transpose(1, 2).contiguous() + return [(key, tensor)] down_match = self._FUSED_EXPERT_DOWN_PATTERN.match(key) if down_match is not None: - return [(key, self._slice_expert_tensor_for_ep(tensor))] + tensor = self._slice_expert_tensor_for_ep(tensor) + if tensor.ndim == 3 and tensor.shape[1] > tensor.shape[2]: + tensor = tensor.transpose(1, 2).contiguous() + return [(key, tensor)] split_match = self._STACKED_EXPERT_SPLIT_PATTERN.match(key) if split_match is None: @@ -169,7 +175,7 @@ def _handle_stacked_expert_weights( tensor = self._slice_expert_tensor_for_ep(tensor) if proj == "down": - return [(f"model.layers.{layer_idx}.mlp.experts.down_proj", tensor)] + return [(f"model.layers.{layer_idx}.mlp.experts.down_proj", tensor.transpose(1, 2).contiguous())] pending = self._stacked_gate_up_pending.setdefault(layer_idx, {}) pending[proj] = tensor @@ -180,7 +186,7 @@ def _handle_stacked_expert_weights( return [ ( f"model.layers.{layer_idx}.mlp.experts.gate_up_proj", - torch.cat([gate, up], dim=2), + torch.cat([gate.transpose(1, 2), up.transpose(1, 2)], dim=2).contiguous(), ) ] return [] diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index 710d8390..51280374 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -40,7 +40,10 @@ from xorl.server.runner.adapters import LoRAAdapterManager from xorl.server.runner.checkpoint import CheckpointManager from xorl.server.runner.utils import MoeMetricsTracker, RoutingReplayHandler, run_self_test, validate_token_ids -from xorl.trainers.model_builder import build_training_model +from xorl.trainers.model_builder import ( + build_training_model, + resolve_training_model_dtype, +) from xorl.trainers.training_utils import ( clip_gradients, count_valid_tokens, @@ -448,13 +451,11 @@ def _initialize_model(self): # Resolve target_modules from Tinker-style or flat config target_modules = self._resolve_lora_target_modules() if lora_enabled else None - # Determine model dtype - if (lora_enabled or enable_qlora) and enable_mixed_precision: - model_dtype = "bfloat16" - elif enable_mixed_precision: - model_dtype = "float32" - else: - model_dtype = "bfloat16" + model_dtype = resolve_training_model_dtype( + enable_lora=lora_enabled, + enable_qlora=enable_qlora, + enable_mixed_precision=enable_mixed_precision, + ) pp_size = self.train_config.get("pipeline_parallel_size", 1) pp_schedule_name = self.train_config.get("pipeline_parallel_schedule", "1F1B") if pp_size > 1 else None diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index e236184c..3dd54f39 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -34,6 +34,51 @@ class TrainingModelResult: exclude_modules: Set[str] = field(default_factory=set) +def resolve_training_model_dtype( + *, + enable_lora: bool, + enable_qlora: bool, + enable_mixed_precision: bool, +) -> str: + """Return the foundation-model dtype for the requested training mode. + + Full-weight mixed-precision training keeps parameters in fp32 before FSDP + wrapping. LoRA/QLoRA instead keep the frozen base weights in bf16 and only + upcast the trainable adapter weights to fp32. + """ + if (enable_lora or enable_qlora) and enable_mixed_precision: + return "bfloat16" + if enable_mixed_precision: + return "float32" + return "bfloat16" + + +def should_skip_generic_param_upcast( + *, + enable_lora: bool, + enable_qlora: bool, +) -> bool: + """Whether the generic full-model fp32 upcast should be skipped.""" + return enable_lora or enable_qlora + + +def maybe_upcast_trainable_adapter_params( + model: nn.Module, + *, + enable_lora: bool, + enable_qlora: bool, + enable_mixed_precision: bool, +) -> None: + """Upcast trainable adapter weights to fp32 while leaving the frozen base in bf16.""" + if not enable_mixed_precision or not (enable_lora or enable_qlora): + return + + for param in model.parameters(): + if param.requires_grad: + param.data = param.data.to(torch.float32) + logger.info_rank0("Upcast trainable LoRA params to float32") + + def build_training_model( *, # --- Model --- @@ -178,11 +223,12 @@ def build_training_model( # ------------------------------------------------------------------ # 4. LoRA + mixed precision: upcast trainable params to fp32 # ------------------------------------------------------------------ - if (enable_lora or enable_qlora) and enable_mixed_precision: - for param in model.parameters(): - if param.requires_grad: - param.data = param.data.to(torch.float32) - logger.info_rank0("Upcast trainable LoRA params to float32") + maybe_upcast_trainable_adapter_params( + model, + enable_lora=enable_lora, + enable_qlora=enable_qlora, + enable_mixed_precision=enable_mixed_precision, + ) # ------------------------------------------------------------------ # 5. Save optimizer pre-hook (some models register hooks) @@ -207,7 +253,10 @@ def build_training_model( load_weights_mode=load_weights_mode, pp_schedule=pp_schedule, reshard_after_forward=reshard_after_forward, - skip_param_upcast=enable_qlora, + skip_param_upcast=should_skip_generic_param_upcast( + enable_lora=enable_lora, + enable_qlora=enable_qlora, + ), ) pp_enabled = isinstance(build_result, dict) diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index 1450490a..1cea27ea 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -34,6 +34,11 @@ from xorl.models.layers.moe.routing_replay import RoutingReplay, set_replay_stage from xorl.models.module_utils import compute_loss from xorl.optim import build_lr_scheduler, build_optimizer +from xorl.trainers.model_builder import ( + maybe_upcast_trainable_adapter_params, + resolve_training_model_dtype, + should_skip_generic_param_upcast, +) from xorl.trainers.training_utils import ( clip_gradients, count_valid_tokens, @@ -50,6 +55,17 @@ logger = helper.create_logger(__name__) +_trainer_cpu_group: Optional[dist.ProcessGroup] = None + + +def _get_trainer_cpu_group() -> Optional[dist.ProcessGroup]: + """Return a cached CPU/Gloo group for object collectives in trainer bootstrap.""" + global _trainer_cpu_group + if not dist.is_available() or not dist.is_initialized() or dist.get_world_size() <= 1: + return None + if _trainer_cpu_group is None: + _trainer_cpu_group = dist.new_group(backend="gloo") + return _trainer_cpu_group # --------------------------------------------------------------------------- @@ -237,8 +253,10 @@ def _collect_host_inventory(self) -> List[Dict[str, Any]]: "local_rank": self.args.train.local_rank, "hostname": socket.gethostname(), } + if not dist.is_available() or not dist.is_initialized() or self.args.train.world_size <= 1: + return [payload] gathered: List[Optional[Dict[str, Any]]] = [None] * self.args.train.world_size - dist.all_gather_object(gathered, payload) + dist.all_gather_object(gathered, payload, group=_get_trainer_cpu_group()) return [item for item in gathered if item is not None] def _log_host_inventory(self) -> None: @@ -349,10 +367,15 @@ def _build_model(self) -> None: """Build foundation model and inject LoRA/QLoRA if configured.""" args = self.args logger.info_rank0("Prepare model") + model_dtype = resolve_training_model_dtype( + enable_lora=args.lora.enable_lora, + enable_qlora=args.lora.enable_qlora, + enable_mixed_precision=args.train.enable_mixed_precision, + ) self.model = build_foundation_model( config_path=args.model.config_path, weights_path=args.model.model_path, - torch_dtype="float32" if args.train.enable_mixed_precision else "bfloat16", + torch_dtype=model_dtype, attn_implementation=args.model.attn_implementation, moe_implementation=args.model.moe_implementation, ep_dispatch=args.model.ep_dispatch, @@ -382,6 +405,13 @@ def _build_model(self) -> None: elif args.lora.enable_lora: self._inject_lora() + maybe_upcast_trainable_adapter_params( + self.model, + enable_lora=args.lora.enable_lora, + enable_qlora=args.lora.enable_qlora, + enable_mixed_precision=args.train.enable_mixed_precision, + ) + # Save pre-hook before parallelization (some models register optimizer hooks) self._optimizer_pre_hook_fn = getattr(self.model, "get_optimizer_pre_hook", None) @@ -471,7 +501,10 @@ def _parallelize(self) -> None: load_weights_mode=args.train.load_weights_mode, pp_schedule=args.train.pipeline_parallel_schedule if args.train.pipeline_parallel_size > 1 else None, reshard_after_forward=args.train.reshard_after_forward, - skip_param_upcast=args.lora.enable_qlora, + skip_param_upcast=should_skip_generic_param_upcast( + enable_lora=args.lora.enable_lora, + enable_qlora=args.lora.enable_qlora, + ), ) logger.info_rank0(f"Model weights loaded in {time.time() - _t0:.1f}s") diff --git a/tests/trainers/test_lora_mixed_precision_dtype_policy.py b/tests/trainers/test_lora_mixed_precision_dtype_policy.py new file mode 100644 index 00000000..34154005 --- /dev/null +++ b/tests/trainers/test_lora_mixed_precision_dtype_policy.py @@ -0,0 +1,110 @@ +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from xorl.trainers.model_builder import ( + build_training_model, + resolve_training_model_dtype, + should_skip_generic_param_upcast, +) + + +class TinyModel(nn.Module): + _no_split_modules = [] + + def __init__(self, dtype: torch.dtype): + super().__init__() + self.config = SimpleNamespace(model_type="tiny") + self.proj = nn.Linear(4, 4, bias=False, dtype=dtype) + + +def test_lora_mixed_precision_keeps_base_bf16_and_skips_generic_upcast(monkeypatch): + captured = {} + + def fake_build_foundation_model(**kwargs): + return TinyModel(getattr(torch, kwargs["torch_dtype"])) + + def fake_parallelize(model, **kwargs): + captured["skip_param_upcast"] = kwargs["skip_param_upcast"] + captured["base_dtype"] = model.proj.weight.dtype + captured["lora_a_dtype"] = model.proj.lora_A.dtype + captured["lora_b_dtype"] = model.proj.lora_B.dtype + return model + + monkeypatch.setattr("xorl.models.build_foundation_model", fake_build_foundation_model) + monkeypatch.setattr( + "xorl.distributed.torch_parallelize.build_parallelize_model", + fake_parallelize, + ) + monkeypatch.setattr( + "xorl.trainers.model_builder.helper.print_device_mem_info", + lambda *args, **kwargs: None, + ) + + result = build_training_model( + config_path="unused", + weights_path="unused", + torch_dtype="bfloat16", + enable_lora=True, + lora_rank=8, + lora_alpha=16, + lora_target_modules=["proj"], + enable_mixed_precision=True, + ) + + assert captured["skip_param_upcast"] is True + assert captured["base_dtype"] == torch.bfloat16 + assert captured["lora_a_dtype"] == torch.float32 + assert captured["lora_b_dtype"] == torch.float32 + assert result.model.proj.weight.requires_grad is False + assert result.model.proj.lora_A.requires_grad is True + assert result.model.proj.lora_B.requires_grad is True + + +def test_lora_dtype_policy_matches_intended_training_modes(): + assert ( + resolve_training_model_dtype( + enable_lora=True, + enable_qlora=False, + enable_mixed_precision=True, + ) + == "bfloat16" + ) + assert ( + resolve_training_model_dtype( + enable_lora=False, + enable_qlora=True, + enable_mixed_precision=True, + ) + == "bfloat16" + ) + assert ( + resolve_training_model_dtype( + enable_lora=False, + enable_qlora=False, + enable_mixed_precision=True, + ) + == "float32" + ) + assert ( + should_skip_generic_param_upcast( + enable_lora=True, + enable_qlora=False, + ) + is True + ) + assert ( + should_skip_generic_param_upcast( + enable_lora=False, + enable_qlora=True, + ) + is True + ) + assert ( + should_skip_generic_param_upcast( + enable_lora=False, + enable_qlora=False, + ) + is False + ) From aca3f53300aaae241d0b3d0e740116537f374820 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Mon, 6 Apr 2026 15:21:48 -0700 Subject: [PATCH 14/41] Respect deferred QLoRA MoE expert loading (#68) * Respect deferred QLoRA MoE expert loading * Format qwen3 moe checkpoint handler * Filter skipped weights in rank0 broadcast load path * Update broadcast loader tests for weight-load group (cherry picked from commit 0184c0a8feb27e7dd18c1937d4d774a3d526210d) --- src/xorl/models/module_utils.py | 15 ++- .../qwen3_5_moe/checkpoint_handler.py | 35 ++++++- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 9 ++ .../qwen3_moe/checkpoint_handler.py | 31 +++++- .../qwen3_moe/modeling_qwen3_moe.py | 9 ++ tests/models/test_module_utils_broadcast.py | 97 ++++++++++++++++++- tests/models/test_moe_fused_gate_up_proj.py | 44 +++++++++ 7 files changed, 231 insertions(+), 9 deletions(-) diff --git a/src/xorl/models/module_utils.py b/src/xorl/models/module_utils.py index ad472b01..eead1eed 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -1102,6 +1102,7 @@ def rank0_load_and_broadcast_weights( is_broadcast=True, weights_path=weights_path, ) + skip_key_fn = handler.get_skip_key_fn() if handler is not None else None # When PP is enabled, gather expected keys from all ranks so we can distinguish # "key belongs to another PP stage" (expected) from "truly unexpected key" (warning). @@ -1130,7 +1131,11 @@ def rank0_load_and_broadcast_weights( # background threads while the current shard is being broadcast. # prefetch_count=2 allows 2 concurrent NFS reads for higher aggregate throughput. if global_rank == 0: - prefetched = _prefetch_shards(state_dict_iterators, prefetch_count=2) + if skip_key_fn is not None: + logger.info_rank0("Filtered broadcast loading enabled on rank 0 for handler-skipped checkpoint keys") + prefetched = _prefetch_shards_filtered(state_dict_iterators, skip_key_fn, prefetch_count=2) + else: + prefetched = _prefetch_shards(state_dict_iterators, prefetch_count=2) if global_rank == 0: shard_range = tqdm( @@ -1149,7 +1154,13 @@ def rank0_load_and_broadcast_weights( # all tensors through the checkpoint handler. While this runs, the # background thread is already loading the *next* shard from disk. if global_rank == 0: - state_dict = next(prefetched) + skipped_keys = [] + if skip_key_fn is not None: + state_dict, skipped_keys = next(prefetched) + for skipped_key in skipped_keys: + merged_queue.extend(handler.on_skip_weight(skipped_key)) + else: + state_dict = next(prefetched) for key, tensor in state_dict.items(): key = _convert_weight_key(key, model) if handler is not None: diff --git a/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py b/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py index 0ee076f6..f70b85f8 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py @@ -17,6 +17,7 @@ ExpertWeightBuffer, GateUpMergeBuffer, QKVMergeBuffer, + parse_expert_full_key, parse_expert_key, ) from ..qwen3_5_shared import ( @@ -64,13 +65,14 @@ def __init__( checkpoint_has_per_expert: bool = True, skip_qkv_merge: bool = False, skip_gate_up_merge: bool = False, + skip_expert_loading: bool = False, is_prequantized: bool = False, exclude_modules: Optional[Set[str]] = None, ): self._expert_buffer: Optional[ExpertWeightBuffer] = None # Disable expert buffer when pre-quantized: expert weights are loaded # directly by QLoRAMoeExperts, not merged through the checkpoint handler. - if checkpoint_has_per_expert and not is_prequantized: + if checkpoint_has_per_expert and not is_prequantized and not skip_expert_loading: self._expert_buffer = ExpertWeightBuffer(num_experts, ep_rank=ep_rank, ep_size=ep_size) self._qkv_buffer: Optional[QKVMergeBuffer] = None if not skip_qkv_merge: @@ -83,6 +85,7 @@ def __init__( self._linear_key_dim = linear_key_dim self._linear_value_dim = linear_value_dim self._is_prequantized = is_prequantized + self._skip_expert_loading = skip_expert_loading self._exclude_modules = exclude_modules or set() self._ep_rank = ep_rank self._ep_size = ep_size @@ -213,15 +216,29 @@ def get_skip_key_fn(self) -> Optional[Callable[[str], bool]]: self._expert_buffer.expert_start == 0 and self._expert_buffer.expert_end == self._expert_buffer.num_experts ) - if not has_ep_filter and not self._is_prequantized: + if not has_ep_filter and not self._is_prequantized and not self._skip_expert_loading: return None ep_start = self._expert_buffer.expert_start if has_ep_filter else 0 ep_end = self._expert_buffer.expert_end if has_ep_filter else 0 is_prequantized = self._is_prequantized + skip_expert_loading = self._skip_expert_loading exclude_modules = self._exclude_modules def _should_skip(key: str) -> bool: + if skip_expert_loading: + if parse_expert_full_key(key) is not None: + return True + if self._STACKED_EXPERT_SPLIT_PATTERN.match(key) is not None: + return True + if self._HF_FUSED_EXPERT_GATE_UP_PATTERN.match(key) is not None: + return True + if self._HF_FUSED_EXPERT_DOWN_PATTERN.match(key) is not None: + return True + if self._INTERNAL_FUSED_EXPERT_GATE_UP_PATTERN.match(key) is not None: + return True + if self._INTERNAL_FUSED_EXPERT_DOWN_PATTERN.match(key) is not None: + return True if is_prequantized: # Don't skip keys belonging to excluded modules — they're bf16, load normally. if exclude_modules: @@ -267,6 +284,20 @@ def _is_excluded_module(self, key: str) -> bool: return is_excluded_module_key(key, self._exclude_modules) def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: + if self._skip_expert_loading: + if parse_expert_full_key(key) is not None: + return [] + if self._STACKED_EXPERT_SPLIT_PATTERN.match(key) is not None: + return [] + if self._HF_FUSED_EXPERT_GATE_UP_PATTERN.match(key) is not None: + return [] + if self._HF_FUSED_EXPERT_DOWN_PATTERN.match(key) is not None: + return [] + if self._INTERNAL_FUSED_EXPERT_GATE_UP_PATTERN.match(key) is not None: + return [] + if self._INTERNAL_FUSED_EXPERT_DOWN_PATTERN.match(key) is not None: + return [] + # 0. Pre-quantized: skip quantized auxiliary keys and quantized weight keys # that weren't caught by get_skip_key_fn (e.g., when skip_key_fn wasn't used) # Excluded modules pass through as bf16 — don't drop them. diff --git a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py index a6d2b90f..0c6a8ab0 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -392,6 +392,14 @@ def get_checkpoint_handler(self, **kwargs): ep_rank, ep_size = 0, 1 unfused = getattr(self, "_unfused_for_tp", False) head_dim = getattr(self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads) + skip_expert_loading = False + if not is_prequantized: + from xorl.qlora.modules.moe_experts import QLoRAMoeExperts + + skip_expert_loading = any( + isinstance(module, QLoRAMoeExperts) and not getattr(module, "_weights_loaded", False) + for module in self.modules() + ) return Qwen3_5MoeCheckpointHandler( num_experts=self.config.num_experts, num_attention_heads=self.config.num_attention_heads, @@ -404,6 +412,7 @@ def get_checkpoint_handler(self, **kwargs): checkpoint_has_per_expert=has_per_expert, skip_qkv_merge=True, skip_gate_up_merge=unfused, + skip_expert_loading=skip_expert_loading, is_prequantized=is_prequantized, exclude_modules=exclude_modules, ) diff --git a/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py b/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py index a48eb537..63ae7320 100644 --- a/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py @@ -66,6 +66,7 @@ def __init__( checkpoint_has_per_expert: bool = True, skip_qkv_merge: bool = False, skip_gate_up_merge: bool = False, + skip_expert_loading: bool = False, is_prequantized: bool = False, exclude_modules: Optional[Set[str]] = None, device: Optional["torch.device"] = None, @@ -73,7 +74,7 @@ def __init__( ): self._expert_buffer: Optional[ExpertWeightBuffer] = None # Non-QLoRA expert stacking (disabled when pre-quantized — use QLoRAExpertBuffer instead) - if checkpoint_has_per_expert and not is_prequantized: + if checkpoint_has_per_expert and not is_prequantized and not skip_expert_loading: self._expert_buffer = ExpertWeightBuffer( num_experts, ep_rank=ep_rank, @@ -95,6 +96,7 @@ def __init__( self._expert_start = ep_rank * self._local_num_experts self._expert_end = self._expert_start + self._local_num_experts self._is_prequantized = is_prequantized + self._skip_expert_loading = skip_expert_loading self._exclude_modules = exclude_modules or set() # Inline QLoRA weight loading for dense modules (attention, shared expert) self._qlora_buffer: Optional[QLoRAWeightBuffer] = None @@ -209,12 +211,18 @@ def get_skip_key_fn(self) -> Optional[Callable[[str], bool]]: and self._qlora_expert_buffer.expert_end == self._qlora_expert_buffer._num_experts ) - if not has_ep_filter and not has_expert_ep_filter and not self._is_prequantized: + if ( + not has_ep_filter + and not has_expert_ep_filter + and not self._is_prequantized + and not self._skip_expert_loading + ): return None ep_start = self._expert_buffer.expert_start if has_ep_filter else 0 ep_end = self._expert_buffer.expert_end if has_ep_filter else 0 is_prequantized = self._is_prequantized + skip_expert_loading = self._skip_expert_loading exclude_modules = self._exclude_modules has_qlora_buffer = self._qlora_buffer is not None has_qlora_expert_buffer = self._qlora_expert_buffer is not None @@ -223,6 +231,15 @@ def get_skip_key_fn(self) -> Optional[Callable[[str], bool]]: qe_end = self._qlora_expert_buffer.expert_end def _should_skip(key: str) -> bool: + if skip_expert_loading: + if parse_expert_full_key(key) is not None: + return True + if self._STACKED_EXPERT_SPLIT_PATTERN.match(key) is not None: + return True + if self._FUSED_EXPERT_GATE_UP_PATTERN.match(key) is not None: + return True + if self._FUSED_EXPERT_DOWN_PATTERN.match(key) is not None: + return True if is_prequantized: # Don't skip keys belonging to excluded modules — they're bf16, load normally. if exclude_modules: @@ -281,6 +298,16 @@ def _is_excluded_module(self, key: str) -> bool: return module_short_name in self._exclude_modules def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: + if self._skip_expert_loading: + if parse_expert_full_key(key) is not None: + return [] + if self._STACKED_EXPERT_SPLIT_PATTERN.match(key) is not None: + return [] + if self._FUSED_EXPERT_GATE_UP_PATTERN.match(key) is not None: + return [] + if self._FUSED_EXPERT_DOWN_PATTERN.match(key) is not None: + return [] + # Drop input_scale (unused by our quantization) if key.endswith(".input_scale"): return [] diff --git a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py index 84cf861f..889b5c72 100644 --- a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py +++ b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py @@ -354,6 +354,14 @@ def get_checkpoint_handler(self, **kwargs): # When unfused for TP, skip QKV and gate/up merging — checkpoint keys # already match the model's parameter names. Expert merging is still needed. unfused = getattr(self, "_unfused_for_tp", False) + skip_expert_loading = False + if not is_prequantized: + from xorl.qlora.modules.moe_experts import QLoRAMoeExperts + + skip_expert_loading = any( + isinstance(module, QLoRAMoeExperts) and not getattr(module, "_weights_loaded", False) + for module in self.modules() + ) head_dim = getattr(self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads) return Qwen3MoeCheckpointHandler( @@ -366,6 +374,7 @@ def get_checkpoint_handler(self, **kwargs): checkpoint_has_per_expert=has_per_expert, skip_qkv_merge=unfused, skip_gate_up_merge=unfused, + skip_expert_loading=skip_expert_loading, is_prequantized=is_prequantized, exclude_modules=exclude_modules, device=kwargs.get("device"), diff --git a/tests/models/test_module_utils_broadcast.py b/tests/models/test_module_utils_broadcast.py index e9d440bc..bb4681c5 100644 --- a/tests/models/test_module_utils_broadcast.py +++ b/tests/models/test_module_utils_broadcast.py @@ -1,6 +1,7 @@ from types import SimpleNamespace import pytest +import torch from xorl.models import module_utils @@ -22,18 +23,20 @@ def to_empty(self, device): def test_rank0_broadcast_path_calls_load_state_dict_on_nonzero_ranks(monkeypatch): calls = [] - def fake_broadcast_object_list(obj, src=0): + def fake_broadcast_object_list(obj, src=0, group=None, device=None): if obj[0] is None: obj[0] = [] fake_dist = SimpleNamespace( is_available=lambda: True, is_initialized=lambda: True, - broadcast=lambda tensor, src=0: None, + broadcast=lambda tensor, src=0, group=None: None, broadcast_object_list=fake_broadcast_object_list, ) monkeypatch.setattr(module_utils, "dist", fake_dist) + monkeypatch.setattr(module_utils, "_get_weight_load_group", lambda: None) + monkeypatch.setattr(module_utils, "_get_weight_load_object_device", lambda: None) monkeypatch.setattr( module_utils, "get_parallel_state", @@ -54,10 +57,96 @@ def fake_load_state_dict(weights_path): assert calls == ["dummy-weights"] +def test_rank0_broadcast_path_uses_filtered_prefetch_for_handler_skips(monkeypatch): + batch_meta_calls = [] + dispatched = [] + handler_calls = {"loaded": [], "skipped": []} + + class _Handler: + def get_skip_key_fn(self): + return lambda key: key == "skip.expert.weight" + + def on_skip_weight(self, key): + handler_calls["skipped"].append(key) + return [("merged.skip.weight", torch.tensor([1.0]))] + + def on_load_weight(self, key, tensor): + handler_calls["loaded"].append(key) + return [(key, tensor)] + + def on_load_complete(self): + return [] + + class _BroadcastModel: + def named_buffers(self): + return [] + + def named_parameters(self): + return [ + ("merged.skip.weight", None), + ("keep.weight", None), + ] + + def to_empty(self, device): + self.device = device + + def get_checkpoint_handler(self, **kwargs): + return _Handler() + + def fake_broadcast_object_list(obj, src=0, group=None, device=None): + batch_meta_calls.append(obj[0]) + + fake_dist = SimpleNamespace( + is_available=lambda: True, + is_initialized=lambda: True, + broadcast=lambda tensor, src=0, group=None: None, + broadcast_object_list=fake_broadcast_object_list, + ) + + def fake_prefetch_filtered(state_dict_iterators, skip_key_fn, prefetch_count): + assert state_dict_iterators == ["shard-0"] + assert prefetch_count == 2 + assert skip_key_fn("skip.expert.weight") + yield ({"keep.weight": torch.tensor([2.0])}, ["skip.expert.weight"]) + + def fail_prefetch(*args, **kwargs): + raise AssertionError("_prefetch_shards should not be used when handler skip filtering is available") + + monkeypatch.setattr(module_utils, "dist", fake_dist) + monkeypatch.setattr(module_utils, "_get_weight_load_group", lambda: None) + monkeypatch.setattr(module_utils, "_get_weight_load_object_device", lambda: None) + monkeypatch.setattr( + module_utils, + "get_parallel_state", + lambda: SimpleNamespace(global_rank=0, pp_enabled=False), + ) + monkeypatch.setattr(module_utils, "_build_compiled_key_map", lambda *args, **kwargs: {}) + monkeypatch.setattr(module_utils, "_shrink_expert_params_for_ep", lambda model: None) + monkeypatch.setattr( + module_utils, "_get_checkpoint_keys", lambda weights_path: {"skip.expert.weight", "keep.weight"} + ) + monkeypatch.setattr(module_utils, "_load_state_dict", lambda weights_path: ["shard-0"]) + monkeypatch.setattr(module_utils, "_prefetch_shards_filtered", fake_prefetch_filtered) + monkeypatch.setattr(module_utils, "_prefetch_shards", fail_prefetch) + monkeypatch.setattr(module_utils, "_dispatch_parameter", lambda *args, **kwargs: dispatched.append(args[1])) + monkeypatch.setattr(module_utils, "post_process_after_weight_loading", lambda *args, **kwargs: None) + monkeypatch.setattr(module_utils, "empty_cache", lambda: None) + + module_utils.rank0_load_and_broadcast_weights(_BroadcastModel(), "dummy-weights", init_device="cpu") + + assert handler_calls["skipped"] == ["skip.expert.weight"] + assert handler_calls["loaded"] == ["keep.weight"] + assert dispatched == ["merged.skip.weight", "keep.weight"] + assert batch_meta_calls[0] == [ + ("merged.skip.weight", torch.Size([1]), torch.float32, "broadcast"), + ("keep.weight", torch.Size([1]), torch.float32, "broadcast"), + ] + + def test_try_load_state_dict_uses_rank0_for_local_resolution(monkeypatch): local_resolution_calls = [] - def fake_broadcast_object_list(obj, src=0, group=None): + def fake_broadcast_object_list(obj, src=0, group=None, device=None): assert src == 0 obj[0] = ["shard-0.safetensors", "shard-1.safetensors"] @@ -73,6 +162,8 @@ def fake_try_load_state_dict_local(weights_path, **kwargs): return [module_utils.StateDictIterator("local-only.safetensors")] monkeypatch.setattr(module_utils, "dist", fake_dist) + monkeypatch.setattr(module_utils, "_get_weight_load_group", lambda: None) + monkeypatch.setattr(module_utils, "_get_weight_load_object_device", lambda: None) monkeypatch.setattr(module_utils, "_get_cpu_group", lambda: object()) monkeypatch.setattr(module_utils, "_try_load_state_dict_local", fake_try_load_state_dict_local) diff --git a/tests/models/test_moe_fused_gate_up_proj.py b/tests/models/test_moe_fused_gate_up_proj.py index 9e0bd0a2..c996d57e 100644 --- a/tests/models/test_moe_fused_gate_up_proj.py +++ b/tests/models/test_moe_fused_gate_up_proj.py @@ -132,3 +132,47 @@ def test_qwen3_5_moe_checkpoint_handler_round_trips_fused_experts(): torch.testing.assert_close(saved["model.layers.0.mlp.experts.1.up_proj.weight"], up[1]) torch.testing.assert_close(saved["model.layers.0.mlp.experts.0.down_proj.weight"], down_weight[0]) torch.testing.assert_close(saved["model.layers.0.mlp.experts.1.down_proj.weight"], down_weight[1]) + + +def test_qwen3_moe_checkpoint_handler_skips_deferred_qlora_expert_loading(): + handler = Qwen3MoeCheckpointHandler( + num_experts=2, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=2, + skip_expert_loading=True, + ) + + skip_fn = handler.get_skip_key_fn() + assert skip_fn is not None + assert skip_fn("model.layers.0.mlp.experts.0.gate_proj.weight") + assert skip_fn("model.layers.0.mlp.experts.0.up_proj.weight") + assert skip_fn("model.layers.0.mlp.experts.0.down_proj.weight") + assert skip_fn("model.layers.0.mlp.experts.gate_up_proj") + assert skip_fn("model.layers.0.mlp.experts.down_proj") + + assert handler.on_load_weight("model.layers.0.mlp.experts.0.gate_proj.weight", torch.randn(3, 4)) == [] + assert handler.on_load_weight("model.layers.0.mlp.experts.gate_up_proj", torch.randn(2, 4, 6)) == [] + + +def test_qwen3_5_moe_checkpoint_handler_skips_deferred_qlora_expert_loading(): + handler = Qwen3_5MoeCheckpointHandler( + num_experts=2, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=2, + linear_key_dim=2, + linear_value_dim=2, + skip_expert_loading=True, + ) + + skip_fn = handler.get_skip_key_fn() + assert skip_fn is not None + assert skip_fn("model.layers.0.mlp.experts.0.gate_proj.weight") + assert skip_fn("model.layers.0.mlp.experts.0.up_proj.weight") + assert skip_fn("model.layers.0.mlp.experts.0.down_proj.weight") + assert skip_fn("model.layers.0.mlp.experts.gate_up_proj.weight") + assert skip_fn("model.layers.0.mlp.experts.down_proj.weight") + + assert handler.on_load_weight("model.layers.0.mlp.experts.0.gate_proj.weight", torch.randn(3, 4)) == [] + assert handler.on_load_weight("model.layers.0.mlp.experts.gate_up_proj.weight", torch.randn(2, 6, 4)) == [] From 814861269e7884c319ad6497afa02a92807ca21a Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Mon, 6 Apr 2026 22:58:27 -0700 Subject: [PATCH 15/41] Add SignSGD optimizer support (#86) * Add SignSGD optimizer support * Remove import churn from SignSGD PR * Trim unrelated SignSGD diff (cherry picked from commit 2607fc6db06cfa09244610562ef35ae1c885c555) --- .../content/docs/config-reference/local.md | 2 +- .../content/docs/config-reference/server.md | 2 +- src/xorl/arguments.py | 7 +- src/xorl/optim/__init__.py | 4 +- src/xorl/optim/optimizer.py | 61 ++++++++++- src/xorl/server/server_arguments.py | 7 +- tests/optim/test_signsgd.py | 103 ++++++++++++++++++ tests/server/test_server_arguments.py | 30 +++++ tests/test_arguments.py | 43 ++++++++ 9 files changed, 247 insertions(+), 12 deletions(-) create mode 100644 tests/optim/test_signsgd.py create mode 100644 tests/server/test_server_arguments.py create mode 100644 tests/test_arguments.py diff --git a/docs/src/content/docs/config-reference/local.md b/docs/src/content/docs/config-reference/local.md index 65f8f309..7314047b 100644 --- a/docs/src/content/docs/config-reference/local.md +++ b/docs/src/content/docs/config-reference/local.md @@ -126,7 +126,7 @@ Each entry in `datasets` (or `test_datasets`) is a dict: | Field | Default | Description | |---|---|---| -| `optimizer` | `adamw` | Optimizer: `adamw`, `anyprecision_adamw`, `sgd`, `muon`. | +| `optimizer` | `adamw` | Optimizer: `adamw`, `anyprecision_adamw`, `sgd`, `signsgd`, `muon`. | | `optimizer_dtype` | `bf16` | Dtype for optimizer states in `anyprecision_adamw` and `muon`: `fp32` or `bf16`. BF16 halves optimizer memory. | | `lr` | `5e-5` | Peak learning rate. | | `lr_min` | `1e-7` | Minimum learning rate at the end of decay. | diff --git a/docs/src/content/docs/config-reference/server.md b/docs/src/content/docs/config-reference/server.md index d24655b7..c4d41f7d 100644 --- a/docs/src/content/docs/config-reference/server.md +++ b/docs/src/content/docs/config-reference/server.md @@ -95,7 +95,7 @@ These flags align the training model's numerics with the inference engine (SGLan | Field | Default | Description | |---|---|---| -| `optimizer` | `adamw` | Optimizer: `adamw`, `anyprecision_adamw`, `sgd`, `muon`. | +| `optimizer` | `adamw` | Optimizer: `adamw`, `anyprecision_adamw`, `sgd`, `signsgd`, `muon`. | | `optimizer_dtype` | `bf16` | Dtype for optimizer states: `fp32` or `bf16`. BF16 halves optimizer memory. | | `muon_lr` | `0.02` | Learning rate for Muon matrix parameter groups. Only used when `optimizer: muon`. | | `muon_momentum` | `0.95` | Muon momentum coefficient. | diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 149cd2b8..1a7e4280 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -542,9 +542,12 @@ class TrainingArguments: metadata={"help": "Parameters without weight decay, for example, bias."}, ) - optimizer: Literal["adamw", "anyprecision_adamw", "sgd", "muon"] = field( + optimizer: Literal["adamw", "anyprecision_adamw", "sgd", "signsgd", "muon"] = field( default="adamw", - metadata={"help": "Optimizer type. 'muon' uses Newton-Schulz orthogonalization for 2D+ weight matrices."}, + metadata={ + "help": "Optimizer type. 'signsgd' is a state-free sign update; 'muon' uses " + "Newton-Schulz orthogonalization for 2D+ weight matrices." + }, ) optimizer_dtype: Literal["fp32", "bf16"] = field( default="bf16", diff --git a/src/xorl/optim/__init__.py b/src/xorl/optim/__init__.py index b8c4f12a..1b7f3111 100644 --- a/src/xorl/optim/__init__.py +++ b/src/xorl/optim/__init__.py @@ -1,6 +1,6 @@ from .lr_scheduler import build_lr_scheduler from .muon import Muon -from .optimizer import build_optimizer +from .optimizer import SignSGD, build_optimizer -__all__ = ["build_lr_scheduler", "build_optimizer", "Muon"] +__all__ = ["build_lr_scheduler", "build_optimizer", "Muon", "SignSGD"] diff --git a/src/xorl/optim/optimizer.py b/src/xorl/optim/optimizer.py index 857c1ad0..af14733a 100644 --- a/src/xorl/optim/optimizer.py +++ b/src/xorl/optim/optimizer.py @@ -29,6 +29,53 @@ # https://github.com/meta-llama/llama-recipes/blob/v0.0.4/src/llama_recipes/policies/anyprecision_optimizer.py +class SignSGD(Optimizer): + """Sign-based SGD optimizer with no optimizer-state tensors.""" + + def __init__( + self, + params, + lr: float = 1e-3, + weight_decay: float = 0.0, + ): + if lr < 0.0: + raise ValueError(f"Invalid learning rate: {lr}") + if weight_decay < 0.0: + raise ValueError(f"Invalid weight_decay: {weight_decay}") + + defaults = { + "lr": lr, + "weight_decay": weight_decay, + } + super().__init__(params, defaults) + + @torch.no_grad() + def step(self, closure=None): + """Perform a single optimization step.""" + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + lr = group["lr"] + weight_decay = group["weight_decay"] + + for p in group["params"]: + grad = p.grad + if grad is None: + continue + if grad.is_sparse: + raise RuntimeError("SignSGD does not support sparse gradients.") + + if weight_decay: + p.add_(p, alpha=-lr * weight_decay) + + p.add_(torch.sign(grad), alpha=-lr) + + return loss + + class AnyPrecisionAdamW(Optimizer): def __init__( self, @@ -223,6 +270,7 @@ def _should_build_ep_aware(model: "nn.Module", param_groups: Optional[Sequence[D # Only auto-split when using FSDP2 with EP and no explicit param_groups if param_groups is not None: return False + ps = get_parallel_state() if ps.dp_mode != "fsdp2" or not ps.ep_enabled: return False @@ -333,6 +381,9 @@ def _get_optimizer_cls_and_kwargs( sgd_defaults.update(kwargs) ctor_kwargs = dict(lr=lr, weight_decay=weight_decay, **sgd_defaults) return torch.optim.SGD, ctor_kwargs + elif optimizer_type == "signsgd": + ctor_kwargs = dict(lr=lr, weight_decay=weight_decay, **kwargs) + return SignSGD, ctor_kwargs elif optimizer_type == "muon": adamw_state_dtype = _ANYPRECISION_STATE_DTYPES.get(optimizer_dtype) ctor_kwargs = dict( @@ -350,7 +401,7 @@ def _get_optimizer_cls_and_kwargs( return Muon, ctor_kwargs else: raise ValueError( - f"Unsupported optimizer type: '{optimizer_type}'. Supported: adamw, anyprecision_adamw, sgd, muon." + f"Unsupported optimizer type: '{optimizer_type}'. Supported: adamw, anyprecision_adamw, sgd, signsgd, muon." ) @@ -372,6 +423,7 @@ def _create_optimizer( Common args (lr, betas, eps, weight_decay) are passed directly. Optimizer-specific args are passed via optimizer_kwargs: - sgd: {"momentum": 0.9, "nesterov": True} + - signsgd: no optimizer-specific kwargs - muon: {"muon_lr": 0.02, "muon_momentum": 0.95, ...} - adamw/anyprecision_adamw: any extra kwargs forwarded to constructor """ @@ -492,13 +544,14 @@ def build_optimizer( eps: AdamW epsilon. weight_decay: Weight decay coefficient. fused: Use fused AdamW kernel (mutually exclusive with foreach). - optimizer_type: One of "adamw", "anyprecision_adamw", "sgd", "muon". + optimizer_type: One of "adamw", "anyprecision_adamw", "sgd", "signsgd", "muon". optimizer_dtype: State dtype for anyprecision_adamw / muon ("fp32" or "bf16"). param_groups: Custom param groups. If None, auto-built with weight decay splitting. no_decay_modules: Module class names to exclude from weight decay. no_decay_params: Parameter name patterns to exclude from weight decay. optimizer_kwargs: Optimizer-specific keyword arguments passed to the constructor. - sgd: {"momentum": 0.9, "nesterov": True} + - signsgd: no optimizer-specific kwargs - muon: {"muon_lr": 0.02, "muon_momentum": 0.95, "muon_nesterov": True, "muon_ns_steps": 5, "muon_adjust_lr_fn": None, "muon_momentum_dtype": None} - adamw/anyprecision_adamw: any extra kwargs forwarded to constructor @@ -618,8 +671,8 @@ def build_ep_fsdp2_optimizer( if optimizer_type == "muon": # For Muon + EP: classify within each EP/non-EP subset, then build Muon optimizers - ep_param_set = set(id(p) for p in ep_params) - non_ep_param_set = set(id(p) for p in non_ep_params) + ep_param_set = {id(p) for p in ep_params} + non_ep_param_set = {id(p) for p in non_ep_params} # Classify all model params into Muon vs AdamW all_muon, _, all_adamw, _ = _classify_muon_params(model) diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index f6ab2ec0..ab2c916d 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -232,9 +232,12 @@ class ServerArguments: # Optimizer # ======================================================================== - optimizer: Literal["adamw", "anyprecision_adamw", "sgd", "muon"] = field( + optimizer: Literal["adamw", "anyprecision_adamw", "sgd", "signsgd", "muon"] = field( default="adamw", - metadata={"help": "Optimizer type. 'muon' uses Newton-Schulz orthogonalization for 2D+ weight matrices."}, + metadata={ + "help": "Optimizer type. 'signsgd' is a state-free sign update; 'muon' uses " + "Newton-Schulz orthogonalization for 2D+ weight matrices." + }, ) optimizer_dtype: Literal["fp32", "bf16"] = field( diff --git a/tests/optim/test_signsgd.py b/tests/optim/test_signsgd.py new file mode 100644 index 00000000..421c88a6 --- /dev/null +++ b/tests/optim/test_signsgd.py @@ -0,0 +1,103 @@ +import pytest +import torch +import torch.nn as nn + +from xorl.optim import SignSGD, build_optimizer + + +pytestmark = [pytest.mark.cpu] + + +class TinyModule(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(3, 2) + + +def test_signsgd_applies_sign_updates(): + param = nn.Parameter(torch.tensor([1.0, -2.0, 3.0])) + optimizer = SignSGD([param], lr=0.1) + + param.grad = torch.tensor([2.0, -0.5, 0.0]) + optimizer.step() + + expected = torch.tensor([0.9, -1.9, 3.0]) + assert torch.allclose(param, expected) + + +def test_signsgd_applies_decoupled_weight_decay_before_sign_update(): + param = nn.Parameter(torch.tensor([2.0, -2.0])) + optimizer = SignSGD([param], lr=0.1, weight_decay=0.5) + + param.grad = torch.tensor([3.0, -4.0]) + optimizer.step() + + expected = torch.tensor([1.8, -1.8]) + assert torch.allclose(param, expected) + + +def test_signsgd_skips_parameters_without_gradients(): + param = nn.Parameter(torch.tensor([1.5, -1.5])) + optimizer = SignSGD([param], lr=0.1, weight_decay=0.5) + + optimizer.step() + + assert torch.allclose(param, torch.tensor([1.5, -1.5])) + + +def test_signsgd_rejects_sparse_gradients(): + param = nn.Parameter(torch.ones(4)) + optimizer = SignSGD([param], lr=0.1) + param.grad = torch.sparse_coo_tensor(indices=[[0, 2]], values=torch.tensor([1.0, -1.0]), size=(4,)) + + with pytest.raises(RuntimeError, match="does not support sparse gradients"): + optimizer.step() + + +def test_signsgd_keeps_optimizer_state_empty_across_steps(): + param = nn.Parameter(torch.tensor([1.0])) + optimizer = SignSGD([param], lr=0.1) + + for grad in (torch.tensor([1.0]), torch.tensor([-1.0])): + param.grad = grad + optimizer.step() + assert len(optimizer.state) == 0 + + +def test_signsgd_state_dict_round_trips_without_state_tensors(): + source_param = nn.Parameter(torch.tensor([1.0, -1.0])) + source_optimizer = SignSGD([source_param], lr=0.1, weight_decay=0.25) + source_param.grad = torch.tensor([1.0, -1.0]) + source_optimizer.step() + + state_dict = source_optimizer.state_dict() + + target_param = nn.Parameter(torch.tensor([0.0, 0.0])) + target_optimizer = SignSGD([target_param], lr=1.0, weight_decay=0.0) + target_optimizer.load_state_dict(state_dict) + + assert state_dict["state"] == {} + assert target_optimizer.state_dict()["state"] == {} + assert target_optimizer.param_groups[0]["lr"] == pytest.approx(0.1) + assert target_optimizer.param_groups[0]["weight_decay"] == pytest.approx(0.25) + + +def test_build_optimizer_supports_signsgd_and_preserves_weight_decay_split(): + model = TinyModule() + + optimizer = build_optimizer( + model, + lr=0.1, + weight_decay=0.01, + optimizer_type="signsgd", + no_decay_params=["bias"], + ) + + assert isinstance(optimizer, SignSGD) + assert len(optimizer.param_groups) == 2 + + decay_group, no_decay_group = optimizer.param_groups + assert decay_group["weight_decay"] == pytest.approx(0.01) + assert no_decay_group["weight_decay"] == pytest.approx(0.0) + assert decay_group["params"] == [model.linear.weight] + assert no_decay_group["params"] == [model.linear.bias] diff --git a/tests/server/test_server_arguments.py b/tests/server/test_server_arguments.py new file mode 100644 index 00000000..9ccedd25 --- /dev/null +++ b/tests/server/test_server_arguments.py @@ -0,0 +1,30 @@ +import pytest +import yaml + +from xorl.server.launcher import load_server_arguments + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +def test_load_server_arguments_threads_signsgd_through_nested_config(tmp_path): + config_path = tmp_path / "server_config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model": { + "model_path": "Qwen/Qwen3-8B", + }, + "train": { + "optimizer": "signsgd", + "output_dir": str(tmp_path / "outputs"), + }, + } + ), + encoding="utf-8", + ) + + args = load_server_arguments(str(config_path)) + + assert args.optimizer == "signsgd" + assert args.to_config_dict()["train"]["optimizer"] == "signsgd" diff --git a/tests/test_arguments.py b/tests/test_arguments.py new file mode 100644 index 00000000..11728439 --- /dev/null +++ b/tests/test_arguments.py @@ -0,0 +1,43 @@ +import sys + +import pytest +import yaml + +from xorl.arguments import Arguments, parse_args + + +pytestmark = [pytest.mark.cpu] + + +def test_parse_args_accepts_signsgd_from_yaml(tmp_path, monkeypatch): + config_path = tmp_path / "config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model": { + "model_path": "Qwen/Qwen3-8B", + }, + "data": { + "datasets": [{"path": "dummy", "type": "tokenized"}], + }, + "train": { + "init_device": "meta", + "output_dir": str(tmp_path / "outputs"), + "optimizer": "signsgd", + "use_wandb": False, + }, + } + ), + encoding="utf-8", + ) + + monkeypatch.setenv("WORLD_SIZE", "1") + monkeypatch.setenv("LOCAL_WORLD_SIZE", "1") + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("LOCAL_RANK", "0") + monkeypatch.setattr(sys, "argv", ["train.py", str(config_path)]) + + args = parse_args(Arguments) + + assert args.train.optimizer == "signsgd" + assert args.train.optimizer_kwargs == {} From 1648ebf87dca3d21694260b310c86e7c7a5dfc67 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Tue, 7 Apr 2026 13:06:16 -0700 Subject: [PATCH 16/41] [Feat] Move MoE routing weights before down_proj (#87) * Move MoE routing weights before down_proj * Fix EP routing score dispatch and gradients * Apply Ruff formatting for EP routing score changes (cherry picked from commit 67a2c921544fee18534343e234985a63b343657e) --- src/xorl/distributed/moe/alltoall.py | 64 +++++-- src/xorl/distributed/moe/deepep.py | 24 +-- src/xorl/distributed/moe/utils.py | 17 +- src/xorl/models/layers/moe/backend/native.py | 46 ++++- src/xorl/models/layers/moe/experts.py | 4 + src/xorl/ops/moe/quack.py | 70 +++++-- src/xorl/ops/moe/triton.py | 55 ++++-- tests/distributed/test_deepep_correctness.py | 3 +- .../test_moe_memory_efficient_permutation.py | 149 +++++++++++++++ tests/ops/test_ep_routing_scores.py | 180 ++++++++++++++++++ 10 files changed, 535 insertions(+), 77 deletions(-) create mode 100644 tests/distributed/test_moe_memory_efficient_permutation.py create mode 100644 tests/ops/test_ep_routing_scores.py diff --git a/src/xorl/distributed/moe/alltoall.py b/src/xorl/distributed/moe/alltoall.py index baf871d2..35d1e4fc 100644 --- a/src/xorl/distributed/moe/alltoall.py +++ b/src/xorl/distributed/moe/alltoall.py @@ -6,7 +6,17 @@ import torch.nn.functional as F from .comm import all_to_all -from .utils import generate_weights_idx, permute, sort_chunks_by_idxs, unpermute +from .utils import permute, permuted_weights, sort_chunks_by_idxs, unpermute + + +def _expert_chunk_permute_order(num_experts: int, ep_size: int) -> List[int]: + num_local_experts = num_experts // ep_size + return torch.arange(num_experts).reshape(-1, num_local_experts).T.ravel().tolist() + + +def _expert_chunk_unpermute_order(num_experts: int, ep_size: int) -> List[int]: + num_local_experts = num_experts // ep_size + return torch.arange(num_experts).reshape(num_local_experts, -1).T.ravel().tolist() def preprocess( @@ -80,21 +90,37 @@ def token_pre_all2all( global_permuted_hidden_states = all_to_all(ep_group, local_permuted_hidden_states, output_splits, input_splits) # group tokens together by expert - num_local_experts = num_experts // ep_group.size() - permute_order = torch.arange(num_experts).reshape(-1, num_local_experts).T.ravel().tolist() global_permuted_hidden_states = sort_chunks_by_idxs( global_permuted_hidden_states, num_global_tokens_per_local_expert.ravel(), - permute_order, + _expert_chunk_permute_order(num_experts, ep_group.size()), ) return global_permuted_hidden_states, routing_map, local_input_permutation_mapping, org_hidden_states_shape +def score_pre_all2all( + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, + num_experts: int, + input_splits: torch.Tensor, + output_splits: torch.Tensor, + num_global_tokens_per_local_expert: torch.Tensor, + ep_group: Optional[dist.ProcessGroup] = None, +) -> torch.Tensor: + """Dispatch routing weights into the same expert-sorted order as tokens.""" + local_scores = permuted_weights(routing_weights, selected_experts, num_experts).unsqueeze(-1) + global_scores = all_to_all(ep_group, local_scores, output_splits, input_splits) + global_scores = sort_chunks_by_idxs( + global_scores, + num_global_tokens_per_local_expert.ravel(), + _expert_chunk_permute_order(num_experts, ep_group.size()), + ) + return global_scores.squeeze(-1) + + def tokens_post_all2all( expert_outputs: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: int, num_experts: int, input_splits: torch.Tensor, output_splits: torch.Tensor, @@ -105,25 +131,18 @@ def tokens_post_all2all( ep_group: Optional[dist.ProcessGroup] = None, ) -> torch.Tensor: # group tokens together by expert - num_local_experts = num_experts // ep_group.size() - unpermute_order = torch.arange(num_experts).reshape(num_local_experts, -1).T.ravel().tolist() expert_outputs = sort_chunks_by_idxs( expert_outputs, num_global_tokens_per_local_expert.T.ravel(), - unpermute_order, + _expert_chunk_unpermute_order(num_experts, ep_group.size()), ) unpermute_outputs = all_to_all(ep_group, expert_outputs, input_splits, output_splits) - # [tokens, experts] - weights_idx = generate_weights_idx(routing_weights, selected_experts, num_experts) - unpermute_outputs = unpermute( unpermute_outputs, - weights_idx, org_hidden_states_shape, local_input_permutation_mapping, - routing_map, ) return unpermute_outputs @@ -150,9 +169,8 @@ class AllToAllDispatchContext: num_tokens_per_expert: torch.Tensor routing_map: torch.Tensor perm_mapping: torch.Tensor + expert_scores: torch.Tensor orig_shape: torch.Size - routing_weights: torch.Tensor - selected_experts: torch.Tensor num_experts: int @@ -200,6 +218,15 @@ def alltoall_pre_dispatch( ) cumsum = torch.cumsum(sum_tokens, dim=0).to(permuted_tokens.device) + expert_scores = score_pre_all2all( + routing_weights=routing_weights, + selected_experts=selected_experts, + num_experts=num_experts, + input_splits=input_splits, + output_splits=output_splits, + num_global_tokens_per_local_expert=num_tokens_per_expert, + ep_group=ep_group, + ).to(permuted_tokens.dtype) ctx = AllToAllDispatchContext( input_splits=input_splits, @@ -207,9 +234,8 @@ def alltoall_pre_dispatch( num_tokens_per_expert=num_tokens_per_expert, routing_map=routing_map, perm_mapping=perm_mapping, + expert_scores=expert_scores, orig_shape=orig_shape, - routing_weights=routing_weights, - selected_experts=selected_experts, num_experts=num_experts, ) @@ -236,8 +262,6 @@ def alltoall_post_combine( """ return tokens_post_all2all( expert_outputs=expert_output, - routing_weights=ctx.routing_weights, - selected_experts=ctx.selected_experts, num_experts=ctx.num_experts, input_splits=ctx.input_splits, output_splits=ctx.output_splits, diff --git a/src/xorl/distributed/moe/deepep.py b/src/xorl/distributed/moe/deepep.py index 88b9b3e9..117f5f7c 100644 --- a/src/xorl/distributed/moe/deepep.py +++ b/src/xorl/distributed/moe/deepep.py @@ -352,16 +352,16 @@ def backward(ctx, grad_expert_input, grad_cumsum, grad_dispatch_ctx): class _FusedUnpermuteAndCombine(torch.autograd.Function): - """Fused weighted scatter-add (unpermute) + combine in a single autograd boundary. + """Fused scatter-add (unpermute) + combine in a single autograd boundary. Forward: - 1. Weighted scatter-add: ``output[idx[i]] += score[i] * expert_output[i]`` + 1. Scatter-add: ``output[idx[i]] += expert_output[i]`` 2. ``buffer.combine()`` to send results back to original ranks 3. Optionally defer sync (async_combine) for overlap with next layer Backward: 1. ``buffer.dispatch()`` to reverse combine → grad in recv order - 2. ``grad_expert[i] = score[i] * grad[idx[i]]`` — reverse of scatter-add + 2. ``grad_expert[i] = grad[idx[i]]`` — reverse of scatter-add """ @staticmethod @@ -376,7 +376,7 @@ def forward( hidden_dim = expert_output.shape[1] if expert_output.shape[0] > 0 else dispatch_ctx.hidden_dim device = expert_output.device - # Step 1: Weighted scatter-add (unpermute) + # Step 1: Scatter-add (unpermute) if expert_output.shape[0] == 0: gather_output = torch.zeros( dispatch_ctx.num_recv_tokens, @@ -385,7 +385,6 @@ def forward( device=device, ) else: - weighted = expert_output * dispatch_ctx.permuted_scores.unsqueeze(1).to(expert_output.dtype) gather_output = torch.zeros( dispatch_ctx.num_recv_tokens, hidden_dim, @@ -393,7 +392,7 @@ def forward( device=device, ) idx_2d = dispatch_ctx.permuted_indices.unsqueeze(1).expand(-1, hidden_dim) - gather_output.scatter_add_(0, idx_2d, weighted) + gather_output.scatter_add_(0, idx_2d, expert_output) # Step 2: Combine previous_event = EventOverlap(EventHandle()) @@ -411,9 +410,7 @@ def forward( else: event.current_stream_wait() - # Save for backward — scores don't need grad (not a tensor input), - # so we never save expert_output, saving memory. - ctx.save_for_backward(dispatch_ctx.permuted_scores, dispatch_ctx.permuted_indices) + ctx.save_for_backward(dispatch_ctx.permuted_indices) ctx.buffer = buffer ctx.handle = dispatch_ctx.handle ctx.input_dtype = expert_output.dtype @@ -424,7 +421,7 @@ def backward(ctx, grad_output): if grad_output is None: return None, None, None, None - permuted_scores, permuted_indices = ctx.saved_tensors + (permuted_indices,) = ctx.saved_tensors buffer = ctx.buffer handle = ctx.handle @@ -443,11 +440,8 @@ def backward(ctx, grad_output): if grad_gather.dtype != ctx.input_dtype: grad_gather = grad_gather.to(ctx.input_dtype) - # Step 2: Reverse weighted scatter-add - # Forward: gather_output[idx[i]] += score[i] * expert_output[i] - # Backward: grad_expert[i] = score[i] * grad_gather[idx[i]] - grad_gathered = grad_gather.index_select(0, permuted_indices) - grad_expert_output = grad_gathered * permuted_scores.unsqueeze(1).to(grad_gathered.dtype) + # Step 2: Reverse scatter-add + grad_expert_output = grad_gather.index_select(0, permuted_indices) return grad_expert_output, None, None, None diff --git a/src/xorl/distributed/moe/utils.py b/src/xorl/distributed/moe/utils.py index 8bbe3ad8..57760781 100644 --- a/src/xorl/distributed/moe/utils.py +++ b/src/xorl/distributed/moe/utils.py @@ -28,27 +28,19 @@ def permute(tokens: torch.Tensor, routing_map: torch.Tensor): def unpermute( tokens: torch.Tensor, - routing_weights: torch.Tensor, hidden_states_shape: torch.Size, permutation_mapping: torch.Tensor, - routing_map: torch.Tensor, ): """ - Unpermutes the tokens and apply the weight. + Unpermutes preweighted tokens via scatter-add. Args: tokens (torch.Tensor): The input token tensor, [num_tokens, hidden_dim]. - routing_weights (torch.Tensor): The routing weights, [num_tokens, num_experts]. hidden_states_shape (torch.Size): The shape of the hidden states, [num_tokens, hidden_dim]. - routing_map (torch.Tensor): The sparse token to expert mapping, [num_experts, tokens]. Returns: torch.Tensor: The unpermuted token tensor, [num_tokens, hidden_dim]. """ - tokens_weight = routing_weights.T.contiguous().masked_select(routing_map.bool()) - - # Weight in-place to avoid a second full-sized activation buffer during combine. - tokens.mul_(tokens_weight.unsqueeze(-1)) hidden_dim = hidden_states_shape[-1] unpermuted_tokens = torch.zeros(hidden_states_shape, device=tokens.device, dtype=tokens.dtype) @@ -61,6 +53,13 @@ def unpermute( return unpermuted_tokens +def permuted_weights(routing_weights: torch.Tensor, selected_experts: torch.Tensor, num_experts: int) -> torch.Tensor: + """Return routing weights in the same local expert-sorted order as ``permute()``.""" + dense_weights = generate_weights_idx(routing_weights, selected_experts, num_experts) + routing_map = torch.nn.functional.one_hot(selected_experts, num_classes=num_experts).permute(2, 1, 0).sum(dim=1) + return dense_weights.T.contiguous().masked_select(routing_map.bool()) + + def generate_weights_idx(routing_weights: torch.Tensor, selected_experts: torch.Tensor, num_experts) -> torch.Tensor: """ Generate the weight index for the unpermute operation. diff --git a/src/xorl/models/layers/moe/backend/native.py b/src/xorl/models/layers/moe/backend/native.py index 664dadd8..297c3769 100644 --- a/src/xorl/models/layers/moe/backend/native.py +++ b/src/xorl/models/layers/moe/backend/native.py @@ -127,6 +127,7 @@ def _run_experts_grouped_mm( down_proj: torch.Tensor, x: torch.Tensor, padded_counts: torch.Tensor, + expert_scores: torch.Tensor | None = None, ) -> torch.Tensor: """Run MoE experts using ``torch._grouped_mm``. @@ -159,6 +160,8 @@ def _run_experts_grouped_mm( # SwiGLU: silu(gate) * up h = gate_out * up_out + if expert_scores is not None: + h = h * expert_scores.to(h.dtype).unsqueeze(-1) # down: h @ down_proj -> (tokens, hidden) out = torch._grouped_mm( @@ -210,6 +213,7 @@ def _run_experts_moe_act( down_proj: torch.Tensor, x: torch.Tensor, padded_counts: torch.Tensor, + expert_scores: torch.Tensor | None = None, ) -> torch.Tensor: """Run MoE experts with moe_act: gate+up SwiGLU is checkpointed so gate_output and up_output are recomputed in backward instead of saved. @@ -231,6 +235,8 @@ def _run_experts_moe_act( offsets, use_reentrant=False, ) + if expert_scores is not None: + h = h * expert_scores.to(h.dtype).unsqueeze(-1) out = torch._grouped_mm( h, @@ -310,10 +316,20 @@ def _native_expert_forward_impl( sorted_hidden_padded[pad_dst] = sorted_hidden # 7. Expert compute (backend-specific) - expert_out_padded = compute_fn(gate_proj, up_proj, down_proj, sorted_hidden_padded, padded_counts) + expert_scores_padded = sorted_hidden_padded.new_zeros(total_padded) + expert_scores_padded[pad_dst] = sorted_weights.to(expert_scores_padded.dtype) - # 8. Gather from padded layout (reuse pad_dst) + apply routing weights - expert_out = expert_out_padded[pad_dst] * sorted_weights.unsqueeze(-1) + expert_out_padded = compute_fn( + gate_proj, + up_proj, + down_proj, + sorted_hidden_padded, + padded_counts, + expert_scores_padded, + ) + + # 8. Gather from padded layout (reuse pad_dst) + expert_out = expert_out_padded[pad_dst] # 9. Scatter-add back to original token positions output = hidden_states.new_zeros(num_tokens, hidden_dim) @@ -553,6 +569,7 @@ def native_ep_compute( gate_proj: torch.Tensor, up_proj: torch.Tensor, down_proj: torch.Tensor, + expert_scores: torch.Tensor | None = None, ) -> torch.Tensor: """EP expert compute using ``torch._grouped_mm``. @@ -578,7 +595,17 @@ def native_ep_compute( # Pad for alignment and run compiled grouped GEMM padded_tokens, padded_counts = _pad_to_alignment(permute_tokens, counts, num_local_experts) - out_padded = _run_experts_compiled(gate_proj, up_proj, down_proj, padded_tokens, padded_counts) + expert_scores_padded = None + if expert_scores is not None: + expert_scores_padded = padded_tokens.new_zeros(padded_tokens.shape[0]) + expert_scores_padded[ + _compute_pad_indices( + counts, padded_counts, num_local_experts, permute_tokens.shape[0], padded_tokens.device + ) + ] = expert_scores.to(padded_tokens.dtype) + out_padded = _run_experts_compiled( + gate_proj, up_proj, down_proj, padded_tokens, padded_counts, expert_scores_padded + ) # Unpad back to real token counts return _unpad(out_padded, counts, padded_counts, num_local_experts, permute_tokens.shape[0]) @@ -643,6 +670,7 @@ def native_ep_compute_moe_act( gate_proj: torch.Tensor, up_proj: torch.Tensor, down_proj: torch.Tensor, + expert_scores: torch.Tensor | None = None, ) -> torch.Tensor: """EP expert compute with moe_act using ``torch._grouped_mm``. @@ -657,7 +685,15 @@ def native_ep_compute_moe_act( counts = _cumsum_to_counts(cumsum, num_local_experts) padded_tokens, padded_counts = _pad_to_alignment(permute_tokens, counts, num_local_experts) - out_padded = _run_experts_moe_act(gate_proj, up_proj, down_proj, padded_tokens, padded_counts) + expert_scores_padded = None + if expert_scores is not None: + expert_scores_padded = padded_tokens.new_zeros(padded_tokens.shape[0]) + expert_scores_padded[ + _compute_pad_indices( + counts, padded_counts, num_local_experts, permute_tokens.shape[0], padded_tokens.device + ) + ] = expert_scores.to(padded_tokens.dtype) + out_padded = _run_experts_moe_act(gate_proj, up_proj, down_proj, padded_tokens, padded_counts, expert_scores_padded) return _unpad(out_padded, counts, padded_counts, num_local_experts, permute_tokens.shape[0]) diff --git a/src/xorl/models/layers/moe/experts.py b/src/xorl/models/layers/moe/experts.py index ee3fad9f..2a578dee 100644 --- a/src/xorl/models/layers/moe/experts.py +++ b/src/xorl/models/layers/moe/experts.py @@ -201,12 +201,14 @@ def _ep_forward( torch.cuda.synchronize() # Step 2: Expert computation (backend-specific GEMM only) + expert_scores = getattr(ctx, "expert_scores", getattr(ctx, "permuted_scores", None)) expert_output = compute_fn( permute_tokens, cumsum, gate_proj, up_proj, self.down_proj, + expert_scores, ) # Step 3: Combine expert outputs back to original ranks @@ -233,12 +235,14 @@ def _ep_forward_debug(self, dispatch_fn, combine_fn, compute_fn, dispatch_kwargs # --- compute --- ev[2].record() + expert_scores = getattr(ctx, "expert_scores", getattr(ctx, "permuted_scores", None)) expert_output = compute_fn( permute_tokens, cumsum, self.gate_proj.contiguous(), self.up_proj.contiguous(), self.down_proj, + expert_scores, ) ev[3].record() diff --git a/src/xorl/ops/moe/quack.py b/src/xorl/ops/moe/quack.py index 7128e8f2..83a5dab7 100644 --- a/src/xorl/ops/moe/quack.py +++ b/src/xorl/ops/moe/quack.py @@ -29,9 +29,10 @@ class QuackEPGroupGemmMoeAct(torch.autograd.Function): gate_output and up_output in backward (2 local GEMMs, no EP communication).""" @staticmethod - def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): + def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores=None): max_M = permute_tokens.shape[0] cu_seqlens = cumsum_to_cu_seqlens(cumsum) + ctx.has_expert_scores = expert_scores is not None gate_output = quack_group_gemm_same_nk( a=permute_tokens, b=gate_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=False, cu_seqlens_m=cu_seqlens @@ -42,6 +43,8 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): gate_activation = torch.ops.aten.silu(gate_output) gated_output = gate_activation * up_output + if expert_scores is not None: + gated_output = gated_output * expert_scores.to(gated_output.dtype).unsqueeze(-1) del gate_activation down_output = quack_group_gemm_same_nk( @@ -50,12 +53,14 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): del gated_output # moe_act: save only 5 tensors (drop gate_output, up_output) - ctx.save_for_backward(permute_tokens, cumsum, gate_proj, up_proj, down_proj) + if expert_scores is None: + expert_scores = permute_tokens.new_ones(permute_tokens.shape[0]) + ctx.save_for_backward(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores) return down_output @staticmethod def backward(ctx, grad_output): - permute_tokens, cumsum, gate_proj, up_proj, down_proj = ctx.saved_tensors + permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores = ctx.saved_tensors max_M = grad_output.shape[0] cu_seqlens_m = cumsum_to_cu_seqlens(cumsum) @@ -70,9 +75,12 @@ def backward(ctx, grad_output): # Rest identical to QuackEPGroupGemm.backward gate_activation = torch.ops.aten.silu(gate_output) gated_output = gate_activation * up_output + expert_scores_dtype = expert_scores.dtype + expert_scores = expert_scores.to(gated_output.dtype) + gated_weighted = gated_output * expert_scores.unsqueeze(-1) # dgrad FC2 - grad_gated_output = quack_group_gemm_same_nk( + grad_gated_weighted = quack_group_gemm_same_nk( a=grad_output, b=down_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=True, cu_seqlens_m=cu_seqlens_m ) @@ -81,7 +89,7 @@ def backward(ctx, grad_output): if down_proj.requires_grad: grad_down_proj = torch.empty_like(down_proj) quack_group_gemm_same_mn( - a=gated_output, + a=gated_weighted, b=grad_output, c=grad_down_proj, cumsum_K=cumsum, @@ -90,7 +98,13 @@ def backward(ctx, grad_output): transpose_b=False, cu_seqlens_k=cu_seqlens_m, ) - del gated_output + grad_expert_scores = None + if ctx.has_expert_scores: + grad_expert_scores = (grad_gated_weighted * gated_output).sum(dim=-1).to(expert_scores_dtype) + del gated_output, gated_weighted + + grad_gated_output = grad_gated_weighted * expert_scores.unsqueeze(-1) + del grad_gated_weighted # Activation backward grad_up_output = gate_activation * grad_gated_output @@ -137,7 +151,7 @@ def backward(ctx, grad_output): ) del grad_up_output - return grad_permute_tokens, None, grad_gate_proj, grad_up_proj, grad_down_proj + return grad_permute_tokens, None, grad_gate_proj, grad_up_proj, grad_down_proj, grad_expert_scores class QuackMoeExpertsFunction(torch.autograd.Function): @@ -296,9 +310,10 @@ class QuackEPGroupGemm(torch.autograd.Function): """Memory-optimized EP expert GEMM. Recomputes cheap intermediates, explicit del.""" @staticmethod - def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): + def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores=None): max_M = permute_tokens.shape[0] cu_seqlens = cumsum_to_cu_seqlens(cumsum) + ctx.has_expert_scores = expert_scores is not None if _DEBUG_EP: return QuackEPGroupGemm._forward_debug( @@ -308,6 +323,7 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): gate_proj, up_proj, down_proj, + expert_scores, max_M, cu_seqlens, ) @@ -321,6 +337,8 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): gate_activation = torch.ops.aten.silu(gate_output) gated_output = gate_activation * up_output + if expert_scores is not None: + gated_output = gated_output * expert_scores.to(gated_output.dtype).unsqueeze(-1) del gate_activation down_output = quack_group_gemm_same_nk( @@ -328,14 +346,19 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): ) del gated_output - ctx.save_for_backward(permute_tokens, cumsum, gate_proj, up_proj, down_proj, gate_output, up_output) + if expert_scores is None: + expert_scores = permute_tokens.new_ones(permute_tokens.shape[0]) + ctx.save_for_backward( + permute_tokens, cumsum, gate_proj, up_proj, down_proj, gate_output, up_output, expert_scores + ) return down_output @staticmethod - def _forward_debug(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, max_M, cu_seqlens): + def _forward_debug(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores, max_M, cu_seqlens): """Instrumented forward with per-GEMM CUDA event timing.""" rank = dist.get_rank() if dist.is_initialized() else 0 ev = [torch.cuda.Event(enable_timing=True) for _ in range(8)] + ctx.has_expert_scores = expert_scores is not None ev[0].record() gate_output = quack_group_gemm_same_nk( @@ -350,6 +373,8 @@ def _forward_debug(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, m gate_activation = torch.ops.aten.silu(gate_output) gated_output = gate_activation * up_output + if expert_scores is not None: + gated_output = gated_output * expert_scores.to(gated_output.dtype).unsqueeze(-1) del gate_activation ev[3].record() @@ -377,21 +402,28 @@ def _forward_debug(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, m flush=True, ) - ctx.save_for_backward(permute_tokens, cumsum, gate_proj, up_proj, down_proj, gate_output, up_output) + if expert_scores is None: + expert_scores = permute_tokens.new_ones(permute_tokens.shape[0]) + ctx.save_for_backward( + permute_tokens, cumsum, gate_proj, up_proj, down_proj, gate_output, up_output, expert_scores + ) return down_output @staticmethod def backward(ctx, grad_output): - permute_tokens, cumsum, gate_proj, up_proj, down_proj, gate_output, up_output = ctx.saved_tensors + permute_tokens, cumsum, gate_proj, up_proj, down_proj, gate_output, up_output, expert_scores = ctx.saved_tensors max_M = grad_output.shape[0] cu_seqlens_m = cumsum_to_cu_seqlens(cumsum) # Recompute cheap intermediates gate_activation = torch.ops.aten.silu(gate_output) gated_output = gate_activation * up_output + expert_scores_dtype = expert_scores.dtype + expert_scores = expert_scores.to(gated_output.dtype) + gated_weighted = gated_output * expert_scores.unsqueeze(-1) # dgrad FC2 - grad_gated_output = quack_group_gemm_same_nk( + grad_gated_weighted = quack_group_gemm_same_nk( a=grad_output, b=down_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=True, cu_seqlens_m=cu_seqlens_m ) @@ -400,7 +432,7 @@ def backward(ctx, grad_output): if down_proj.requires_grad: grad_down_proj = torch.empty_like(down_proj) quack_group_gemm_same_mn( - a=gated_output, + a=gated_weighted, b=grad_output, c=grad_down_proj, cumsum_K=cumsum, @@ -409,7 +441,13 @@ def backward(ctx, grad_output): transpose_b=False, cu_seqlens_k=cu_seqlens_m, ) - del gated_output + grad_expert_scores = None + if ctx.has_expert_scores: + grad_expert_scores = (grad_gated_weighted * gated_output).sum(dim=-1).to(expert_scores_dtype) + del gated_output, gated_weighted + + grad_gated_output = grad_gated_weighted * expert_scores.unsqueeze(-1) + del grad_gated_weighted # Activation backward grad_up_output = gate_activation * grad_gated_output @@ -456,7 +494,7 @@ def backward(ctx, grad_output): ) del grad_up_output - return grad_permute_tokens, None, grad_gate_proj, grad_up_proj, grad_down_proj + return grad_permute_tokens, None, grad_gate_proj, grad_up_proj, grad_down_proj, grad_expert_scores class QuackTPMoeExpertsFunction(torch.autograd.Function): diff --git a/src/xorl/ops/moe/triton.py b/src/xorl/ops/moe/triton.py index 6db8c8a5..3249936d 100644 --- a/src/xorl/ops/moe/triton.py +++ b/src/xorl/ops/moe/triton.py @@ -25,8 +25,9 @@ class TritonEPGroupGemm(torch.autograd.Function): """ @staticmethod - def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): + def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores=None): max_M = permute_tokens.shape[0] + ctx.has_expert_scores = expert_scores is not None gate_output = group_gemm_same_nk( a=permute_tokens, @@ -47,6 +48,8 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): gate_activation = torch.ops.aten.silu(gate_output) gated_output = gate_activation * up_output + if expert_scores is not None: + gated_output = gated_output * expert_scores.to(gated_output.dtype).unsqueeze(-1) del gate_activation down_output = group_gemm_same_nk( @@ -59,7 +62,11 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): ) del gated_output - ctx.save_for_backward(permute_tokens, cumsum, gate_proj, up_proj, down_proj, gate_output, up_output) + if expert_scores is None: + expert_scores = permute_tokens.new_ones(permute_tokens.shape[0]) + ctx.save_for_backward( + permute_tokens, cumsum, gate_proj, up_proj, down_proj, gate_output, up_output, expert_scores + ) return down_output @staticmethod @@ -72,15 +79,19 @@ def backward(ctx, grad_output): down_proj, gate_output, up_output, + expert_scores, ) = ctx.saved_tensors max_M = grad_output.shape[0] # Recompute activation gate_activation = torch.ops.aten.silu(gate_output) gated_output = gate_activation * up_output + expert_scores_dtype = expert_scores.dtype + expert_scores = expert_scores.to(gated_output.dtype) + gated_weighted = gated_output * expert_scores.unsqueeze(-1) # dgrad FC2 - grad_gated_output = group_gemm_same_nk( + grad_gated_weighted = group_gemm_same_nk( a=grad_output, b=down_proj, cumsum_M=cumsum, @@ -93,7 +104,7 @@ def backward(ctx, grad_output): if down_proj.requires_grad: grad_down_proj = torch.empty_like(down_proj) group_gemm_same_mn( - a=gated_output, + a=gated_weighted, b=grad_output, c=grad_down_proj, cumsum_K=cumsum, @@ -101,7 +112,13 @@ def backward(ctx, grad_output): transpose_a=True, transpose_b=False, ) - del gated_output + grad_expert_scores = None + if ctx.has_expert_scores: + grad_expert_scores = (grad_gated_weighted * gated_output).sum(dim=-1).to(expert_scores_dtype) + del gated_output, gated_weighted + + grad_gated_output = grad_gated_weighted * expert_scores.unsqueeze(-1) + del grad_gated_weighted # Activation backward grad_up_output = gate_activation * grad_gated_output @@ -162,6 +179,7 @@ def backward(ctx, grad_output): grad_gate_proj, grad_up_proj, grad_down_proj, + grad_expert_scores, ) @@ -173,8 +191,9 @@ class TritonEPGroupGemmMoeAct(torch.autograd.Function): """ @staticmethod - def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): + def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores=None): max_M = permute_tokens.shape[0] + ctx.has_expert_scores = expert_scores is not None gate_output = group_gemm_same_nk( a=permute_tokens, @@ -195,6 +214,8 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): gate_activation = torch.ops.aten.silu(gate_output) gated_output = gate_activation * up_output + if expert_scores is not None: + gated_output = gated_output * expert_scores.to(gated_output.dtype).unsqueeze(-1) del gate_activation down_output = group_gemm_same_nk( @@ -208,12 +229,14 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj): del gated_output # moe_act: save only 5 tensors (drop gate_output, up_output) - ctx.save_for_backward(permute_tokens, cumsum, gate_proj, up_proj, down_proj) + if expert_scores is None: + expert_scores = permute_tokens.new_ones(permute_tokens.shape[0]) + ctx.save_for_backward(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores) return down_output @staticmethod def backward(ctx, grad_output): - permute_tokens, cumsum, gate_proj, up_proj, down_proj = ctx.saved_tensors + permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores = ctx.saved_tensors max_M = grad_output.shape[0] # Recompute gate_output and up_output from saved inputs + weights @@ -237,9 +260,12 @@ def backward(ctx, grad_output): # Rest identical to TritonEPGroupGemm.backward gate_activation = torch.ops.aten.silu(gate_output) gated_output = gate_activation * up_output + expert_scores_dtype = expert_scores.dtype + expert_scores = expert_scores.to(gated_output.dtype) + gated_weighted = gated_output * expert_scores.unsqueeze(-1) # dgrad FC2 - grad_gated_output = group_gemm_same_nk( + grad_gated_weighted = group_gemm_same_nk( a=grad_output, b=down_proj, cumsum_M=cumsum, @@ -252,7 +278,7 @@ def backward(ctx, grad_output): if down_proj.requires_grad: grad_down_proj = torch.empty_like(down_proj) group_gemm_same_mn( - a=gated_output, + a=gated_weighted, b=grad_output, c=grad_down_proj, cumsum_K=cumsum, @@ -260,7 +286,13 @@ def backward(ctx, grad_output): transpose_a=True, transpose_b=False, ) - del gated_output + grad_expert_scores = None + if ctx.has_expert_scores: + grad_expert_scores = (grad_gated_weighted * gated_output).sum(dim=-1).to(expert_scores_dtype) + del gated_output, gated_weighted + + grad_gated_output = grad_gated_weighted * expert_scores.unsqueeze(-1) + del grad_gated_weighted # Activation backward grad_up_output = gate_activation * grad_gated_output @@ -321,6 +353,7 @@ def backward(ctx, grad_output): grad_gate_proj, grad_up_proj, grad_down_proj, + grad_expert_scores, ) diff --git a/tests/distributed/test_deepep_correctness.py b/tests/distributed/test_deepep_correctness.py index 9686ca36..a015493b 100644 --- a/tests/distributed/test_deepep_correctness.py +++ b/tests/distributed/test_deepep_correctness.py @@ -104,7 +104,8 @@ def run_one_pass( dispatch_kwargs["num_local_experts"] = num_local_experts permuted, cumsum, ctx = dispatch_fn(**dispatch_kwargs) - expert_out = compute_fn(permuted, cumsum, gate_proj, up_proj, down_proj) + expert_scores = getattr(ctx, "expert_scores", getattr(ctx, "permuted_scores", None)) + expert_out = compute_fn(permuted, cumsum, gate_proj, up_proj, down_proj, expert_scores) if ep_dispatch == "alltoall": combine_kwargs = dict(expert_output=expert_out, ctx=ctx, ep_group=ep_group) diff --git a/tests/distributed/test_moe_memory_efficient_permutation.py b/tests/distributed/test_moe_memory_efficient_permutation.py new file mode 100644 index 00000000..873fb9bf --- /dev/null +++ b/tests/distributed/test_moe_memory_efficient_permutation.py @@ -0,0 +1,149 @@ +import pytest +import torch +import torch.nn.functional as F + +import xorl.distributed.moe.alltoall as alltoall_module +from xorl.distributed.moe.utils import permute, permuted_weights, sort_chunks_by_idxs, unpermute + + +pytestmark = [pytest.mark.cpu] + + +def test_permuted_weights_follow_expert_sorted_token_order(): + tokens = torch.arange(12, dtype=torch.float32).view(4, 3) + selected_experts = torch.tensor( + [ + [2, 0], + [1, 2], + [0, 1], + [2, 1], + ], + dtype=torch.long, + ) + routing_weights = torch.tensor( + [ + [0.9, 0.1], + [0.6, 0.4], + [0.7, 0.3], + [0.8, 0.2], + ], + dtype=torch.float32, + ) + + expert_mask = F.one_hot(selected_experts, num_classes=3).permute(2, 1, 0) + routing_map = expert_mask.sum(dim=1) + _, permutation_mapping = permute(tokens, routing_map) + + expected = torch.tensor([0.1, 0.7, 0.6, 0.3, 0.2, 0.9, 0.4, 0.8], dtype=torch.float32) + + torch.testing.assert_close(permuted_weights(routing_weights, selected_experts, 3), expected) + torch.testing.assert_close(permutation_mapping, torch.tensor([0, 2, 1, 2, 3, 0, 1, 3])) + + +def test_unpermute_only_scatter_adds_preweighted_outputs(): + expert_outputs = torch.tensor( + [ + [1.0, 1.0], + [2.0, 2.0], + [3.0, 3.0], + ] + ) + permutation_mapping = torch.tensor([0, 1, 0], dtype=torch.long) + + output = unpermute( + expert_outputs, + hidden_states_shape=torch.Size([2, 2]), + permutation_mapping=permutation_mapping, + ) + + expected = torch.tensor( + [ + [4.0, 4.0], + [2.0, 2.0], + ] + ) + torch.testing.assert_close(output, expected) + + +def test_alltoall_pre_dispatch_routes_scores_with_received_token_order(monkeypatch): + class FakeGroup: + def size(self): + return 2 + + num_experts = 4 + selection = torch.tensor([0, 2, 3, 5, 7], dtype=torch.long) + input_splits = [4, 4] + output_splits = [3, 2] + num_tokens_per_expert = torch.tensor([[1, 2], [1, 1]], dtype=torch.int64) + sum_tokens = torch.tensor([2, 3], dtype=torch.int64) + + def fake_preprocess(*, expert_mask, num_experts, ep_group): + return input_splits, output_splits, num_tokens_per_expert, sum_tokens + + def fake_all_to_all(group, input, output_split_sizes, input_split_sizes): + return input.index_select(0, selection.to(input.device)) + + monkeypatch.setattr(alltoall_module, "preprocess", fake_preprocess) + monkeypatch.setattr(alltoall_module, "all_to_all", fake_all_to_all) + + hidden_states = torch.tensor( + [ + [1.0, 10.0], + [2.0, 20.0], + [3.0, 30.0], + [4.0, 40.0], + ] + ) + selected_experts = torch.tensor( + [ + [2, 0], + [1, 2], + [0, 1], + [2, 1], + ], + dtype=torch.long, + ) + routing_weights = torch.tensor( + [ + [0.9, 0.1], + [0.6, 0.4], + [0.7, 0.3], + [0.8, 0.2], + ], + dtype=torch.float32, + requires_grad=True, + ) + + expert_mask = F.one_hot(selected_experts, num_classes=num_experts).permute(2, 1, 0) + routing_map = expert_mask.sum(dim=1) + local_tokens, _ = permute(hidden_states, routing_map) + local_scores = permuted_weights(routing_weights, selected_experts, num_experts) + expected_tokens = sort_chunks_by_idxs( + local_tokens.index_select(0, selection), num_tokens_per_expert.ravel(), [0, 2, 1, 3] + ) + expected_scores = sort_chunks_by_idxs( + local_scores.index_select(0, selection), num_tokens_per_expert.ravel(), [0, 2, 1, 3] + ) + + permuted_tokens, cumsum, ctx = alltoall_module.alltoall_pre_dispatch( + hidden_states=hidden_states, + routing_weights=routing_weights, + selected_experts=selected_experts, + num_experts=num_experts, + ep_group=FakeGroup(), + ) + + torch.testing.assert_close(permuted_tokens, expected_tokens) + torch.testing.assert_close(ctx.expert_scores, expected_scores) + torch.testing.assert_close(cumsum, torch.tensor([2, 5], dtype=cumsum.dtype)) + + ctx.expert_scores.sum().backward() + expected_grad = torch.tensor( + [ + [1.0, 1.0], + [1.0, 0.0], + [0.0, 1.0], + [1.0, 0.0], + ] + ) + torch.testing.assert_close(routing_weights.grad, expected_grad) diff --git a/tests/ops/test_ep_routing_scores.py b/tests/ops/test_ep_routing_scores.py new file mode 100644 index 00000000..ae80baac --- /dev/null +++ b/tests/ops/test_ep_routing_scores.py @@ -0,0 +1,180 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest +import torch +import torch.nn.functional as F + + +pytestmark = pytest.mark.cpu + +_MODULE_PATHS = { + "xorl.ops.moe.triton": Path(__file__).resolve().parents[2] / "src/xorl/ops/moe/triton.py", + "xorl.ops.moe.quack": Path(__file__).resolve().parents[2] / "src/xorl/ops/moe/quack.py", +} + + +def _counts_from_cumsum(cumsum: torch.Tensor) -> list[int]: + counts = torch.empty_like(cumsum) + counts[0] = cumsum[0] + counts[1:] = cumsum[1:] - cumsum[:-1] + return counts.tolist() + + +def _naive_group_gemm_same_nk(a, b, cumsum_M, max_M, transpose_a=False, transpose_b=False, **kwargs): + del max_M, kwargs + assert not transpose_a + + outputs = [] + start = 0 + for expert_idx, count in enumerate(_counts_from_cumsum(cumsum_M)): + end = start + count + weight = b[expert_idx] + if transpose_b: + outputs.append(a[start:end] @ weight.transpose(0, 1)) + else: + outputs.append(a[start:end] @ weight) + start = end + + return torch.cat(outputs, dim=0) + + +def _naive_group_gemm_same_mn(a, b, c, cumsum_K, max_K, transpose_a=False, transpose_b=False, **kwargs): + del max_K, kwargs + + start = 0 + for expert_idx, count in enumerate(_counts_from_cumsum(cumsum_K)): + end = start + count + lhs = a[start:end].transpose(0, 1) if transpose_a else a[start:end] + rhs = b[start:end].transpose(0, 1) if transpose_b else b[start:end] + c[expert_idx].copy_(lhs @ rhs) + start = end + + return c + + +def _patch_ep_kernels(monkeypatch, module_name: str): + import xorl.utils.import_utils as import_utils + + moe_stub = types.ModuleType("xorl.ops.group_gemm.kernel.moe") + moe_stub.expert_histogram = None + moe_stub.moe_gather = None + moe_stub.moe_index_compute = None + moe_stub.moe_scatter = None + monkeypatch.setattr(import_utils, "is_fused_moe_available", lambda: True) + sys.modules.pop("xorl.ops.group_gemm.kernel.moe", None) + sys.modules.pop("xorl.ops.group_gemm.kernel.group_gemm", None) + sys.modules.pop("xorl.ops.group_gemm.kernel.quack", None) + monkeypatch.setitem( + sys.modules, + "xorl.ops.group_gemm.kernel.moe", + moe_stub, + ) + if module_name.endswith("triton"): + group_gemm_stub = types.ModuleType("xorl.ops.group_gemm.kernel.group_gemm") + group_gemm_stub.group_gemm_same_nk = _naive_group_gemm_same_nk + group_gemm_stub.group_gemm_same_mn = _naive_group_gemm_same_mn + monkeypatch.setitem( + sys.modules, + "xorl.ops.group_gemm.kernel.group_gemm", + group_gemm_stub, + ) + else: + quack_stub = types.ModuleType("xorl.ops.group_gemm.kernel.quack") + quack_stub.cumsum_to_cu_seqlens = lambda cumsum: cumsum + quack_stub.quack_group_gemm_same_nk = _naive_group_gemm_same_nk + quack_stub.quack_group_gemm_same_mn = _naive_group_gemm_same_mn + monkeypatch.setitem( + sys.modules, + "xorl.ops.group_gemm.kernel.quack", + quack_stub, + ) + spec = importlib.util.spec_from_file_location( + f"codex_test_{module_name.rsplit('.', maxsplit=1)[-1]}", _MODULE_PATHS[module_name] + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def reference_ep_forward( + permute_tokens: torch.Tensor, + cumsum: torch.Tensor, + gate_proj: torch.Tensor, + up_proj: torch.Tensor, + down_proj: torch.Tensor, + expert_scores: torch.Tensor, +) -> torch.Tensor: + outputs = [] + start = 0 + for expert_idx, count in enumerate(_counts_from_cumsum(cumsum)): + end = start + count + x = permute_tokens[start:end] + h = F.silu(x @ gate_proj[expert_idx]) * (x @ up_proj[expert_idx]) + h = h * expert_scores[start:end].to(h.dtype).unsqueeze(-1) + outputs.append(h @ down_proj[expert_idx]) + start = end + + return torch.cat(outputs, dim=0) + + +@pytest.mark.parametrize( + ("module_name", "class_name"), + [ + pytest.param("xorl.ops.moe.triton", "TritonEPGroupGemm", id="triton"), + pytest.param("xorl.ops.moe.triton", "TritonEPGroupGemmMoeAct", id="triton-moe-act"), + pytest.param("xorl.ops.moe.quack", "QuackEPGroupGemm", id="quack"), + pytest.param("xorl.ops.moe.quack", "QuackEPGroupGemmMoeAct", id="quack-moe-act"), + ], +) +def test_ep_group_gemm_propagates_routing_score_gradients(monkeypatch, module_name, class_name): + try: + module = _patch_ep_kernels(monkeypatch, module_name) + except ImportError as exc: + pytest.skip(f"{module_name} unavailable: {exc}") + + fn = getattr(module, class_name) + + torch.manual_seed(0) + dtype = torch.float32 + num_local_experts = 2 + hidden_dim = 8 + intermediate_size = 12 + counts = torch.tensor([3, 2]) + cumsum = torch.cumsum(counts, dim=0) + num_tokens = int(cumsum[-1].item()) + + permute_tokens = torch.randn(num_tokens, hidden_dim, dtype=dtype) + gate_proj = torch.randn(num_local_experts, hidden_dim, intermediate_size, dtype=dtype) + up_proj = torch.randn(num_local_experts, hidden_dim, intermediate_size, dtype=dtype) + down_proj = torch.randn(num_local_experts, intermediate_size, hidden_dim, dtype=dtype) + expert_scores = torch.rand(num_tokens, dtype=dtype, requires_grad=True) + upstream = torch.randn(num_tokens, hidden_dim, dtype=dtype) + + output = fn.apply( + permute_tokens, + cumsum, + gate_proj, + up_proj, + down_proj, + expert_scores, + ) + output.backward(upstream) + grad_scores = expert_scores.grad.detach().clone() + + expert_scores_ref = expert_scores.detach().clone().requires_grad_(True) + ref_output = reference_ep_forward( + permute_tokens, + cumsum, + gate_proj, + up_proj, + down_proj, + expert_scores_ref, + ) + ref_output.backward(upstream) + + torch.testing.assert_close(output, ref_output) + torch.testing.assert_close(grad_scores, expert_scores_ref.grad) From 03bd37376aca4341aad069a7a2aef600e57c5182 Mon Sep 17 00:00:00 2001 From: Zhongzhu Zhou Date: Wed, 8 Apr 2026 11:36:03 +1000 Subject: [PATCH 17/41] qwen3.5 correction (#59) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: add Qwen3.5 MoE support and fix multi-node training stability - Add Qwen3.5-35B-A3B and 397B-A17B server configs (full + LoRA) - Fix weight loading deadlock: remove Gloo broadcast from _try_load_state_dict - Fix MoE numeric stability: FP32 scatter_add accumulation, chunked DeepEP combine - Fix Ulysses SP position_ids generation for correct RoPE on ranks > 0 - Handle variable experts-per-tok in routing replay (Qwen3.5 uses topk=10) - Fix LoRA adapter DTensor shard boundary, add fast checkpoint save path - Increase timeouts for large model init and weight sync * Fix fused gate_up_proj OOM: fused forward GEMM + in-place backward dgrad + Muon split MoE memory optimization for fused gate_up_proj [E, H, 2I]: - Fused forward GEMM: single x @ gate_up_proj instead of 2 separate GEMMs, zero .contiguous() copies, save_for_backward stores param reference (zero-copy) - In-place backward dgrad: accumulate 2nd dgrad into grad_permute_tokens via c= parameter (ACCUMULATE_TO_C), avoiding 7 GiB temp allocation - Kernel fix: cast ACCUMULATE_TO_C loaded values to float32 for tl.dot compat - Muon split: detect fused [E, H, 2I] params and run Newton-Schulz on each [E, H, I] half independently (4x faster, same math) - Warmup: pre-compile ACCUMULATE_TO_C=True kernel variant before training - EP compute: fused interface (gate_up_proj, down_proj, intermediate_size), quack/native adapters for backward compat - Checkpoint handler: fix internal pattern transpose for HF Qwen3.5 weights - Config: moe_act mode, no autotune, 60min NCCL timeout * Adjust group gemm to avoid explicit memory copy * Fuse EP backward GEMMs and eliminate .contiguous() weight copies Backward FC1: 4 split GEMMs + 2 weight copies → 2 fused GEMMs. Expand kernel warmup to all backward variants to avoid compilation OOM. Add XORL_TRITON_NO_AUTOTUNE=1 to skip autotune benchmarking on large models. * Fix ruff formatting * Fix review feedback: LoRA save gather bug and Muon fused_split heuristic - Remove incorrect EP gather in LoRA checkpoint save: lora_params already stores global-shape tensors, gather was duplicating data and dst=0 as world rank could hang on non-root EP groups. - Replace shape-based fused_split heuristic with explicit param name check: build _fused_gate_up_ids set from param names containing "gate_up_proj" as @kiddyboots216 suggest. - add 397B benchmark config * Fix ruff-format: break long lines in optimizer.py * Fix train_router not propagated to MoE subclasses, causing crash with deepep and fix issued mentioned by Ashwinee * Merge main: resolve MoE conflicts, combine fused gate_up with expert_scores (cherry picked from commit b5631452f53de8231784cc3047267ebd5b6d8eec) --- README.md | 2 + docs/src/content/docs/models.md | 23 +- .../full/qwen3_5_35b_a3b_ep1_fsdp8.yaml | 39 ++++ .../configs/full/qwen3_5_35b_a3b_ep8_cp1.yaml | 40 ++++ .../configs/full/qwen3_5_35b_a3b_full.yaml | 1 + .../full/qwen3_5_35b_a3b_full_no_deepep.yaml | 45 ++++ .../configs/full/qwen3_5_397b_a17b_full.yaml | 54 +++++ .../configs/lora/qwen3_5_35b_a3b_lora.yaml | 49 +++++ .../configs/lora/qwen3_5_397b_a17b_lora.yaml | 49 +++++ src/xorl/cli/train.py | 12 ++ src/xorl/distributed/moe/deepep.py | 7 +- src/xorl/distributed/moe/utils.py | 10 +- .../models/layers/moe/backend/__init__.py | 48 +++-- src/xorl/models/layers/moe/experts.py | 58 +++-- src/xorl/models/layers/moe/lora.py | 1 + src/xorl/models/module_utils.py | 13 +- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 1 + .../qwen3_moe/modeling_qwen3_moe.py | 4 + src/xorl/ops/group_gemm/kernel/group_gemm.py | 21 +- src/xorl/ops/moe/triton.py | 200 ++++++------------ src/xorl/optim/muon.py | 16 +- src/xorl/optim/optimizer.py | 26 ++- src/xorl/server/api_server/weights.py | 2 +- src/xorl/server/backend/remote.py | 2 +- src/xorl/server/launcher.py | 1 + src/xorl/server/runner/adapters/manager.py | 13 +- src/xorl/server/runner/checkpoint/manager.py | 30 ++- src/xorl/server/runner/model_runner.py | 37 +++- src/xorl/server/runner/runner_dispatcher.py | 16 ++ .../runner/utils/routing_replay_handler.py | 34 ++- src/xorl/trainers/trainer.py | 4 +- src/xorl/utils/helper.py | 2 +- 32 files changed, 662 insertions(+), 198 deletions(-) create mode 100644 examples/server/configs/full/qwen3_5_35b_a3b_ep1_fsdp8.yaml create mode 100644 examples/server/configs/full/qwen3_5_35b_a3b_ep8_cp1.yaml create mode 100644 examples/server/configs/full/qwen3_5_35b_a3b_full_no_deepep.yaml create mode 100644 examples/server/configs/full/qwen3_5_397b_a17b_full.yaml create mode 100644 examples/server/configs/lora/qwen3_5_35b_a3b_lora.yaml create mode 100644 examples/server/configs/lora/qwen3_5_397b_a17b_lora.yaml diff --git a/README.md b/README.md index 64604525..b663c727 100644 --- a/README.md +++ b/README.md @@ -78,6 +78,8 @@ See the [quick start guide](https://togethercomputer.github.io/xorl/getting-star |---|---|---| | Qwen3 | Dense | `Qwen/Qwen3-8B`, `Qwen/Qwen3-32B`, ... | | Qwen3-MoE | Mixture-of-Experts | `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-235B-A22B`, ... | +| Qwen3.5 | Dense | `Qwen/Qwen3.5-7B`, ... | +| Qwen3.5-MoE | Mixture-of-Experts | `Qwen/Qwen3.5-35B-A3B`, `Qwen/Qwen3.5-397B-A17B`, ... | Models are loaded directly from HuggingFace checkpoints — no preprocessing needed. See the [supported models](https://togethercomputer.github.io/xorl/models/) page for details. diff --git a/docs/src/content/docs/models.md b/docs/src/content/docs/models.md index a17343c6..65cf930b 100644 --- a/docs/src/content/docs/models.md +++ b/docs/src/content/docs/models.md @@ -6,10 +6,12 @@ xorl currently supports the following model architectures. ## Architectures -| Architecture | HuggingFace class | Notes | -|---|---|---| -| Qwen3 (dense) | `Qwen3ForCausalLM` | Standard transformer with GQA, SwiGLU activations, RoPE positional embeddings. | -| Qwen3-MoE | `Qwen3MoeForCausalLM` | Same as Qwen3 but FFN layers are replaced with Mixture-of-Experts blocks. Weight conversion is automatic — see below. | +| Architecture | HuggingFace class | Example models | Notes | +|---|---|---|---| +| Qwen3 (dense) | `Qwen3ForCausalLM` | `Qwen/Qwen3-8B`, `Qwen/Qwen3-32B` | Standard transformer with GQA, SwiGLU, RoPE. | +| Qwen3-MoE | `Qwen3MoeForCausalLM` | `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-235B-A22B` | Mixture-of-Experts with top-k routing. Weight conversion is automatic — see below. | +| Qwen3.5 (dense) | `Qwen3_5ForCausalLM` | `Qwen/Qwen3.5-7B` | Qwen3.5 dense with hybrid full/linear attention layers. | +| Qwen3.5-MoE | `Qwen3_5MoeForCausalLM` | `Qwen/Qwen3.5-35B-A3B`, `Qwen/Qwen3.5-397B-A17B` | Qwen3.5 MoE with hybrid attention and grouped expert routing. | Model selection is config-based: xorl reads `model_type` from `config.json` inside `model_path` and instantiates the appropriate class automatically. @@ -33,6 +35,19 @@ Specify the checkpoint with `model_path` (local path or HF Hub ID). Use `config_ | `attn_implementation` | Attention backend: `flash_attention_3`, `flash_attention_4`, `native`, `sdpa`, `eager`. | | `moe_implementation` | MoE kernel: `null` (auto), `triton`, `native`, `quack`, `eager`. | +## Tested configurations + +The following model + training-mode combinations have pre-built example configs under `examples/server/configs/`: + +| Model | Full weights | LoRA | +|---|---|---| +| Qwen3-8B | `qwen3_8b_full.yaml` | `qwen3_8b_lora.yaml` | +| Qwen3-30B-A3B (MoE) | `qwen3_30b_a3b_full.yaml` | — | +| Qwen3-Coder-30B-A3B (MoE) | `qwen3_coder_30b_a3b_full.yaml` | `qwen3_coder_30b_a3b_lora.yaml` | +| Qwen3-235B-A22B (MoE) | `qwen3_235b_a22b_8node_ep64.yaml` | — | +| Qwen3.5-35B-A3B (MoE) | `qwen3_5_35b_a3b_full.yaml` | `qwen3_5_35b_a3b_lora.yaml` | +| Qwen3.5-397B-A17B (MoE) | `qwen3_5_397b_a17b_full.yaml` | `qwen3_5_397b_a17b_lora.yaml` | + ## MoE models: automatic weight conversion MoE checkpoints from HuggingFace store experts as a `ModuleList` (one module per expert). xorl uses fused grouped-kernel (GKN) tensors for efficient expert dispatch. This conversion happens **automatically during model loading** — no separate preprocessing step is needed. Simply point `model_path` at the standard HuggingFace checkpoint and xorl will fuse the expert weights on the fly. diff --git a/examples/server/configs/full/qwen3_5_35b_a3b_ep1_fsdp8.yaml b/examples/server/configs/full/qwen3_5_35b_a3b_ep1_fsdp8.yaml new file mode 100644 index 00000000..ccbb7187 --- /dev/null +++ b/examples/server/configs/full/qwen3_5_35b_a3b_ep1_fsdp8.yaml @@ -0,0 +1,39 @@ +# EP=1, FSDP=8 — isolate whether EP/quack causes logprob divergence +model_path: Qwen/Qwen3.5-35B-A3B +tokenizer_path: Qwen/Qwen3.5-35B-A3B +attn_implementation: flash_attention_3 +moe_implementation: eager + +data_parallel_mode: fsdp2 +expert_parallel_size: 1 +ulysses_parallel_size: 1 +data_parallel_replicate_size: 1 +data_parallel_shard_size: 8 + +enable_mixed_precision: true +enable_gradient_checkpointing: true +enable_full_shard: true +enable_activation_offload: false +init_device: meta +load_weights_mode: all_ranks + +output_dir: outputs/Qwen3.5-35B-A3B-server-ep1-fsdp8 +load_checkpoint_path: "" +ckpt_manager: dcp + +log_level: INFO + +worker_connection_timeout: 60.0 +worker_max_retries: 3 + +sample_packing_sequence_len: 128000 +enable_packing: true + +skip_initial_checkpoint: true +ce_mode: compiled + +optimizer: muon +optimizer_dtype: bf16 +muon_lr: 0.02 +muon_momentum: 0.95 +muon_adjust_lr_fn: match_rms_adamw diff --git a/examples/server/configs/full/qwen3_5_35b_a3b_ep8_cp1.yaml b/examples/server/configs/full/qwen3_5_35b_a3b_ep8_cp1.yaml new file mode 100644 index 00000000..5331cb42 --- /dev/null +++ b/examples/server/configs/full/qwen3_5_35b_a3b_ep8_cp1.yaml @@ -0,0 +1,40 @@ +# EP=8, CP=1 — isolate whether CP causes logprob divergence +model_path: Qwen/Qwen3.5-35B-A3B +tokenizer_path: Qwen/Qwen3.5-35B-A3B +attn_implementation: flash_attention_3 +moe_implementation: quack +ep_dispatch: deepep + +data_parallel_mode: fsdp2 +expert_parallel_size: 8 +ulysses_parallel_size: 1 +data_parallel_replicate_size: 1 +data_parallel_shard_size: 1 + +enable_mixed_precision: true +enable_gradient_checkpointing: true +enable_full_shard: true +enable_activation_offload: false +init_device: meta +load_weights_mode: all_ranks + +output_dir: outputs/Qwen3.5-35B-A3B-server-ep8-cp1 +load_checkpoint_path: "" +ckpt_manager: dcp + +log_level: INFO + +worker_connection_timeout: 60.0 +worker_max_retries: 3 + +sample_packing_sequence_len: 128000 +enable_packing: true + +skip_initial_checkpoint: true +ce_mode: compiled + +optimizer: muon +optimizer_dtype: bf16 +muon_lr: 0.02 +muon_momentum: 0.95 +muon_adjust_lr_fn: match_rms_adamw diff --git a/examples/server/configs/full/qwen3_5_35b_a3b_full.yaml b/examples/server/configs/full/qwen3_5_35b_a3b_full.yaml index d58083c9..ef6bb58a 100644 --- a/examples/server/configs/full/qwen3_5_35b_a3b_full.yaml +++ b/examples/server/configs/full/qwen3_5_35b_a3b_full.yaml @@ -21,6 +21,7 @@ enable_gradient_checkpointing: true enable_full_shard: true enable_activation_offload: false init_device: meta +load_weights_mode: all_ranks output_dir: outputs/Qwen3.5-35B-A3B-server-full-rl load_checkpoint_path: "" diff --git a/examples/server/configs/full/qwen3_5_35b_a3b_full_no_deepep.yaml b/examples/server/configs/full/qwen3_5_35b_a3b_full_no_deepep.yaml new file mode 100644 index 00000000..932d8d6b --- /dev/null +++ b/examples/server/configs/full/qwen3_5_35b_a3b_full_no_deepep.yaml @@ -0,0 +1,45 @@ +# Server-side configuration for XORL Training Server (Qwen3.5-35B-A3B full weights, bf16) +# +# Same as qwen3_5_35b_a3b_full.yaml but WITHOUT DeepEP. +# Uses alltoall EP dispatch + triton MoE kernels to better align with SGLang. +# Purpose: test if DeepEP + Quack kernels cause the logprob divergence vs SGLang. + +model_path: Qwen/Qwen3.5-35B-A3B +tokenizer_path: Qwen/Qwen3.5-35B-A3B +attn_implementation: flash_attention_3 +moe_implementation: triton +ep_dispatch: alltoall + +data_parallel_mode: fsdp2 +expert_parallel_size: 8 +ulysses_parallel_size: 8 +data_parallel_replicate_size: 1 +data_parallel_shard_size: 1 + +enable_mixed_precision: true +enable_gradient_checkpointing: true +enable_full_shard: true +enable_activation_offload: false +init_device: meta +load_weights_mode: all_ranks + +output_dir: outputs/Qwen3.5-35B-A3B-server-full-rl-no-deepep +load_checkpoint_path: "" +ckpt_manager: dcp + +log_level: INFO + +worker_connection_timeout: 60.0 +worker_max_retries: 3 + +sample_packing_sequence_len: 128000 +enable_packing: true + +skip_initial_checkpoint: true +ce_mode: compiled + +optimizer: muon +optimizer_dtype: bf16 +muon_lr: 0.02 +muon_momentum: 0.95 +muon_adjust_lr_fn: match_rms_adamw diff --git a/examples/server/configs/full/qwen3_5_397b_a17b_full.yaml b/examples/server/configs/full/qwen3_5_397b_a17b_full.yaml new file mode 100644 index 00000000..cab3a3b2 --- /dev/null +++ b/examples/server/configs/full/qwen3_5_397b_a17b_full.yaml @@ -0,0 +1,54 @@ +# Server-side configuration for XORL Training Server (Qwen3.5-397B-A17B full weights, bf16) +# +# Full weight RL fine-tuning of the 397B MoE model (512 experts, 17B activated). +# 64 GPUs (8 nodes x 8): EP=64 (8 experts/rank), Ulysses SP=8, FSDP shard=8. +# +# Memory budget per GPU (H100 80GB): +# Expert params (387B/64): ~12 GB +# Dense params (10B/8): ~2.5 GB +# Muon optimizer (1x): ~14.5 GB +# Gradients: ~14.5 GB +# Activations + overhead: ~36 GB headroom +# +# Uses triton MoE + alltoall dispatch for reliable cross-node expert communication. +# For DeepEP-optimized clusters, switch to moe_implementation: quack / ep_dispatch: deepep. + +model_path: Qwen/Qwen3.5-397B-A17B +tokenizer_path: Qwen/Qwen3.5-397B-A17B +attn_implementation: flash_attention_3 +moe_implementation: triton +ep_dispatch: alltoall + +data_parallel_mode: fsdp2 +expert_parallel_size: 64 +ulysses_parallel_size: 8 +data_parallel_replicate_size: 1 +data_parallel_shard_size: 8 + +enable_mixed_precision: true +enable_gradient_checkpointing: true +enable_full_shard: true +enable_activation_offload: false +init_device: meta +load_weights_mode: all_ranks + +output_dir: outputs/Qwen3.5-397B-A17B-server-full-rl +load_checkpoint_path: "" +ckpt_manager: dcp + +log_level: INFO + +worker_connection_timeout: 180.0 +worker_max_retries: 5 + +sample_packing_sequence_len: 65536 +enable_packing: true + +skip_initial_checkpoint: true +ce_mode: compiled + +optimizer: muon +optimizer_dtype: bf16 +muon_lr: 0.02 +muon_momentum: 0.95 +muon_adjust_lr_fn: match_rms_adamw diff --git a/examples/server/configs/lora/qwen3_5_35b_a3b_lora.yaml b/examples/server/configs/lora/qwen3_5_35b_a3b_lora.yaml new file mode 100644 index 00000000..6020b72c --- /dev/null +++ b/examples/server/configs/lora/qwen3_5_35b_a3b_lora.yaml @@ -0,0 +1,49 @@ +# Server-side configuration for XORL Training Server +# Qwen3.5-35B-A3B LoRA (bf16 base + LoRA rank 32) +# +# 1 node (8 GPUs): EP=8, Ulysses SP=8, FSDP shard=1 +# +# Memory budget per GPU (H100 80GB): +# Expert params (35B/8): ~8.75 GB (8 experts/GPU, 64 experts total) +# Dense params (3B/1): ~6 GB +# LoRA params + grads: ~1 GB +# Activations + overhead: ~64 GB headroom + +model_path: Qwen/Qwen3.5-35B-A3B +tokenizer_path: Qwen/Qwen3.5-35B-A3B +attn_implementation: flash_attention_3 +moe_implementation: quack +ep_dispatch: deepep + +data_parallel_mode: fsdp2 +expert_parallel_size: 8 +ulysses_parallel_size: 8 +data_parallel_replicate_size: 1 +data_parallel_shard_size: 1 + +enable_mixed_precision: true +enable_gradient_checkpointing: true +enable_full_shard: true +enable_activation_offload: false +init_device: meta +load_weights_mode: all_ranks + +output_dir: outputs/Qwen3.5-35B-A3B-server-lora-rl +load_checkpoint_path: "" +ckpt_manager: dcp + +log_level: INFO + +worker_connection_timeout: 60.0 +worker_max_retries: 3 + +sample_packing_sequence_len: 32768 +enable_packing: true + +enable_lora: true +lora_rank: 32 +lora_alpha: 32 +lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] + +skip_initial_checkpoint: true +ce_mode: compiled diff --git a/examples/server/configs/lora/qwen3_5_397b_a17b_lora.yaml b/examples/server/configs/lora/qwen3_5_397b_a17b_lora.yaml new file mode 100644 index 00000000..4a812b83 --- /dev/null +++ b/examples/server/configs/lora/qwen3_5_397b_a17b_lora.yaml @@ -0,0 +1,49 @@ +# Server-side configuration for XORL Training Server +# Qwen3.5-397B-A17B LoRA (bf16 base + LoRA rank 32) +# +# 4 nodes (32 GPUs): EP=32, Ulysses SP=8, FSDP shard=4 +# +# Memory budget per GPU (H100 80GB): +# Expert params (387B/32): ~24 GB (16 experts/GPU, 512 experts total) +# Dense params (10B/4): ~5 GB (sharded by dp_shard=4) +# LoRA params + grads: ~1 GB +# Activations + overhead: ~50 GB headroom + +model_path: Qwen/Qwen3.5-397B-A17B +tokenizer_path: Qwen/Qwen3.5-397B-A17B +attn_implementation: flash_attention_3 +moe_implementation: triton +ep_dispatch: alltoall + +data_parallel_mode: fsdp2 +expert_parallel_size: 32 +ulysses_parallel_size: 8 +data_parallel_replicate_size: 1 +data_parallel_shard_size: 4 + +enable_mixed_precision: true +enable_gradient_checkpointing: true +enable_full_shard: true +enable_activation_offload: false +init_device: meta +load_weights_mode: all_ranks + +output_dir: outputs/Qwen3.5-397B-A17B-server-lora-rl +load_checkpoint_path: "" +ckpt_manager: dcp + +log_level: INFO + +worker_connection_timeout: 180.0 +worker_max_retries: 5 + +sample_packing_sequence_len: 65536 +enable_packing: true + +enable_lora: true +lora_rank: 32 +lora_alpha: 32 +lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] + +skip_initial_checkpoint: true +ce_mode: compiled diff --git a/src/xorl/cli/train.py b/src/xorl/cli/train.py index 80d91d94..88b7cc78 100644 --- a/src/xorl/cli/train.py +++ b/src/xorl/cli/train.py @@ -10,6 +10,18 @@ os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") configure_rank_local_compile_caches() +# When XORL_TRITON_NO_AUTOTUNE=1, force ALL Triton autotune decorators to +# use only the first config (skip do_bench benchmarking that can OOM). +if os.environ.get("XORL_TRITON_NO_AUTOTUNE", "0") == "1": + import triton + + _orig_autotune = triton.autotune + + def _single_config_autotune(configs, *args, **kwargs): + return _orig_autotune([configs[0]], *args, **kwargs) + + triton.autotune = _single_config_autotune + from xorl.arguments import Arguments, parse_args from xorl.trainers import Trainer diff --git a/src/xorl/distributed/moe/deepep.py b/src/xorl/distributed/moe/deepep.py index 117f5f7c..beb31f2c 100644 --- a/src/xorl/distributed/moe/deepep.py +++ b/src/xorl/distributed/moe/deepep.py @@ -388,11 +388,14 @@ def forward( gather_output = torch.zeros( dispatch_ctx.num_recv_tokens, hidden_dim, - dtype=expert_output.dtype, + dtype=dtype, device=device, ) idx_2d = dispatch_ctx.permuted_indices.unsqueeze(1).expand(-1, hidden_dim) - gather_output.scatter_add_(0, idx_2d, expert_output) + _CHUNK = 4096 + for _i in range(0, expert_output.shape[0], _CHUNK): + _end = min(_i + _CHUNK, expert_output.shape[0]) + gather_output.scatter_add_(0, idx_2d[_i:_end], expert_output[_i:_end]) # Step 2: Combine previous_event = EventOverlap(EventHandle()) diff --git a/src/xorl/distributed/moe/utils.py b/src/xorl/distributed/moe/utils.py index 57760781..4f044e39 100644 --- a/src/xorl/distributed/moe/utils.py +++ b/src/xorl/distributed/moe/utils.py @@ -43,14 +43,12 @@ def unpermute( """ hidden_dim = hidden_states_shape[-1] - unpermuted_tokens = torch.zeros(hidden_states_shape, device=tokens.device, dtype=tokens.dtype) - - # Scatter add the permuted_input back to the original positions expanded_mapping = permutation_mapping.unsqueeze(1).expand(-1, hidden_dim) - unpermuted_tokens.scatter_add_(0, expanded_mapping, tokens) - - return unpermuted_tokens + # FP32 accumulation to reduce non-determinism from BF16 atomic scatter_add_ + unpermuted_fp32 = torch.zeros(hidden_states_shape, device=tokens.device, dtype=torch.float32) + unpermuted_fp32.scatter_add_(0, expanded_mapping, tokens.float()) + return unpermuted_fp32.to(tokens.dtype) def permuted_weights(routing_weights: torch.Tensor, selected_experts: torch.Tensor, num_experts: int) -> torch.Tensor: diff --git a/src/xorl/models/layers/moe/backend/__init__.py b/src/xorl/models/layers/moe/backend/__init__.py index 30f2c3c2..271ef343 100644 --- a/src/xorl/models/layers/moe/backend/__init__.py +++ b/src/xorl/models/layers/moe/backend/__init__.py @@ -43,12 +43,12 @@ # --------------------------------------------------------------------------- # EP expert compute registry # Maps implementation name -> compute callable for Expert Parallelism. -# All compute functions share: (permute_tokens, cumsum, gate_proj, up_proj, down_proj) -> output +# Fused signature: (permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size) -> output # --------------------------------------------------------------------------- EP_EXPERT_COMPUTE: Dict[str, Callable] = {} -# Triton EP compute (custom Triton group GEMM with autograd) +# Triton EP compute (fused gate+up GEMM with autograd) try: from xorl.ops.moe.triton import TritonEPGroupGemm @@ -56,19 +56,29 @@ except ImportError: pass -# Quack EP compute (quack group GEMM with autograd) +# Quack EP compute — adapt fused interface to old (gate_proj, up_proj) signature try: - from xorl.ops.moe.quack import QuackEPGroupGemm + from xorl.ops.moe.quack import QuackEPGroupGemm as _QuackEPGroupGemm - EP_EXPERT_COMPUTE["quack"] = QuackEPGroupGemm.apply + def _quack_ep_fused(permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size): + gate_proj = gate_up_proj[..., :intermediate_size].contiguous() + up_proj = gate_up_proj[..., intermediate_size:].contiguous() + return _QuackEPGroupGemm.apply(permute_tokens, cumsum, gate_proj, up_proj, down_proj) + + EP_EXPERT_COMPUTE["quack"] = _quack_ep_fused except ImportError: pass -# Native EP compute (torch._grouped_mm with alignment padding) +# Native EP compute — adapt fused interface try: - from .native import native_ep_compute + from .native import native_ep_compute as _native_ep_compute + + def _native_ep_fused(permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size): + gate_proj = gate_up_proj[..., :intermediate_size].contiguous() + up_proj = gate_up_proj[..., intermediate_size:].contiguous() + return _native_ep_compute(permute_tokens, cumsum, gate_proj, up_proj, down_proj) - EP_EXPERT_COMPUTE["native"] = native_ep_compute + EP_EXPERT_COMPUTE["native"] = _native_ep_fused except ImportError: pass @@ -191,19 +201,29 @@ except ImportError: pass -# Quack EP moe_act compute +# Quack EP moe_act compute — adapt fused interface try: - from xorl.ops.moe.quack import QuackEPGroupGemmMoeAct + from xorl.ops.moe.quack import QuackEPGroupGemmMoeAct as _QuackEPGroupGemmMoeAct - EP_EXPERT_COMPUTE_MOE_ACT["quack"] = QuackEPGroupGemmMoeAct.apply + def _quack_ep_moe_act_fused(permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size): + gate_proj = gate_up_proj[..., :intermediate_size].contiguous() + up_proj = gate_up_proj[..., intermediate_size:].contiguous() + return _QuackEPGroupGemmMoeAct.apply(permute_tokens, cumsum, gate_proj, up_proj, down_proj) + + EP_EXPERT_COMPUTE_MOE_ACT["quack"] = _quack_ep_moe_act_fused except ImportError: pass -# Native EP moe_act compute +# Native EP moe_act compute — adapt fused interface try: - from .native import native_ep_compute_moe_act + from .native import native_ep_compute_moe_act as _native_ep_compute_moe_act + + def _native_ep_moe_act_fused(permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size): + gate_proj = gate_up_proj[..., :intermediate_size].contiguous() + up_proj = gate_up_proj[..., intermediate_size:].contiguous() + return _native_ep_compute_moe_act(permute_tokens, cumsum, gate_proj, up_proj, down_proj) - EP_EXPERT_COMPUTE_MOE_ACT["native"] = native_ep_compute_moe_act + EP_EXPERT_COMPUTE_MOE_ACT["native"] = _native_ep_moe_act_fused except ImportError: pass diff --git a/src/xorl/models/layers/moe/experts.py b/src/xorl/models/layers/moe/experts.py index 2a578dee..99c2654e 100644 --- a/src/xorl/models/layers/moe/experts.py +++ b/src/xorl/models/layers/moe/experts.py @@ -59,6 +59,7 @@ def __init__( torch.empty(num_experts, hidden_dim, 2 * intermediate_size), requires_grad=True, ) + self.gate_up_proj._fused_gate_up = True self.down_proj = nn.Parameter( torch.empty(num_experts, intermediate_size, hidden_dim), requires_grad=True, @@ -104,8 +105,6 @@ def forward( use the unified dispatch → compute → combine path via ``_ep_forward()``. """ _moe_act = self._moe_act - gate_proj = self.gate_proj.contiguous() - up_proj = self.up_proj.contiguous() if self.moe_implementation == "eager": fn = MOE_EXPERT_BACKENDS[self.moe_implementation] @@ -113,8 +112,8 @@ def forward( return fn( hidden_states, expert_idx, - gate_proj, - up_proj, + self.gate_proj.contiguous(), + self.up_proj.contiguous(), self.down_proj, self.act_fn, ) @@ -127,7 +126,9 @@ def forward( if parallel_state.ep_enabled: return self._ep_forward(hidden_states, routing_weights, selected_experts, parallel_state) - # Local single-GPU path — select moe_act variant when available + # Local single-GPU path + gate_proj = self.gate_proj.contiguous() + up_proj = self.up_proj.contiguous() if _moe_act and self.moe_implementation in MOE_EXPERT_BACKENDS_MOE_ACT: fn = MOE_EXPERT_BACKENDS_MOE_ACT[self.moe_implementation] else: @@ -180,8 +181,6 @@ def _ep_forward( compute_fn = EP_EXPERT_COMPUTE_MOE_ACT[self.moe_implementation] else: compute_fn = EP_EXPERT_COMPUTE[self.moe_implementation] - gate_proj = self.gate_proj.contiguous() - up_proj = self.up_proj.contiguous() # Step 1: Dispatch tokens to expert-owning ranks dispatch_kwargs = self._build_dispatch_kwargs(hidden_states, routing_weights, selected_experts, parallel_state) @@ -200,14 +199,49 @@ def _ep_forward( if _FORCE_SYNC: torch.cuda.synchronize() - # Step 2: Expert computation (backend-specific GEMM only) + # Warmup: pre-compile all backward GEMM kernel variants to avoid + # first-use compilation memory spikes during training. + if not getattr(type(self), "_kernel_warmed_up", False): + from xorl.ops.group_gemm.kernel.group_gemm import group_gemm_same_mn as _warmup_mn + from xorl.ops.group_gemm.kernel.group_gemm import group_gemm_same_nk as _warmup_gemm + + _d = permute_tokens.device + _dt = permute_tokens.dtype + _H = self.gate_up_proj.shape[1] + _I = self.intermediate_size + _E = self.gate_up_proj.shape[0] + _M = _E * 2 + _cum = torch.arange(2, _M + 2, 2, dtype=torch.int32, device=_d) + + # Forward GEMM: x @ gate_up_proj + _x = torch.zeros(_M, _H, dtype=_dt, device=_d) + _w = torch.zeros(_E, _H, 2 * _I, dtype=_dt, device=_d) + _warmup_gemm(a=_x, b=_w, cumsum_M=_cum, max_M=2) + + # Backward dgrad FC1: grad_gate_up_act @ gate_up_proj^T + _g = torch.zeros(_M, 2 * _I, dtype=_dt, device=_d) + _warmup_gemm(a=_g, b=_w, cumsum_M=_cum, max_M=2, transpose_b=True) + + # Backward dgrad FC2: grad @ down_proj^T + _wd = torch.zeros(_E, _I, _H, dtype=_dt, device=_d) + _gd = torch.zeros(_M, _I, dtype=_dt, device=_d) + _warmup_gemm(a=_gd, b=_wd, cumsum_M=_cum, max_M=2, transpose_b=True) + + # Backward wgrad FC1: permute_tokens^T @ grad_gate_up_act + _c = torch.zeros(_E, _H, 2 * _I, dtype=_dt, device=_d) + _warmup_mn(a=_x, b=_g, c=_c, cumsum_K=_cum, max_K=2, transpose_a=True) + + del _x, _w, _g, _gd, _wd, _c, _cum + torch.cuda.empty_cache() + type(self)._kernel_warmed_up = True + expert_scores = getattr(ctx, "expert_scores", getattr(ctx, "permuted_scores", None)) expert_output = compute_fn( permute_tokens, cumsum, - gate_proj, - up_proj, + self.gate_up_proj, self.down_proj, + self.intermediate_size, expert_scores, ) @@ -239,9 +273,9 @@ def _ep_forward_debug(self, dispatch_fn, combine_fn, compute_fn, dispatch_kwargs expert_output = compute_fn( permute_tokens, cumsum, - self.gate_proj.contiguous(), - self.up_proj.contiguous(), + self.gate_up_proj, self.down_proj, + self.intermediate_size, expert_scores, ) ev[3].record() diff --git a/src/xorl/models/layers/moe/lora.py b/src/xorl/models/layers/moe/lora.py index 5405d2cf..83fdbca2 100644 --- a/src/xorl/models/layers/moe/lora.py +++ b/src/xorl/models/layers/moe/lora.py @@ -86,6 +86,7 @@ def __init__( torch.empty(self.num_experts, hidden_dim, 2 * intermediate_size), requires_grad=False, ) + self.gate_up_proj._fused_gate_up = True self.down_proj = nn.Parameter( torch.empty(self.num_experts, intermediate_size, hidden_dim), requires_grad=False, diff --git a/src/xorl/models/module_utils.py b/src/xorl/models/module_utils.py index eead1eed..8080d65a 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -348,6 +348,10 @@ def _load_state_dict(weights_path: str, **kwargs) -> List["StateDictIterator"]: def _try_load_state_dict(weights_path: str, **kwargs): """ Single attempt to load state dict. Returns list of iterators or None if not found. + + Rank 0 resolves all file paths (cached_file + get_checkpoint_shard_files) + and broadcasts the result to other ranks. This avoids N-way filesystem + hammering and keeps the broadcast-loading path rank-0-driven. """ cache_kwargs = {"_raise_exceptions_for_missing_entries": False, **kwargs} resolved_weight_file = cached_file(weights_path, SAFE_WEIGHTS_NAME, **cache_kwargs) @@ -1522,9 +1526,16 @@ def save_model_weights( def save_model_assets(output_dir: Union[str, "os.PathLike"], model_assets: Sequence["ModelAssets"]): + from transformers import PretrainedConfig, PreTrainedTokenizerBase + for model_asset in model_assets: if hasattr(model_asset, "save_pretrained"): - model_asset.save_pretrained(output_dir) + try: + model_asset.save_pretrained(output_dir) + except TypeError as e: + if isinstance(model_asset, (PretrainedConfig, PreTrainedTokenizerBase)): + raise + logger.warning(f"Skipping {type(model_asset).__name__}.save_pretrained(): {e}") else: logger.warning(f"Model asset {model_asset} should implement `save_pretrained`.") diff --git a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py index 0c6a8ab0..a3a7fe62 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -225,6 +225,7 @@ def __init__(self, config, moe_implementation="triton"): hidden_act=config.hidden_act, norm_topk_prob=config.norm_topk_prob, moe_implementation=moe_implementation, + train_router=getattr(config, "train_router", False), ) self.config = config self.experts.ep_dispatch = getattr(config, "_ep_dispatch", "alltoall") diff --git a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py index 889b5c72..33a61fe1 100644 --- a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py +++ b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py @@ -135,6 +135,7 @@ def __init__(self, config): hidden_act=config.hidden_act, norm_topk_prob=config.norm_topk_prob, moe_implementation="eager", + train_router=getattr(config, "train_router", False), ) self.config = config self.experts.ep_dispatch = getattr(config, "_ep_dispatch", "alltoall") @@ -155,6 +156,7 @@ def __init__(self, config): hidden_act=config.hidden_act, norm_topk_prob=config.norm_topk_prob, moe_implementation="triton", + train_router=getattr(config, "train_router", False), ) self.config = config self.experts.ep_dispatch = getattr(config, "_ep_dispatch", "alltoall") @@ -175,6 +177,7 @@ def __init__(self, config): hidden_act=config.hidden_act, norm_topk_prob=config.norm_topk_prob, moe_implementation="quack", + train_router=getattr(config, "train_router", False), ) self.config = config self.experts.ep_dispatch = getattr(config, "_ep_dispatch", "alltoall") @@ -195,6 +198,7 @@ def __init__(self, config): hidden_act=config.hidden_act, norm_topk_prob=config.norm_topk_prob, moe_implementation="native", + train_router=getattr(config, "train_router", False), ) self.config = config self.experts.ep_dispatch = getattr(config, "_ep_dispatch", "alltoall") diff --git a/src/xorl/ops/group_gemm/kernel/group_gemm.py b/src/xorl/ops/group_gemm/kernel/group_gemm.py index aaed7cef..ee49d002 100644 --- a/src/xorl/ops/group_gemm/kernel/group_gemm.py +++ b/src/xorl/ops/group_gemm/kernel/group_gemm.py @@ -4,12 +4,17 @@ multiplications with the same N,K or M,N dimensions are batched together. """ +import inspect +import os from typing import Optional import torch import triton import triton.language as tl + +_AUTOTUNE_CACHE_KW = {"cache_results": True} if "cache_results" in inspect.signature(triton.autotune).parameters else {} + from .triton_utils.activation import ( ActivationType, activation_fwd, @@ -28,6 +33,14 @@ def _get_cuda_autotune_config(): + if os.environ.get("XORL_TRITON_NO_AUTOTUNE", "0") == "1": + return [ + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "GROUP": 8}, + num_stages=3, + num_warps=8, + ), + ] return [ triton.Config( {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 64, "GROUP": 8}, @@ -50,6 +63,7 @@ def _get_cuda_autotune_config(): @triton.autotune( configs=_get_cuda_autotune_config(), key=["N", "K"], + **_AUTOTUNE_CACHE_KW, ) # @pretuned( # algo_key=algo_key_scaled(["total_M", "N", "K"], [5000, 1, 1], ["TRANSPOSE_A", "TRANSPOSE_B"]), @@ -114,7 +128,7 @@ def group_gemm_same_nk_kernel( a_ptrs = a_ptr + (offs_am[:, None] * stride_am + blk_k[None, :] * stride_ak) b_ptrs = b_ptr + (blk_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) - c_ptrs = c_ptr + N * offs_m[:, None] + 1 * offs_n[None, :] + c_ptrs = c_ptr + N * offs_am[:, None] + 1 * offs_bn[None, :] if ACCUMULATE_TO_C: c = load_with_pred_2d( @@ -124,7 +138,7 @@ def group_gemm_same_nk_kernel( offs_m[:, None] < m_size, offs_n[None, :] < N, other=0, - ) + ).to(tl.float32) else: c = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for k in range(0, tl.cdiv(K, BLOCK_K)): @@ -142,7 +156,7 @@ def group_gemm_same_nk_kernel( c = make_blocked(c, c_ptr.dtype.element_ty) if SAVE_ACTIVATION: store_with_pred_2d( - act_ptr + gtid_start * N + N * offs_m[:, None] + offs_n[None, :], + act_ptr + gtid_start * N + N * offs_am[:, None] + offs_bn[None, :], c, False, N_ALIGNED, @@ -236,6 +250,7 @@ def group_gemm_same_nk( @triton.autotune( configs=_get_cuda_autotune_config(), key=["M", "N"], + **_AUTOTUNE_CACHE_KW, ) # @pretuned( # algo_key=algo_key_scaled(["M", "N", "total_K"], [1, 1, 5000], ["TRANSPOSE_A", "TRANSPOSE_B"]), diff --git a/src/xorl/ops/moe/triton.py b/src/xorl/ops/moe/triton.py index 3249936d..f7c5f854 100644 --- a/src/xorl/ops/moe/triton.py +++ b/src/xorl/ops/moe/triton.py @@ -14,37 +14,30 @@ class TritonEPGroupGemm(torch.autograd.Function): - """EP expert MLP computation using Triton group GEMM. + """EP expert MLP with fused gate+up GEMM. Zero-copy weight references. - Pure compute kernel — no dispatch/combine logic. Tokens have already - been dispatched via all-to-all or DeepEP; this only handles the - expert SwiGLU MLP: ``down_proj(SiLU(gate_proj(x)) * up_proj(x))``. - - Moved from ``distributed/moe/moe_layer.py`` to collocate with - the sibling local-path ``TritonMoeExpertsFunction``. + Forward: single ``x @ gate_up_proj`` GEMM → split → SwiGLU → down GEMM. + Backward: fused dgrad/wgrad for gate+up (2x fewer GEMMs than split version). + ``save_for_backward`` stores the original ``gate_up_proj`` parameter + reference (zero extra memory) plus the fused activation output. """ @staticmethod - def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores=None): + def forward(ctx, permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores=None): max_M = permute_tokens.shape[0] + I = intermediate_size ctx.has_expert_scores = expert_scores is not None - gate_output = group_gemm_same_nk( - a=permute_tokens, - b=gate_proj, - cumsum_M=cumsum, - max_M=max_M, - transpose_a=False, - transpose_b=False, - ) - up_output = group_gemm_same_nk( + gate_up_output = group_gemm_same_nk( a=permute_tokens, - b=up_proj, + b=gate_up_proj, cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, ) + gate_output = gate_up_output[..., :I] + up_output = gate_up_output[..., I:] gate_activation = torch.ops.aten.silu(gate_output) gated_output = gate_activation * up_output @@ -64,26 +57,20 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_s if expert_scores is None: expert_scores = permute_tokens.new_ones(permute_tokens.shape[0]) - ctx.save_for_backward( - permute_tokens, cumsum, gate_proj, up_proj, down_proj, gate_output, up_output, expert_scores - ) + ctx.save_for_backward(permute_tokens, cumsum, gate_up_proj, down_proj, gate_up_output, expert_scores) + ctx.intermediate_size = I + return down_output @staticmethod def backward(ctx, grad_output): - ( - permute_tokens, - cumsum, - gate_proj, - up_proj, - down_proj, - gate_output, - up_output, - expert_scores, - ) = ctx.saved_tensors + permute_tokens, cumsum, gate_up_proj, down_proj, gate_up_output, expert_scores = ctx.saved_tensors + I = ctx.intermediate_size max_M = grad_output.shape[0] - # Recompute activation + gate_output = gate_up_output[..., :I] + up_output = gate_up_output[..., I:] + gate_activation = torch.ops.aten.silu(gate_output) gated_output = gate_activation * up_output expert_scores_dtype = expert_scores.dtype @@ -123,100 +110,74 @@ def backward(ctx, grad_output): # Activation backward grad_up_output = gate_activation * grad_gated_output grad_gate_activation = grad_gated_output * up_output - del grad_gated_output, gate_activation, up_output - - grad_scatter_output_2 = group_gemm_same_nk( - a=grad_up_output, - b=up_proj, - cumsum_M=cumsum, - max_M=max_M, - transpose_b=True, - ) - - grad_up_proj = None - if up_proj.requires_grad: - grad_up_proj = torch.empty_like(up_proj) - group_gemm_same_mn( - a=permute_tokens, - b=grad_up_output, - c=grad_up_proj, - cumsum_K=cumsum, - max_K=max_M, - transpose_a=True, - transpose_b=False, - ) + del grad_gated_output, gate_activation, up_output, gate_up_output grad_gate_output = torch.ops.aten.silu_backward(grad_gate_activation, gate_output) del grad_gate_activation, gate_output + # Fused dgrad FC1: cat grads → single GEMM with gate_up_proj (no .contiguous() copies) + grad_gate_up_act = torch.cat([grad_gate_output, grad_up_output], dim=-1) + del grad_gate_output, grad_up_output grad_permute_tokens = group_gemm_same_nk( - a=grad_gate_output, - b=gate_proj, + a=grad_gate_up_act, + b=gate_up_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=True, ) - grad_permute_tokens += grad_scatter_output_2 - del grad_scatter_output_2 - grad_gate_proj = None - if gate_proj.requires_grad: - grad_gate_proj = torch.empty_like(gate_proj) + # Fused wgrad FC1: single GEMM produces grad_gate_up_proj directly + grad_gate_up_proj = None + if gate_up_proj.requires_grad: + grad_gate_up_proj = torch.empty_like(gate_up_proj) group_gemm_same_mn( a=permute_tokens, - b=grad_gate_output, - c=grad_gate_proj, + b=grad_gate_up_act, + c=grad_gate_up_proj, cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, ) - del grad_gate_output + del grad_gate_up_act return ( grad_permute_tokens, None, # cumsum - grad_gate_proj, - grad_up_proj, + grad_gate_up_proj, grad_down_proj, + None, # intermediate_size grad_expert_scores, ) class TritonEPGroupGemmMoeAct(torch.autograd.Function): - """EP expert MLP with moe_act: saves only inputs + weights, recomputes - gate_output and up_output in backward (2 local GEMMs, no EP communication). + """EP expert MLP with moe_act and fused gate+up GEMM. - Memory savings: 2 * tokens * intermediate_size * dtype_size per MoE layer. + Saves only inputs + weights (no activation outputs), recomputes via + fused GEMM in backward. Zero extra memory from weight copies. """ @staticmethod - def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores=None): + def forward(ctx, permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores=None): max_M = permute_tokens.shape[0] + I = intermediate_size ctx.has_expert_scores = expert_scores is not None - gate_output = group_gemm_same_nk( + gate_up_output = group_gemm_same_nk( a=permute_tokens, - b=gate_proj, - cumsum_M=cumsum, - max_M=max_M, - transpose_a=False, - transpose_b=False, - ) - up_output = group_gemm_same_nk( - a=permute_tokens, - b=up_proj, + b=gate_up_proj, cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, ) - gate_activation = torch.ops.aten.silu(gate_output) - gated_output = gate_activation * up_output + gate_activation = torch.ops.aten.silu(gate_up_output[..., :I]) + gated_output = gate_activation * gate_up_output[..., I:] if expert_scores is not None: gated_output = gated_output * expert_scores.to(gated_output.dtype).unsqueeze(-1) - del gate_activation + del gate_activation, gate_up_output down_output = group_gemm_same_nk( a=gated_output, @@ -228,36 +189,30 @@ def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_s ) del gated_output - # moe_act: save only 5 tensors (drop gate_output, up_output) if expert_scores is None: expert_scores = permute_tokens.new_ones(permute_tokens.shape[0]) - ctx.save_for_backward(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores) + ctx.save_for_backward(permute_tokens, cumsum, gate_up_proj, down_proj, expert_scores) + ctx.intermediate_size = I return down_output @staticmethod def backward(ctx, grad_output): - permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores = ctx.saved_tensors + permute_tokens, cumsum, gate_up_proj, down_proj, expert_scores = ctx.saved_tensors + I = ctx.intermediate_size max_M = grad_output.shape[0] - # Recompute gate_output and up_output from saved inputs + weights - gate_output = group_gemm_same_nk( + # Recompute via fused GEMM + gate_up_output = group_gemm_same_nk( a=permute_tokens, - b=gate_proj, - cumsum_M=cumsum, - max_M=max_M, - transpose_a=False, - transpose_b=False, - ) - up_output = group_gemm_same_nk( - a=permute_tokens, - b=up_proj, + b=gate_up_proj, cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, ) + gate_output = gate_up_output[..., :I] + up_output = gate_up_output[..., I:] - # Rest identical to TritonEPGroupGemm.backward gate_activation = torch.ops.aten.silu(gate_output) gated_output = gate_activation * up_output expert_scores_dtype = expert_scores.dtype @@ -297,62 +252,43 @@ def backward(ctx, grad_output): # Activation backward grad_up_output = gate_activation * grad_gated_output grad_gate_activation = grad_gated_output * up_output - del grad_gated_output, gate_activation, up_output - - grad_scatter_output_2 = group_gemm_same_nk( - a=grad_up_output, - b=up_proj, - cumsum_M=cumsum, - max_M=max_M, - transpose_b=True, - ) - - grad_up_proj = None - if up_proj.requires_grad: - grad_up_proj = torch.empty_like(up_proj) - group_gemm_same_mn( - a=permute_tokens, - b=grad_up_output, - c=grad_up_proj, - cumsum_K=cumsum, - max_K=max_M, - transpose_a=True, - transpose_b=False, - ) + del grad_gated_output, gate_activation, up_output, gate_up_output grad_gate_output = torch.ops.aten.silu_backward(grad_gate_activation, gate_output) del grad_gate_activation, gate_output + # Fused dgrad FC1: cat grads → single GEMM with gate_up_proj (no .contiguous() copies) + grad_gate_up_act = torch.cat([grad_gate_output, grad_up_output], dim=-1) + del grad_gate_output, grad_up_output grad_permute_tokens = group_gemm_same_nk( - a=grad_gate_output, - b=gate_proj, + a=grad_gate_up_act, + b=gate_up_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=True, ) - grad_permute_tokens += grad_scatter_output_2 - del grad_scatter_output_2 - grad_gate_proj = None - if gate_proj.requires_grad: - grad_gate_proj = torch.empty_like(gate_proj) + # Fused wgrad FC1: single GEMM produces grad_gate_up_proj directly + grad_gate_up_proj = None + if gate_up_proj.requires_grad: + grad_gate_up_proj = torch.empty_like(gate_up_proj) group_gemm_same_mn( a=permute_tokens, - b=grad_gate_output, - c=grad_gate_proj, + b=grad_gate_up_act, + c=grad_gate_up_proj, cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, ) - del grad_gate_output + del grad_gate_up_act return ( grad_permute_tokens, None, # cumsum - grad_gate_proj, - grad_up_proj, + grad_gate_up_proj, grad_down_proj, + None, # intermediate_size grad_expert_scores, ) diff --git a/src/xorl/optim/muon.py b/src/xorl/optim/muon.py index c318795d..90ba0b01 100644 --- a/src/xorl/optim/muon.py +++ b/src/xorl/optim/muon.py @@ -156,10 +156,14 @@ def _muon_step(self, group: dict) -> None: p_local = p.data # Handle 3D+ MoE expert tensors: [num_experts, hidden, intermediate] - # Reshape to 2D for Newton-Schulz, then restore. + # For fused gate_up_proj [E, H, 2I], split into two halves for NS orig_shape = None + fused_split = None if grad_local.ndim >= 3: orig_shape = grad_local.shape + fused_gate_up_ids = group.get("_fused_gate_up_ids", set()) + if id(p) in fused_gate_up_ids: + fused_split = grad_local.shape[-1] // 2 grad_local = grad_local.reshape(-1, grad_local.shape[-1]) # --- PyTorch Muon algorithm (aligned with torch.optim._muon) --- @@ -196,8 +200,14 @@ def _muon_step(self, group: dict) -> None: # Work in buf dtype (may be bf16 even if grad is fp32) update = grad_local.to(buf.dtype).lerp(buf, momentum) if nesterov else buf - # Newton-Schulz orthogonalization (uses PyTorch's implementation) - update = _zeropower_via_newtonschulz(update, ns_coefficients, ns_steps, eps) + # Newton-Schulz orthogonalization + if fused_split is not None: + u1 = _zeropower_via_newtonschulz(update[..., :fused_split], ns_coefficients, ns_steps, eps) + u2 = _zeropower_via_newtonschulz(update[..., fused_split:], ns_coefficients, ns_steps, eps) + update = torch.cat([u1, u2], dim=-1) + del u1, u2 + else: + update = _zeropower_via_newtonschulz(update, ns_coefficients, ns_steps, eps) # LR adjustment based on the 2D shape fed to NS adjusted_lr = _adjust_lr(lr, adjust_lr_fn, grad_local.shape) diff --git a/src/xorl/optim/optimizer.py b/src/xorl/optim/optimizer.py index af14733a..72a4165e 100644 --- a/src/xorl/optim/optimizer.py +++ b/src/xorl/optim/optimizer.py @@ -499,15 +499,37 @@ def _make_muon_param_groups( decay_param_names = set(get_parameter_names(model, no_decay_modules, no_decay_params)) name_by_param = {p: n for n, p in model.named_parameters()} + fused_gate_up_ids = { + id(p) + for p, n in zip(model.parameters(), (name_by_param.get(p, "") for p in model.parameters())) + if "gate_up_proj" in name_by_param.get(p, "") + } + groups: List[Dict[str, Any]] = [] # Muon groups muon_decay = [p for p in muon_params if name_by_param.get(p) in decay_param_names] muon_no_decay = [p for p in muon_params if name_by_param.get(p) not in decay_param_names] if muon_decay: - groups.append({"params": muon_decay, "lr": muon_lr, "weight_decay": weight_decay, "use_muon": True}) + groups.append( + { + "params": muon_decay, + "lr": muon_lr, + "weight_decay": weight_decay, + "use_muon": True, + "_fused_gate_up_ids": fused_gate_up_ids, + } + ) if muon_no_decay: - groups.append({"params": muon_no_decay, "lr": muon_lr, "weight_decay": 0.0, "use_muon": True}) + groups.append( + { + "params": muon_no_decay, + "lr": muon_lr, + "weight_decay": 0.0, + "use_muon": True, + "_fused_gate_up_ids": fused_gate_up_ids, + } + ) # AdamW groups adamw_decay = [p for p in adamw_params if name_by_param.get(p) in decay_param_names] diff --git a/src/xorl/server/api_server/weights.py b/src/xorl/server/api_server/weights.py index 34605044..cb69175f 100644 --- a/src/xorl/server/api_server/weights.py +++ b/src/xorl/server/api_server/weights.py @@ -701,7 +701,7 @@ async def save_weights_for_sampler(self, request: SaveWeightsForSamplerRequest) # Determine training mode from model config model_config = self.model_configs.get(request.model_id, {}) lora_config = model_config.get("lora_config", {}) - is_lora = lora_config.get("enable_lora", False) + is_lora = lora_config.get("enable_lora", False) or "rank" in lora_config merge_lora_interval = lora_config.get("merge_lora_interval", 0) if is_lora and merge_lora_interval == 0: diff --git a/src/xorl/server/backend/remote.py b/src/xorl/server/backend/remote.py index 51a90e0d..60c18083 100644 --- a/src/xorl/server/backend/remote.py +++ b/src/xorl/server/backend/remote.py @@ -47,7 +47,7 @@ def __init__( worker_address: str = "tcp://127.0.0.1:5556", operation_timeout: float = 3600.0, connection_timeout: float = 120.0, - ack_timeout: float = 10.0, + ack_timeout: float = 120.0, ): self.worker_address = worker_address self.operation_timeout = operation_timeout diff --git a/src/xorl/server/launcher.py b/src/xorl/server/launcher.py index ca8fc2f7..8bbfc3a2 100644 --- a/src/xorl/server/launcher.py +++ b/src/xorl/server/launcher.py @@ -200,6 +200,7 @@ def run_orchestrator( rank0_worker_address=rank0_worker_address, operation_timeout=operation_timeout, connection_timeout=3600.0, # 1 hour for loading large models (235B) + EP sharding + LoRA + Triton compilation + ack_timeout=300.0, # 5 min — weight sync can block workers for 40s+ on large MoE models sample_packing_sequence_len=sample_packing_sequence_len, enable_packing=enable_packing, ) diff --git a/src/xorl/server/runner/adapters/manager.py b/src/xorl/server/runner/adapters/manager.py index 4aa1b4ea..e3b03da5 100644 --- a/src/xorl/server/runner/adapters/manager.py +++ b/src/xorl/server/runner/adapters/manager.py @@ -277,23 +277,24 @@ def prepare_forward(self, model_id: str) -> None: device_mesh = param.device_mesh sliced_data = adapter_data + skip_copy = False for dim, placement in enumerate(placements): if isinstance(placement, Shard): shard_dim = placement.dim mesh_dim_size = device_mesh.size(dim) - # Get local rank in this mesh dimension local_rank = device_mesh.get_local_rank(mesh_dim=dim) - # Compute shard size and slice total_size = sliced_data.shape[shard_dim] shard_size = (total_size + mesh_dim_size - 1) // mesh_dim_size start = local_rank * shard_size end = min(start + shard_size, total_size) length = max(end - start, 0) - if length > 0: - # Slice along the shard dimension - sliced_data = sliced_data.narrow(shard_dim, start, length) + if start >= total_size or length == 0: + skip_copy = True + break + sliced_data = sliced_data.narrow(shard_dim, start, length) - local_tensor.copy_(sliced_data) + if not skip_copy and local_tensor.numel() > 0: + local_tensor.copy_(sliced_data) else: # Regular tensor: direct copy param.data.copy_(adapter_data) diff --git a/src/xorl/server/runner/checkpoint/manager.py b/src/xorl/server/runner/checkpoint/manager.py index 486ae1e1..f411a31d 100644 --- a/src/xorl/server/runner/checkpoint/manager.py +++ b/src/xorl/server/runner/checkpoint/manager.py @@ -111,6 +111,27 @@ def _get_lora_save_config(self): # Adapter save / load (multi-tenancy LoRA) # ------------------------------------------------------------------ + def _gather_adapter_lora_params(self, model_id: str) -> Dict[str, torch.Tensor]: + """Gather LoRA params from adapter manager with EP support. + + Fast alternative to get_lora_state_dict(): uses the adapter manager's + stored full-tensor params directly (no FSDP unshard needed). Only needs + a dist.gather for EP-sharded expert LoRA params. + + Returns complete state dict on rank 0, empty dict on other ranks. + + No EP gather needed: ``AdapterState.lora_params`` already stores + tensors at global shape (see ``register_adapter``). + """ + lora_state_dict: Dict[str, torch.Tensor] = {} + + if self.rank == 0: + state = self._adapter_manager.adapters[model_id] + for name, param in state.lora_params.items(): + lora_state_dict[name] = param.data.cpu() + + return lora_state_dict + def _save_lora_weights(self, save_path: str, model_id: str) -> None: """ Core LoRA saving logic: activate adapter, gather weights, write PEFT checkpoint. @@ -127,8 +148,13 @@ def _save_lora_weights(self, save_path: str, model_id: str) -> None: if self._adapter_manager is not None: self._adapter_manager.switch_adapter(model_id, auto_register=True) - # EP+FSDP2-aware LoRA weight gathering (collective operation) - lora_state_dict = get_lora_state_dict(self.model) + # Use fast adapter-manager path when available (avoids FSDP unshard) + if self._adapter_manager is not None and model_id in self._adapter_manager.adapters: + logger.info(f"Rank {self.rank}: Using fast adapter-manager LoRA save path") + lora_state_dict = self._gather_adapter_lora_params(model_id) + else: + # Fallback: EP+FSDP2-aware LoRA weight gathering (collective operation) + lora_state_dict = get_lora_state_dict(self.model) # Only rank 0 writes files if self.rank == 0: diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index 51280374..bf9f6797 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -163,7 +163,14 @@ class ModelRunner: "_original_position_ids", "rollout_logprobs", }, - "policy_loss": {"labels", "target_tokens", "logprobs", "advantages", "rollout_logprobs"}, + "policy_loss": { + "labels", + "target_tokens", + "logprobs", + "advantages", + "_original_position_ids", + "rollout_logprobs", + }, } def __init__( @@ -231,6 +238,10 @@ def __init__( get_torch_device().set_device(f"{get_device_type()}:{local_rank}") helper.set_seed(self.train_config.get("seed", 42), False) + # Disable TF32 and BF16 reduced-precision accumulation for + # consistent numerics across parallelism strategies. + helper.enable_high_precision_for_bf16() + # Initialize distributed parallel state self._init_parallel_state() @@ -1247,6 +1258,15 @@ def forward_backward( """ self._check_not_sleeping("forward_backward") + # Defragment GPU memory at the top of every forward_backward call. + # After weight-sync + optim_step from the previous step the CUDA + # allocator can have many small free blocks; without this, CUBLAS + # handle creation or Triton autotuner workspace allocs can fail. + import gc + + gc.collect() + torch.cuda.empty_cache() + # Validate single-tenant mode for full-weights training self._validate_single_tenant(model_id) @@ -1285,7 +1305,15 @@ def forward_backward( model_inputs = { k: v for k, v in mb.items() - if k not in ["labels", "target_tokens", "logprobs", "advantages", "rollout_logprobs"] + if k + not in [ + "labels", + "target_tokens", + "logprobs", + "advantages", + "_original_position_ids", + "rollout_logprobs", + ] } with self.model_fwd_context: @@ -1333,6 +1361,11 @@ def forward_backward( logger.info("Reference logprobs computed, replacing old_logprobs") + # Reclaim fragmented GPU memory before the main forward-backward pass. + # Without this, the Triton autotuner's workspace allocations can OOM + # on memory-tight configs (e.g. EP=32, 16 experts/GPU). + torch.cuda.empty_cache() + # R3 (Rollout Routing Replay): Pre-populate routing replay from inference data r3_enabled = self._routing_handler.setup(micro_batches, routed_experts, routed_expert_logits) diff --git a/src/xorl/server/runner/runner_dispatcher.py b/src/xorl/server/runner/runner_dispatcher.py index f3c7a394..9d0f2c6c 100644 --- a/src/xorl/server/runner/runner_dispatcher.py +++ b/src/xorl/server/runner/runner_dispatcher.py @@ -821,6 +821,22 @@ def _shard_and_slice_batches( sharded_batches = [] for i, batch in enumerate(my_batches): try: + # Generate position_ids before sharding if not present. + # Under Ulysses SP, position_ids must be FULL-LENGTH so that + # RoPE cos/sin are computed for global positions and then sliced + # to the local shard inside prepare_position_embeddings. + # Without this, the model falls back to arange(S_local) on each + # rank, producing wrong position encodings for ranks > 0. + if "position_ids" not in batch and "input_ids" in batch: + input_ids = batch["input_ids"] + seq_len = input_ids.shape[-1] if isinstance(input_ids, torch.Tensor) else len(input_ids) + batch_size = ( + input_ids.shape[0] if isinstance(input_ids, torch.Tensor) and input_ids.ndim >= 2 else 1 + ) + batch["position_ids"] = ( + torch.arange(seq_len, dtype=torch.long).unsqueeze(0).expand(batch_size, -1) + ) + # Store original position_ids for unpacking per-token outputs later if "position_ids" in batch: original_pos_ids = batch["position_ids"] diff --git a/src/xorl/server/runner/utils/routing_replay_handler.py b/src/xorl/server/runner/utils/routing_replay_handler.py index b66a5648..407d95fb 100644 --- a/src/xorl/server/runner/utils/routing_replay_handler.py +++ b/src/xorl/server/runner/utils/routing_replay_handler.py @@ -64,6 +64,17 @@ class RoutingReplayHandler: def __init__(self, model: nn.Module) -> None: self.model = model self._moe_blocks: Optional[List[nn.Module]] = None + self._model_topk: Optional[int] = self._extract_topk(model) + + @staticmethod + def _extract_topk(model: nn.Module) -> Optional[int]: + """Extract num_experts_per_tok from the model config, if available.""" + config = getattr(model, "config", None) + if config is not None: + topk = getattr(config, "num_experts_per_tok", None) + if topk is not None: + return int(topk) + return None def get_moe_blocks(self) -> List[nn.Module]: """ @@ -173,15 +184,28 @@ def decode_routed_expert_logits_item( """Decode a single routed_expert_logits item (float32 routing weights).""" return self._decode_routing_array(item, num_moe_layers, np.float32, "R3 weights") - @staticmethod - def _infer_shape(arr: np.ndarray, num_moe_layers: int) -> Optional[np.ndarray]: - """Infer [num_tokens, num_layers, topk] shape from flat array.""" + def _infer_shape(self, arr: np.ndarray, num_moe_layers: int) -> Optional[np.ndarray]: + """Infer [num_tokens, num_layers, topk] shape from flat array. + + Tries the model's known ``num_experts_per_tok`` first, then falls back + to common values. + """ total_elements = len(arr) - for topk in [8, 4, 2, 1]: + candidates = [10, 8, 6, 4, 2, 1, 16] + if self._model_topk is not None and self._model_topk not in candidates: + candidates.insert(0, self._model_topk) + elif self._model_topk is not None: + candidates.remove(self._model_topk) + candidates.insert(0, self._model_topk) + + for topk in candidates: if total_elements % (num_moe_layers * topk) == 0: num_tokens = total_elements // (num_moe_layers * topk) return arr.reshape(num_tokens, num_moe_layers, topk) - logger.warning(f"R3: Cannot infer shape for {total_elements} elements with {num_moe_layers} layers") + logger.warning( + f"R3: Cannot infer shape for {total_elements} elements " + f"with {num_moe_layers} layers (tried topk={candidates})" + ) return None def fill_routing_replay( diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index 1cea27ea..730f1b2d 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -147,7 +147,9 @@ def _timed(phase_name, fn): def _bootstrap(self) -> None: """Initialize distributed, device, seed, parallel state.""" args = self.args - dist.init_process_group(backend=get_nccl_backend()) + from datetime import timedelta + + dist.init_process_group(backend=get_nccl_backend(), timeout=timedelta(minutes=60)) logger.info(f"Process rank: {args.train.global_rank}, world size: {args.train.world_size}") logger.info_rank0(json.dumps(asdict(args), indent=2)) diff --git a/src/xorl/utils/helper.py b/src/xorl/utils/helper.py index 115db65a..f0df2ef3 100644 --- a/src/xorl/utils/helper.py +++ b/src/xorl/utils/helper.py @@ -293,7 +293,7 @@ def enable_high_precision_for_bf16(): torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False - if IS_NPU_AVAILABLE: + if hasattr(torch, "npu") and hasattr(torch.npu, "is_available") and torch.npu.is_available(): torch.npu.matmul.allow_tf32 = False torch.npu.matmul.allow_bf16_reduced_precision_reduction = False From 3c809a5ac1a707e98fdad2cdb714c18570a1767e Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Wed, 8 Apr 2026 09:45:53 -0700 Subject: [PATCH 18/41] Enforce import-outside-top-level with targeted lazy imports (#74) * Enforce import-outside-top-level with targeted lazy imports * Use real imports for PLC0415 fixes * Move safe PLC0415 imports to module scope * Apply ruff formatting * Make pyiceberg warning import non-optional * Revert "Make pyiceberg warning import non-optional" This reverts commit 833fc321cdc917349f693243e8994f8b4a06ee36. * Fix Ruff import spacing in direct_train (cherry picked from commit c02dbf69e91f6bcb277e13915b8d3ec939cfc84c) --- pyproject.toml | 9 ++- src/xorl/arguments.py | 3 +- src/xorl/checkpoint/checkpointer.py | 5 +- src/xorl/cli/direct_train.py | 50 +++++++------- .../data/collators/packing_concat_collator.py | 2 - src/xorl/data/data_loader.py | 2 +- src/xorl/data/prepare/prepare_datasets.py | 4 +- src/xorl/data/prepare/shared.py | 67 +++++++++++-------- src/xorl/distributed/moe/__init__.py | 4 +- src/xorl/distributed/parallel_plan.py | 3 +- src/xorl/distributed/parallel_state.py | 18 ++--- src/xorl/distributed/pipeline_parallel.py | 6 +- .../distributed/sequence_parallel/__init__.py | 2 +- .../distributed/sequence_parallel/strategy.py | 12 ++-- src/xorl/distributed/torch_parallelize.py | 9 +-- src/xorl/lora/__init__.py | 2 +- src/xorl/lora/mapping.py | 3 +- src/xorl/lora/utils.py | 5 +- src/xorl/models/auto.py | 8 +-- .../models/checkpoint_handlers/buffers.py | 48 +++++-------- src/xorl/models/layers/activations.py | 5 +- src/xorl/models/layers/moe/experts.py | 16 +++-- src/xorl/models/layers/moe/lora.py | 7 +- src/xorl/models/layers/moe/moe_block.py | 2 +- src/xorl/models/module_utils.py | 20 ++---- .../transformers/qwen3/checkpoint_handler.py | 3 +- .../qwen3_5/checkpoint_handler.py | 3 +- .../qwen3_5_moe/checkpoint_handler.py | 3 +- .../qwen3_moe/checkpoint_handler.py | 3 +- src/xorl/ops/group_gemm/utils/device.py | 4 +- src/xorl/ops/moe/lora.py | 2 +- src/xorl/ops/quack/cute_dsl_ptxas.py | 2 +- src/xorl/ops/quack/gemm_sm100.py | 6 +- src/xorl/ops/quack/utils.py | 4 +- src/xorl/qlora/modules/linear.py | 22 +++--- src/xorl/qlora/modules/moe_experts.py | 13 ++-- src/xorl/qlora/utils.py | 35 ++++------ .../server/api_server/inference_endpoints.py | 7 +- src/xorl/server/launcher.py | 16 ++--- src/xorl/server/orchestrator/packing.py | 1 - src/xorl/server/server_arguments.py | 6 +- src/xorl/server/utils/network.py | 5 +- .../server/weight_sync/backends/__init__.py | 2 +- .../weight_sync/backends/nccl_broadcast.py | 16 ++--- src/xorl/server/weight_sync/handler.py | 53 +++++---------- src/xorl/trainers/model_builder.py | 43 +++++------- src/xorl/trainers/trainer.py | 43 ++++++------ src/xorl/trainers/training_utils.py | 12 ++-- src/xorl/utils/helper.py | 22 ++++-- src/xorl/utils/import_utils.py | 3 +- src/xorl/utils/recompute_utils.py | 19 ------ tests/data/prepare/test_shared.py | 5 +- tests/data/test_data_loader.py | 4 +- tests/data/test_data_loader_distributed.py | 7 +- tests/distributed/test_deepep_correctness.py | 4 +- .../test_ep_lora_weight_slicing.py | 4 +- tests/distributed/test_parallel_state.py | 4 +- tests/distributed/test_tensor_parallel.py | 7 +- tests/e2e/e2e_utils.py | 19 ++++-- tests/e2e/qwen3_8b/test_tflops_threshold.py | 14 ++-- tests/e2e/server_utils.py | 20 ++++-- tests/models/test_moe_experts_lora.py | 4 +- tests/models/test_moe_routing_cache.py | 8 +-- tests/models/test_moe_routing_replay.py | 3 +- .../test_moe_weight_loading_integration.py | 4 +- tests/models/test_qwen3_moe_fused_lora.py | 3 +- tests/ops/test_attention.py | 2 +- tests/ops/test_eager_vs_native_moe.py | 6 +- tests/ops/test_group_gemm.py | 6 +- tests/ops/test_moe_act.py | 11 +-- tests/ops/test_moe_gkn_format.py | 12 ++-- tests/ops/test_moe_ops.py | 6 +- tests/ops/test_moe_torch_compile.py | 19 +++--- tests/qlora/test_qlora.py | 1 - .../api_server/test_checkpoint_paths.py | 6 +- .../orchestrator/test_cu_seqlens_alignment.py | 10 ++- .../server/orchestrator/test_orchestrator.py | 2 +- .../test_orchestrator_client_communication.py | 4 +- .../weight_sync/test_pp_nccl_transfer.py | 3 +- 79 files changed, 388 insertions(+), 470 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 002cccae..0cd2f9c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,7 +80,7 @@ line-length = 120 [tool.ruff.lint] ignore = ["C901", "E501", "E741", "W605", "C408"] -select = ["C", "E", "F", "I", "W"] +select = ["C", "E", "F", "I", "W", "PLC0415", "TID253"] [tool.ruff.lint.per-file-ignores] "__init__.py" = ["E402", "F401", "F403", "F811"] @@ -90,6 +90,13 @@ lines-after-imports = 2 known-first-party = ["xorl"] known-third-party = ["torch", "transformers", "wandb"] +[tool.ruff.lint.flake8-tidy-imports] +banned-module-level-imports = [ + "wandb", + "flash_attn", + "liger_kernel.transformers.fused_linear_cross_entropy", +] + [tool.pytest.ini_options] addopts = "-v" diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 1a7e4280..d3d44933 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -25,6 +25,7 @@ import torch import yaml +from .checkpoint_utils import get_checkpoint_path from .utils import logging @@ -950,8 +951,6 @@ def __post_init__(self): assert self.init_device == "meta", "Please use init_device: meta for FSDP2 training" if self.load_checkpoint_path == "auto": - from .checkpoint_utils import get_checkpoint_path - self.load_checkpoint_path = get_checkpoint_path( output_dir=self.output_dir, is_local_rank0=self.local_rank == 0, diff --git a/src/xorl/checkpoint/checkpointer.py b/src/xorl/checkpoint/checkpointer.py index 0536a896..c91646c6 100644 --- a/src/xorl/checkpoint/checkpointer.py +++ b/src/xorl/checkpoint/checkpointer.py @@ -1,3 +1,4 @@ +import gc import json import os from abc import ABC, abstractmethod @@ -6,6 +7,7 @@ import torch import torch.distributed as dist from torch.distributed._tensor import DeviceMesh, DTensor, Shard +from torch.distributed.checkpoint.state_dict import StateDictOptions from ..distributed.parallel_state import get_parallel_state from ..utils.checkpoint_utils import _GLOBAL_STEP_PREFIX @@ -224,7 +226,6 @@ def load_state_dict(self, state_dict): parameters (e.g., loading from a base model checkpoint into a LoRA-enabled model). Missing parameters (like lora_A, lora_B) will retain their initialized values. """ - from torch.distributed.checkpoint.state_dict import StateDictOptions model_state_dict = state_dict if self.should_ep_aware: @@ -537,8 +538,6 @@ def save( save_state["optimizer"].optimizer = None del save_state - import gc - gc.collect() gc.collect() # Second pass for cyclic references torch.cuda.empty_cache() diff --git a/src/xorl/cli/direct_train.py b/src/xorl/cli/direct_train.py index e29ca888..292c391d 100644 --- a/src/xorl/cli/direct_train.py +++ b/src/xorl/cli/direct_train.py @@ -1,7 +1,26 @@ # ruff: noqa: E402 import os - +from collections import deque + +import torch.distributed.tensor._random +import torch.nn.functional as F +from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict + +from xorl.distributed.pipeline_parallel import build_pipeline_schedule +from xorl.lora.utils import inject_lora_into_model, save_lora_checkpoint +from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules +from xorl.models.layers.moe.routing_replay import RoutingReplay, set_replay_stage +from xorl.qlora import ( + detect_prequantized_block_fp8, + detect_prequantized_nvfp4, + inject_qlora_into_model, + maybe_load_and_quantize_moe_qlora, + maybe_load_prequantized_qlora, + maybe_quantize_qlora, +) +from xorl.qlora.modules.linear import prefetch_aqn_noise +from xorl.qlora.utils import _deregister_qlora_weights_from_fsdp from xorl.utils.compile_cache import configure_rank_local_compile_caches @@ -83,7 +102,7 @@ def main(): if args.train.global_rank == 0: save_args(args, args.train.output_dir) if args.train.use_wandb: - import wandb + import wandb # noqa: PLC0415 wandb.init( project=args.train.wandb_project, @@ -119,7 +138,7 @@ def main(): ) ) if args.train.use_wandb: - import wandb + import wandb # noqa: PLC0415 wandb.config.update( { @@ -152,8 +171,6 @@ def main(): # Same approach as torchtitan (see torchtitan/distributed/utils.py:set_determinism). ps = get_parallel_state() if ps.device_mesh is not None: - import torch.distributed.tensor._random - torch.distributed.tensor._random.manual_seed(args.train.seed, ps.device_mesh) logger.info_rank0("Prepare data") @@ -226,8 +243,6 @@ def main(): checkpoint_quant_format = None exclude_modules = set() if args.lora.enable_qlora: - from xorl.qlora import detect_prequantized_block_fp8, detect_prequantized_nvfp4, inject_qlora_into_model - if detect_prequantized_nvfp4(args.model.model_path): is_prequantized = True checkpoint_quant_format = "nvfp4" @@ -240,8 +255,6 @@ def main(): exclude_modules = set(args.lora.exclude_modules) logger.info_rank0(f"Using user-specified exclude_modules: {exclude_modules}") elif is_prequantized: - from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules - exclude_modules = get_prequantized_exclude_modules(args.model.model_path) if exclude_modules: logger.info_rank0( @@ -272,8 +285,6 @@ def main(): model._qlora_exclude_modules = exclude_modules helper.print_device_mem_info("VRAM usage after QLoRA injection") elif args.lora.enable_lora: - from xorl.lora.utils import inject_lora_into_model - inject_lora_into_model( model, r=args.lora.lora_rank, @@ -332,16 +343,11 @@ def main(): # For pre-quantized checkpoints, skip quantization — load packed weights directly. if args.lora.enable_qlora: if is_prequantized: - from xorl.qlora import maybe_load_prequantized_qlora - logger.info("Starting pre-quantized NVFP4 weight loading...") helper.print_device_mem_info("VRAM before pre-quantized loading") maybe_load_prequantized_qlora(model, args.model.model_path) logger.info("Done pre-quantized weight loading, freezing non-LoRA params...") else: - from xorl.qlora import maybe_load_and_quantize_moe_qlora, maybe_quantize_qlora - from xorl.qlora.utils import _deregister_qlora_weights_from_fsdp - logger.info("Starting maybe_quantize_qlora...") helper.print_device_mem_info("VRAM before QLoRA quantization") maybe_quantize_qlora(model) @@ -438,9 +444,6 @@ def main(): pp_schedule = None pp_context = {} # mutable container for per-step state used by pp_loss_fn if pp_enabled: - import torch.nn.functional as F - - from xorl.distributed.pipeline_parallel import build_pipeline_schedule @torch.compile def _pp_ce_loss(pred, labels, ntokens): @@ -530,15 +533,12 @@ def pp_loss_fn(pred, labels): # Runs async — overlaps with data prep below, so forward only # pays the cheap addcmul cost per layer. if args.lora.enable_aqn: - from xorl.qlora.modules.linear import prefetch_aqn_noise - prefetch_aqn_noise(model) # Routing replay stage switching for MoE checkpoint determinism. # Only needed with EP — without EP, expert compute has fixed output # shapes regardless of routing, so checkpoint recompute is safe. # See models/layers/moe/routing_replay.py for lifecycle docs. - from xorl.models.layers.moe.routing_replay import RoutingReplay, set_replay_stage use_routing_replay = ps.ep_size > 1 and args.train.moe_recomputed @@ -560,7 +560,6 @@ def pp_loss_fn(pred, labels): # - position_ids: full-length (not SP-sliced) for correct per-document RoPE # - cu_seq_lens/max_length: flash-attention varlen kwargs for document boundaries # Each _pp_forward call pops one entry from the deque. - from collections import deque _PP_FA_KEYS = ("cu_seq_lens_q", "cu_seq_lens_k", "max_length_q", "max_length_k") pp_metadata_list = [] @@ -729,7 +728,7 @@ def pp_loss_fn(pred, labels): if args.train.global_rank == 0: if args.train.use_wandb and global_step % args.train.wandb_log_interval == 0: - import wandb + import wandb # noqa: PLC0415 train_metrics.update( { @@ -792,8 +791,6 @@ def pp_loss_fn(pred, labels): hf_model_state_dict = None if args.train.save_hf_weights and not save_peft_adapter: - from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict - logger.info_rank0("Gathering full model state dict for HF checkpoint via NCCL...") hf_model_state_dict = get_model_state_dict(model, options=StateDictOptions(full_state_dict=True)) @@ -806,7 +803,6 @@ def pp_loss_fn(pred, labels): hf_weights_path = os.path.join(args.train.output_dir, f"global_step_{global_step}", "hf_ckpt") if save_peft_adapter: # Save PEFT adapter format (LoRA-only, base weights unchanged) - from xorl.lora.utils import save_lora_checkpoint save_lora_checkpoint( model, diff --git a/src/xorl/data/collators/packing_concat_collator.py b/src/xorl/data/collators/packing_concat_collator.py index 602554c7..7c54041c 100644 --- a/src/xorl/data/collators/packing_concat_collator.py +++ b/src/xorl/data/collators/packing_concat_collator.py @@ -73,8 +73,6 @@ def __call__(self, features: Sequence[Dict[str, "torch.Tensor"]]) -> Dict[str, " if not features: raise ValueError("PackingConcatCollator received empty features list") - import logging - logger = logging.getLogger(__name__) # Input should be a flat list of dicts from FlattenCollator diff --git a/src/xorl/data/data_loader.py b/src/xorl/data/data_loader.py index 00c1e090..ba2e3ff6 100644 --- a/src/xorl/data/data_loader.py +++ b/src/xorl/data/data_loader.py @@ -12,6 +12,7 @@ PackingConcatCollator, ShiftTokensCollator, TextSequenceShardCollator, + ToTensorCollator, ) @@ -176,7 +177,6 @@ def __init__( def _build_default_collator_list(self) -> List[DataCollator]: """Build the default collator pipeline.""" - from .collators import ToTensorCollator collators = [ ToTensorCollator(), # 1. Convert to tensors diff --git a/src/xorl/data/prepare/prepare_datasets.py b/src/xorl/data/prepare/prepare_datasets.py index f1f5533c..85c50c93 100644 --- a/src/xorl/data/prepare/prepare_datasets.py +++ b/src/xorl/data/prepare/prepare_datasets.py @@ -2,6 +2,8 @@ from typing import Literal +import numpy as np +import pyarrow as pa import torch.distributed as dist from datasets import ( Dataset as HFDataset, @@ -50,8 +52,6 @@ def _create_dummy_dataset(seq_len: int, num_samples: int = 4096, seed: int = 42, 0 is the EOD marker. All ``num_samples`` examples are intentionally identical so local benchmarks can show clean loss convergence. """ - import numpy as np - import pyarrow as pa EOD = 0 VOCAB_SIZE = vocab_size diff --git a/src/xorl/data/prepare/shared.py b/src/xorl/data/prepare/shared.py index 5a641bff..41e3f978 100644 --- a/src/xorl/data/prepare/shared.py +++ b/src/xorl/data/prepare/shared.py @@ -4,6 +4,7 @@ import functools import os +from dataclasses import replace from pathlib import Path from typing import TYPE_CHECKING, Any, Generator @@ -24,6 +25,27 @@ RevisionNotFoundError, ) + +try: + import adlfs +except ImportError: + adlfs = None + +try: + import gcsfs +except ImportError: + gcsfs = None + +try: + import ocifs +except ImportError: + ocifs = None + +try: + import s3fs +except ImportError: + s3fs = None + from ...arguments import Arguments, DatasetConfig from ...data.prepare.constants import DEFAULT_DATASET_PREPARED_PATH from ...data.prepare.hash import generate_split_fingerprints @@ -79,7 +101,6 @@ def datasets_with_name_generator( Yields: Individual dataset configurations, expanded as needed for names or shards. """ - from dataclasses import replace for config in dataset_configs: if config.name and isinstance(config.name, list): @@ -186,40 +207,28 @@ def _get_remote_filesystem( ) -> tuple[S3FileSystem | GCSFileSystem | AzureBlobFileSystem | OCIFileSystem | None, dict]: """Get the appropriate filesystem for a remote path.""" if path.startswith("s3://"): - try: - import s3fs - - storage_options = {"anon": False} - return s3fs.S3FileSystem(**storage_options), storage_options - except ImportError as exc: - raise ImportError("s3:// paths require s3fs to be installed") from exc + if s3fs is None: + raise ImportError("s3:// paths require s3fs to be installed") + storage_options = {"anon": False} + return s3fs.S3FileSystem(**storage_options), storage_options elif path.startswith(("gs://", "gcs://")): - try: - import gcsfs - - storage_options = {"token": None} # type: ignore - return gcsfs.GCSFileSystem(**storage_options), storage_options - except ImportError as exc: - raise ImportError("gs:// or gcs:// paths require gcsfs to be installed") from exc + if gcsfs is None: + raise ImportError("gs:// or gcs:// paths require gcsfs to be installed") + storage_options = {"token": None} # type: ignore + return gcsfs.GCSFileSystem(**storage_options), storage_options elif path.startswith(("adl://", "abfs://", "az://")): - try: - import adlfs - - storage_options = {"anon": False} - return adlfs.AzureBlobFileSystem(**storage_options), storage_options - except ImportError as exc: - raise ImportError("adl:// or abfs:// paths require adlfs to be installed") from exc + if adlfs is None: + raise ImportError("adl:// or abfs:// paths require adlfs to be installed") + storage_options = {"anon": False} + return adlfs.AzureBlobFileSystem(**storage_options), storage_options elif path.startswith("oci://"): - try: - import ocifs - - storage_options = {} - return ocifs.OCIFileSystem(**storage_options), storage_options - except ImportError as exc: - raise ImportError("oci:// paths require ocifs to be installed") from exc + if ocifs is None: + raise ImportError("oci:// paths require ocifs to be installed") + storage_options = {} + return ocifs.OCIFileSystem(**storage_options), storage_options return None, {} diff --git a/src/xorl/distributed/moe/__init__.py b/src/xorl/distributed/moe/__init__.py index 17c8ccda..6159359d 100644 --- a/src/xorl/distributed/moe/__init__.py +++ b/src/xorl/distributed/moe/__init__.py @@ -9,7 +9,7 @@ def __getattr__(name): "alltoall_pre_dispatch", "alltoall_post_combine", ): - from .alltoall import ( + from .alltoall import ( # noqa: PLC0415 AllToAllDispatchContext, alltoall_post_combine, alltoall_pre_dispatch, @@ -37,7 +37,7 @@ def __getattr__(name): "get_default_buffer", "destroy_default_buffer", ): - from .deepep import ( + from .deepep import ( # noqa: PLC0415 DEEPEP_AVAILABLE, DeepEPBuffer, destroy_default_buffer, diff --git a/src/xorl/distributed/parallel_plan.py b/src/xorl/distributed/parallel_plan.py index 596d6f97..cd953a7a 100644 --- a/src/xorl/distributed/parallel_plan.py +++ b/src/xorl/distributed/parallel_plan.py @@ -6,6 +6,7 @@ from torch.distributed._tensor import DeviceMesh, DTensor, Replicate, Shard from ..utils import logging +from .parallel_state import get_parallel_state from .utils import check_fqn_match, get_module_from_path, set_module_from_path @@ -151,8 +152,6 @@ def _slice_expert_tensor_for_ep( ) -> "torch.Tensor": """Slice expert tensor for expert parallelism.""" try: - from .parallel_state import get_parallel_state - parallel_state = get_parallel_state() # Check if we need to slice based on tensor vs target shape mismatch diff --git a/src/xorl/distributed/parallel_state.py b/src/xorl/distributed/parallel_state.py index aa07e4b8..e005c530 100644 --- a/src/xorl/distributed/parallel_state.py +++ b/src/xorl/distributed/parallel_state.py @@ -96,7 +96,7 @@ def __post_init__(self): ) if self.cp_enabled: - from ..distributed.sequence_parallel import ( + from ..distributed.sequence_parallel import ( # noqa: PLC0415 init_sequence_parallel, set_data_parallel_group, set_ringattn_group, @@ -147,7 +147,7 @@ def dp_group(self) -> Optional["ProcessGroup"]: return self.device_mesh.get_group(self._resolve_mesh_name("dp")) if self.cp_enabled: - from ..distributed.sequence_parallel import get_data_parallel_group + from ..distributed.sequence_parallel import get_data_parallel_group # noqa: PLC0415 return get_data_parallel_group() @@ -159,7 +159,7 @@ def dp_rank(self) -> int: return self.device_mesh.get_local_rank(self._resolve_mesh_name("dp")) if self.cp_enabled: - from ..distributed.sequence_parallel import get_data_parallel_rank + from ..distributed.sequence_parallel import get_data_parallel_rank # noqa: PLC0415 return get_data_parallel_rank() @@ -412,7 +412,7 @@ def sp_group(self) -> Optional["ProcessGroup"]: if self.device_mesh is not None: return self.device_mesh.get_group(self._resolve_mesh_name("sp")) - from .sequence_parallel import get_unified_sequence_parallel_group + from .sequence_parallel import get_unified_sequence_parallel_group # noqa: PLC0415 return get_unified_sequence_parallel_group() @@ -424,7 +424,7 @@ def cp_rank(self) -> int: if self.device_mesh is not None: return self.device_mesh.get_local_rank(self._resolve_mesh_name("sp")) - from .sequence_parallel import get_unified_sequence_parallel_rank + from .sequence_parallel import get_unified_sequence_parallel_rank # noqa: PLC0415 return get_unified_sequence_parallel_rank() @@ -444,7 +444,7 @@ def ulysses_group(self) -> Optional["ProcessGroup"]: if self.device_mesh is not None: return self.device_mesh.get_group("ulysses") - from .sequence_parallel import get_ulysses_sequence_parallel_group + from .sequence_parallel import get_ulysses_sequence_parallel_group # noqa: PLC0415 return get_ulysses_sequence_parallel_group() @@ -456,7 +456,7 @@ def ulysses_rank(self) -> int: if self.device_mesh is not None: return self.device_mesh.get_local_rank("ulysses") - from .sequence_parallel import get_ulysses_sequence_parallel_rank + from .sequence_parallel import get_ulysses_sequence_parallel_rank # noqa: PLC0415 return get_ulysses_sequence_parallel_rank() @@ -472,7 +472,7 @@ def ringattn_group(self) -> Optional["ProcessGroup"]: if self.device_mesh is not None: return self.device_mesh.get_group("ringattn") - from .sequence_parallel import get_ringattn_group + from .sequence_parallel import get_ringattn_group # noqa: PLC0415 return get_ringattn_group() @@ -484,7 +484,7 @@ def ringattn_rank(self) -> int: if self.device_mesh is not None: return self.device_mesh.get_local_rank("ringattn") - from .sequence_parallel import get_ringattn_rank + from .sequence_parallel import get_ringattn_rank # noqa: PLC0415 return get_ringattn_rank() diff --git a/src/xorl/distributed/pipeline_parallel.py b/src/xorl/distributed/pipeline_parallel.py index ea96cd11..baf0a0b3 100644 --- a/src/xorl/distributed/pipeline_parallel.py +++ b/src/xorl/distributed/pipeline_parallel.py @@ -33,7 +33,10 @@ get_schedule_class, ) +from xorl.models.layers.moe.routing_replay import get_replay_stage, is_r3_mode, set_replay_stage + from ..utils import logging +from .parallel_state import get_parallel_state logger = logging.get_logger(__name__) @@ -187,9 +190,6 @@ def _pp_forward(self, x): When the queue is not set (e.g. unit tests), falls back to generating sequential position_ids scaled by cp_size for RoPE cache correctness. """ - from xorl.models.layers.moe.routing_replay import get_replay_stage, is_r3_mode, set_replay_stage - - from .parallel_state import get_parallel_state ps = get_parallel_state() diff --git a/src/xorl/distributed/sequence_parallel/__init__.py b/src/xorl/distributed/sequence_parallel/__init__.py index 7688782e..e81d5b35 100644 --- a/src/xorl/distributed/sequence_parallel/__init__.py +++ b/src/xorl/distributed/sequence_parallel/__init__.py @@ -48,7 +48,7 @@ def __getattr__(name): if name == "ring_flash_attention_forward": - from .ring_attention import ring_flash_attention_forward + from .ring_attention import ring_flash_attention_forward # noqa: PLC0415 return ring_flash_attention_forward raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/xorl/distributed/sequence_parallel/strategy.py b/src/xorl/distributed/sequence_parallel/strategy.py index 360e1552..1ab25279 100644 --- a/src/xorl/distributed/sequence_parallel/strategy.py +++ b/src/xorl/distributed/sequence_parallel/strategy.py @@ -19,6 +19,8 @@ import torch import torch.distributed as dist +from ...models.layers.attention.utils import repeat_kv +from ...models.layers.rope import apply_rotary_pos_emb from .async_ulysses import async_ulysses_output_projection, async_ulysses_qkv_projection from .data import slice_position_embedding from .ulysses import gather_heads_scatter_seq, gather_seq_scatter_heads @@ -147,8 +149,6 @@ def __init__(self, group, ulysses_size: int): self.ulysses_size = ulysses_size def project_qkv(self, module, hidden_states, position_embeddings): - from ...models.layers.attention.utils import repeat_kv - # Model-specific QKV projection (MHA, MLA, etc.) q, k, v = module._project_qkv(hidden_states, position_embeddings) @@ -236,8 +236,6 @@ def __init__(self, group, ulysses_size: int): self.ulysses_size = ulysses_size def project_qkv(self, module, hidden_states, position_embeddings): - from ...models.layers.rope import apply_rotary_pos_emb - # QLoRA fallback: weight is None (packed in quantized buffers). # Fall back to sync-style: module._project_qkv() + synchronous a2a. if module.qkv_proj.weight is None: @@ -368,7 +366,7 @@ def project_qkv(self, module, hidden_states, position_embeddings): return module._project_qkv(hidden_states, position_embeddings) def compute_attention(self, module, q, k, v, attention_mask, **kwargs): - from .ring_attention import ring_flash_attention_forward + from .ring_attention import ring_flash_attention_forward # noqa: PLC0415 attn_kwargs = module._attention_kwargs() cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k = _scale_cu_seqlens_for_ringattn( @@ -425,7 +423,7 @@ def project_qkv(self, module, hidden_states, position_embeddings): return self._ulysses.project_qkv(module, hidden_states, position_embeddings) def compute_attention(self, module, q, k, v, attention_mask, **kwargs): - from .ring_attention import ring_flash_attention_forward + from .ring_attention import ring_flash_attention_forward # noqa: PLC0415 attn_kwargs = module._attention_kwargs() cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k = _scale_cu_seqlens_for_ringattn( @@ -476,7 +474,7 @@ def get_cp_strategy(num_kv_heads: Optional[int] = None) -> CPStrategy: num_kv_heads: Number of key-value heads in the model. Required when Ulysses SP is enabled to choose between sync and async variants. """ - from ...distributed.parallel_state import get_parallel_state + from ...distributed.parallel_state import get_parallel_state # noqa: PLC0415 ps = get_parallel_state() if not ps.cp_enabled: diff --git a/src/xorl/distributed/torch_parallelize.py b/src/xorl/distributed/torch_parallelize.py index 5fc1d4c7..4a08790a 100644 --- a/src/xorl/distributed/torch_parallelize.py +++ b/src/xorl/distributed/torch_parallelize.py @@ -4,8 +4,10 @@ import torch import torch.nn as nn +from torch.distributed._composable.fsdp.fully_shard import FSDPModule from torch.distributed._tensor import Replicate, Shard from torch.distributed.fsdp import MixedPrecision +from torch.distributed.tensor import distribute_tensor from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.checkpoint import noop_context_fn @@ -299,7 +301,6 @@ def _experts_shard_placement_fn(param): # (gradient_accumulate_loss for non-PP; explicit grad.mul_(1/gvt) for PP). # Skip EP-sharded expert modules: they manage their own divide factor (ep_size) # for averaging expert gradients across EP ranks. - from torch.distributed._composable.fsdp.fully_shard import FSDPModule for module in model.modules(): if isinstance(module, FSDPModule) and not getattr(module, "_is_ep_fsdp", False): @@ -324,8 +325,6 @@ def _experts_shard_placement_fn(param): if hasattr(m, "reset_lora_parameters"): m.reset_lora_parameters() else: - from torch.distributed.tensor import distribute_tensor - logger.info_rank0("starting to load model weights...") load_weights_mode = kwargs.get("load_weights_mode", "broadcast") if load_weights_mode == "broadcast": @@ -464,8 +463,6 @@ def build_parallelize_model( # With TP + meta init, we must load weights BEFORE FSDP wrapping. if kwargs.get("init_device") == "meta" and weights_path is not None: - from torch.distributed.tensor import distribute_tensor - if i == 0: logger.info_rank0("TP enabled: loading weights before FSDP wrapping...") load_weights_mode = kwargs.get("load_weights_mode", "broadcast") @@ -613,8 +610,6 @@ def _reentrant_ckpt_with_kwargs(fn, *args, **kw): # between logical DTensor shape and local shard). # Load weights now so FSDP wraps materialized TP DTensors. if kwargs.get("init_device") == "meta" and weights_path is not None: - from torch.distributed.tensor import distribute_tensor - logger.info_rank0("TP enabled: loading weights before FSDP wrapping...") load_weights_mode = kwargs.get("load_weights_mode", "broadcast") if load_weights_mode == "broadcast": diff --git a/src/xorl/lora/__init__.py b/src/xorl/lora/__init__.py index 32d54919..74d1e2ff 100644 --- a/src/xorl/lora/__init__.py +++ b/src/xorl/lora/__init__.py @@ -100,7 +100,7 @@ def __getattr__(name): if name in _MOE_LAZY_ATTRS: - from xorl.models.layers.moe.lora import ( + from xorl.models.layers.moe.lora import ( # noqa: PLC0415 MoEExpertsLoRA, MoELoRAConfig, copy_weights_to_lora_experts, diff --git a/src/xorl/lora/mapping.py b/src/xorl/lora/mapping.py index 301172e8..3c207016 100644 --- a/src/xorl/lora/mapping.py +++ b/src/xorl/lora/mapping.py @@ -8,6 +8,8 @@ import torch.nn as nn +from xorl.models.layers.moe import MoEExperts, MoEExpertsLoRA + from .modules.base import LoraModule from .modules.linear import LoraLinear @@ -28,7 +30,6 @@ def _ensure_moe_mapping(): global _moe_registered if not _moe_registered: _moe_registered = True - from xorl.models.layers.moe import MoEExperts, MoEExpertsLoRA LORA_MAPPING[MoEExperts] = MoEExpertsLoRA diff --git a/src/xorl/lora/utils.py b/src/xorl/lora/utils.py index ecc1ca42..78f1df65 100644 --- a/src/xorl/lora/utils.py +++ b/src/xorl/lora/utils.py @@ -11,8 +11,10 @@ from typing import Dict, Iterator, List, Optional, Tuple import torch +import torch.distributed as dist import torch.nn as nn from safetensors.torch import load_file, save_file +from torch.distributed._tensor import DTensor, Replicate, Shard from xorl.lora.mapping import can_apply_lora, get_lora_class_for_module from xorl.lora.modules import LoraLinear @@ -256,8 +258,6 @@ def _gather_ep_tensor(tensor: torch.Tensor, spec_info) -> torch.Tensor: Full tensor with all EP shards concatenated along the shard dimension. Returns tensor unchanged if placement is Replicate. """ - import torch.distributed as dist - from torch.distributed._tensor import Replicate, Shard if isinstance(spec_info.placement, Replicate): return tensor @@ -292,7 +292,6 @@ def get_lora_state_dict( Returns: State dict containing only lora_A and lora_B parameters (on CPU) """ - from torch.distributed._tensor import DTensor fqn2spec_info = getattr(model, "_fqn2spec_info", None) diff --git a/src/xorl/models/auto.py b/src/xorl/models/auto.py index c6dbdc8f..9c0384fe 100644 --- a/src/xorl/models/auto.py +++ b/src/xorl/models/auto.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Union import torch +import torch.distributed as dist import torch.nn as nn from transformers import ( AutoConfig, @@ -15,6 +16,8 @@ from .layers.attention import ATTENTION_FUNCTIONS from .layers.normalization import set_rmsnorm_mode from .loader import ModelLoader, get_loader +from .transformers.qwen3_5.configuration_qwen3_5 import Qwen3_5Config +from .transformers.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig from .transformers.qwen3_5_shared import ( LINEAR_ATTENTION_RING_UNSUPPORTED_MESSAGE, has_linear_attention_layers, @@ -43,13 +46,9 @@ def _load_local_xorl_config( model_type = config_dict.get("model_type") if model_type == "qwen3_5_moe": - from .transformers.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig - return Qwen3_5MoeConfig.from_hf_config(_namespace_from_dict(config_dict)) if model_type == "qwen3_5": - from .transformers.qwen3_5.configuration_qwen3_5 import Qwen3_5Config - return Qwen3_5Config.from_hf_config(_namespace_from_dict(config_dict)) return None @@ -81,7 +80,6 @@ def _load_config_with_rank0_priority( 'Unrecognized model' errors. This function lets rank 0 download first (populating the cache), then other ranks load from the cache. """ - import torch.distributed as dist rank = get_parallel_state().global_rank if get_parallel_state().is_initialized else 0 is_distributed = dist.is_initialized() and dist.get_world_size() > 1 diff --git a/src/xorl/models/checkpoint_handlers/buffers.py b/src/xorl/models/checkpoint_handlers/buffers.py index e74e86d4..ad27e66c 100644 --- a/src/xorl/models/checkpoint_handlers/buffers.py +++ b/src/xorl/models/checkpoint_handlers/buffers.py @@ -6,12 +6,28 @@ """ import json +import math import os import re from collections import defaultdict from typing import Dict, List, Optional, Set, Tuple import torch +from torch.distributed._tensor import DTensor +from transformers.utils import cached_file + +from xorl.ops.quantize import ( + block_fp8_dequantize_gkn, + block_fp8_quantize_gkn, + nvfp4_dequantize, + nvfp4_quantize, +) +from xorl.ops.quantize.fp4_codec import FP4_E2M1_MAX, FP8_E4M3_MAX +from xorl.qlora.modules.linear import QLoRALinear +from xorl.qlora.modules.moe_experts import ( + BlockFP8QLoRAMoeExperts, + QLoRAMoeExperts, +) # ============================================================================= @@ -120,8 +136,6 @@ def _resolve_weights_path(weights_path: Optional[str]) -> Optional[str]: return weights_path # Try resolving as a HF hub model ID try: - from transformers.utils import cached_file - config_path = cached_file(weights_path, "config.json", _raise_exceptions_for_missing_entries=False) if config_path and os.path.isfile(config_path): return os.path.dirname(config_path) @@ -604,8 +618,6 @@ class _QLoRAWeightBufferBase: _suffixes: frozenset def __init__(self, model: "torch.nn.Module"): - from xorl.qlora.modules.linear import QLoRALinear - self._prefix_map: Dict[str, _SourceInfo] = {} self._accums: Dict[str, _ModuleAccum] = {} @@ -734,13 +746,6 @@ def _convert_format( dispatch_pairs: List[Tuple[str, "torch.Tensor"]], ) -> List[Tuple[str, "torch.Tensor"]]: """Dequantize from source format, re-quantize in target format.""" - from xorl.ops.quantize import ( - block_fp8_dequantize_gkn, - block_fp8_quantize_gkn, - nvfp4_dequantize, - nvfp4_quantize, - ) - from xorl.qlora.modules.linear import QLoRALinear module = accum.module M, K = module.out_features, module.in_features @@ -810,8 +815,6 @@ class NvFP4WeightBuffer(_QLoRAWeightBufferBase): _source_format = "nvfp4" def _pack(self, fqn: str, accum: _ModuleAccum) -> List[Tuple[str, "torch.Tensor"]]: - from xorl.qlora.modules.linear import QLoRALinear - module = accum.module merge_sources = accum.merge_sources @@ -863,8 +866,6 @@ class BlockFP8WeightBuffer(_QLoRAWeightBufferBase): _source_format = "block_fp8" def _pack(self, fqn: str, accum: _ModuleAccum) -> List[Tuple[str, "torch.Tensor"]]: - from xorl.qlora.modules.linear import QLoRALinear - module = accum.module merge_sources = accum.merge_sources @@ -896,7 +897,6 @@ def QLoRAWeightBuffer(model: "torch.nn.Module") -> _QLoRAWeightBufferBase: Inspects the first prequantized QLoRALinear module to determine whether the checkpoint is NVFP4 or block FP8, then returns the matching buffer class. """ - from xorl.qlora.modules.linear import QLoRALinear for _fqn, module in model.named_modules(): if isinstance(module, QLoRALinear) and module._is_prequantized: @@ -933,8 +933,6 @@ def __init__( ep_size: int, num_experts: int, ): - from xorl.qlora.modules.moe_experts import QLoRAMoeExperts - self._num_experts = num_experts self._local_num = num_experts // ep_size self._expert_start = ep_rank * self._local_num @@ -1064,9 +1062,6 @@ def get_pending(self) -> List[Tuple[int, str]]: def set_inline_metadata(self) -> None: """Finalize inline-loaded modules: materialize LoRA params from meta device.""" - import math - - from torch.distributed._tensor import DTensor for module in self._layer_modules.values(): if not module._weights_loaded: @@ -1114,9 +1109,6 @@ class NvFP4QLoRAExpertBuffer(_QLoRAExpertBufferBase): _suffixes = _QLORA_SUFFIXES_NVFP4 # {"weight", "weight_scale", "weight_scale_2"} def _process_expert(self, layer, proj, local_idx, pieces): - from xorl.ops.quantize.fp4_codec import FP4_E2M1_MAX, FP8_E4M3_MAX - from xorl.qlora.modules.moe_experts import QLoRAMoeExperts - lp = (layer, proj) packed = pieces["weight"] # [N, K//2] uint8 block_scales = pieces["weight_scale"] # [N, K//bs] fp8 @@ -1137,8 +1129,6 @@ def _process_expert(self, layer, proj, local_idx, pieces): self._meta_lists[lp][local_idx] = amax def _finalize_proj(self, layer, proj): - from xorl.qlora.modules.moe_experts import QLoRAMoeExperts - lp = (layer, proj) module = self._layer_modules[layer] device = torch.device("cuda") @@ -1167,8 +1157,6 @@ class BlockFP8QLoRAExpertBuffer(_QLoRAExpertBufferBase): _suffixes = _QLORA_SUFFIXES_FP8 # {"weight", "weight_scale_inv"} def _process_expert(self, layer, proj, local_idx, pieces): - from xorl.qlora.modules.moe_experts import QLoRAMoeExperts - lp = (layer, proj) fp8_w = pieces["weight"] # [N, K] float8_e4m3fn scales = pieces["weight_scale_inv"] # [N//128, K//128] f32 @@ -1205,10 +1193,6 @@ def QLoRAExpertBuffer( num_experts: int, ) -> Optional[_QLoRAExpertBufferBase]: """Factory: create format-specific QLoRA expert buffer, or None if no MoE experts.""" - from xorl.qlora.modules.moe_experts import ( - BlockFP8QLoRAMoeExperts, - QLoRAMoeExperts, - ) for _, m in model.named_modules(): if isinstance(m, QLoRAMoeExperts) and not m._weights_loaded: diff --git a/src/xorl/models/layers/activations.py b/src/xorl/models/layers/activations.py index 42a95ed5..58a225ed 100644 --- a/src/xorl/models/layers/activations.py +++ b/src/xorl/models/layers/activations.py @@ -12,6 +12,7 @@ import torch from torch import Tensor, nn +from torch._dynamo import allow_in_graph logger = logging.getLogger(__name__) @@ -234,13 +235,11 @@ def __init__( self._xielu_cuda_obj = None try: - import xielu.ops # noqa: F401 + import xielu.ops # noqa: F401, PLC0415 self._xielu_cuda_obj = torch.classes.xielu.XIELU() msg = "Using experimental xIELU CUDA." try: - from torch._dynamo import allow_in_graph - self._xielu_cuda_fn = allow_in_graph(self._xielu_cuda) msg += " Enabled torch._dynamo for xIELU CUDA." except Exception as err: diff --git a/src/xorl/models/layers/moe/experts.py b/src/xorl/models/layers/moe/experts.py index 99c2654e..bf5eca52 100644 --- a/src/xorl/models/layers/moe/experts.py +++ b/src/xorl/models/layers/moe/experts.py @@ -3,10 +3,18 @@ import os import torch +import torch.distributed as dist import torch.nn as nn from ..activations import ACT2FN -from .backend import MOE_EXPERT_BACKENDS, MOE_EXPERT_BACKENDS_MOE_ACT +from .backend import ( + EP_COMBINE, + EP_DISPATCH, + EP_EXPERT_COMPUTE, + EP_EXPERT_COMPUTE_MOE_ACT, + MOE_EXPERT_BACKENDS, + MOE_EXPERT_BACKENDS_MOE_ACT, +) from .common import split_gate_up_proj @@ -119,7 +127,7 @@ def forward( ) # Check EP — use unified dispatch/compute/combine path - from xorl.distributed.parallel_state import get_parallel_state + from xorl.distributed.parallel_state import get_parallel_state # noqa: PLC0415 parallel_state = get_parallel_state() @@ -160,7 +168,6 @@ def _ep_forward( Dispatch strategy is selected by ``self.ep_dispatch`` (``"alltoall"`` or ``"deepep"``). Compute backend by ``self.moe_implementation``. """ - from .backend import EP_COMBINE, EP_DISPATCH, EP_EXPERT_COMPUTE, EP_EXPERT_COMPUTE_MOE_ACT if self.moe_implementation not in EP_EXPERT_COMPUTE: raise ValueError( @@ -256,7 +263,6 @@ def _ep_forward_debug(self, dispatch_fn, combine_fn, compute_fn, dispatch_kwargs times plus tensor metadata to help diagnose performance gaps between different dispatch+compute backend combinations. """ - import torch.distributed as dist rank = dist.get_rank() if dist.is_initialized() else 0 @@ -321,7 +327,7 @@ def _build_dispatch_kwargs(self, hidden_states, routing_weights, selected_expert if self.ep_dispatch == "alltoall": kwargs["ep_group"] = parallel_state.ep_group elif self.ep_dispatch == "deepep": - from xorl.distributed.moe.deepep import get_default_buffer + from xorl.distributed.moe.deepep import get_default_buffer # noqa: PLC0415 kwargs["buffer"] = get_default_buffer( ep_group=parallel_state.ep_group, diff --git a/src/xorl/models/layers/moe/lora.py b/src/xorl/models/layers/moe/lora.py index 83fdbca2..fad35392 100644 --- a/src/xorl/models/layers/moe/lora.py +++ b/src/xorl/models/layers/moe/lora.py @@ -26,6 +26,7 @@ from ....ops.group_gemm.kernel import compute_lora_scaling from ....utils import logging from ..activations import ACT2FN +from .backend import EP_COMBINE, EP_DISPATCH, EP_EXPERT_COMPUTE_LORA, MOE_EXPERT_BACKENDS_LORA from .common import split_gate_up_proj @@ -206,7 +207,7 @@ def forward( return self._eager_lora_forward(hidden_states, expert_idx) # Check EP — use unified dispatch/compute/combine path - from xorl.distributed.parallel_state import get_parallel_state + from xorl.distributed.parallel_state import get_parallel_state # noqa: PLC0415 parallel_state = get_parallel_state() @@ -214,7 +215,6 @@ def forward( return self._ep_forward(hidden_states, routing_weights, selected_experts, parallel_state) # Local path — registry-based - from .backend import MOE_EXPERT_BACKENDS_LORA fn = MOE_EXPERT_BACKENDS_LORA[self.moe_implementation] gate_proj = self.gate_proj.contiguous() @@ -248,7 +248,6 @@ def _ep_forward( Uses the same dispatch/combine as ``MoEExperts._ep_forward()`` but routes to the LoRA-aware EP compute registry. """ - from .backend import EP_COMBINE, EP_DISPATCH, EP_EXPERT_COMPUTE_LORA if self.moe_implementation not in EP_EXPERT_COMPUTE_LORA: raise ValueError( @@ -301,7 +300,7 @@ def _build_dispatch_kwargs(self, hidden_states, routing_weights, selected_expert if self.ep_dispatch == "alltoall": kwargs["ep_group"] = parallel_state.ep_group elif self.ep_dispatch == "deepep": - from xorl.distributed.moe.deepep import get_default_buffer + from xorl.distributed.moe.deepep import get_default_buffer # noqa: PLC0415 kwargs["buffer"] = get_default_buffer( ep_group=parallel_state.ep_group, diff --git a/src/xorl/models/layers/moe/moe_block.py b/src/xorl/models/layers/moe/moe_block.py index 85022102..7255a45b 100644 --- a/src/xorl/models/layers/moe/moe_block.py +++ b/src/xorl/models/layers/moe/moe_block.py @@ -7,6 +7,7 @@ import torch.nn.functional as F from .experts import MoEExperts +from .lora import inject_lora_into_experts from .router import TopKRouter from .routing_replay import RoutingReplay, get_replay_stage @@ -105,7 +106,6 @@ def inject_lora( hybrid_shared: If True, share ``lora_A`` for gate/up and ``lora_B`` for down across experts. """ - from .lora import inject_lora_into_experts inject_lora_into_experts( self, diff --git a/src/xorl/models/module_utils.py b/src/xorl/models/module_utils.py index 8080d65a..00753aac 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -1,8 +1,10 @@ import json +import math import os import re import time -from collections import OrderedDict +from collections import OrderedDict, deque +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass from datetime import timedelta @@ -14,6 +16,7 @@ from safetensors.torch import load_file, save_file from torch import distributed as dist from torch import nn +from torch.distributed._tensor import Shard as DTShard from tqdm import tqdm from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME from transformers.utils.hub import cached_file, get_checkpoint_shard_files @@ -253,8 +256,6 @@ def _prefetch_shards( more CPU memory but can increase NFS throughput via concurrent reads. Default 1 (single background thread, same as before). """ - from collections import deque - from concurrent.futures import ThreadPoolExecutor if not state_dict_iterators: return @@ -290,8 +291,6 @@ def _prefetch_shards_filtered( Yields (state_dict, skipped_keys) tuples. """ - from collections import deque - from concurrent.futures import ThreadPoolExecutor if not state_dict_iterators: return @@ -328,7 +327,6 @@ def _load_state_dict(weights_path: str, **kwargs) -> List["StateDictIterator"]: """ Loads (sharded) state dict in transformers' format. """ - import time max_retries = 5 for attempt in range(max_retries): @@ -368,12 +366,10 @@ def _try_load_state_dict(weights_path: str, **kwargs): return [StateDictIterator(shard_file) for shard_file in shard_files] except OSError as e: if attempt < max_retries - 1: - import time - retry_delay = 5 logger.warning( f"OSError resolving shard files (attempt {attempt + 1}/{max_retries}): {e}. Retrying in 5s..." ) - time.sleep(retry_delay) + time.sleep(5) else: logger.error(f"Failed to get shard files after {max_retries} attempts") raise @@ -653,7 +649,6 @@ def _dispatch_parameter( if orig_tensor.device.type == "cpu": # CPU DTensor: copy shard directly into local tensor. # distribute_tensor doesn't support CPU mesh, so we manually shard and copy. - from torch.distributed._tensor import Shard as DTShard shard_dim = None for p in placements: @@ -714,9 +709,6 @@ def _init_parameter( - lora_A: kaiming uniform initialization - lora_B: zeros (so LoRA has no effect at start) """ - import math - - import torch.nn as nn # Check if this is a LoRA parameter and handle specially if "lora_A" in name or "lora_B" in name: @@ -895,8 +887,6 @@ def all_ranks_load_weights( f"OSError loading weights (attempt {attempt + 1}/{max_retries}): {e}. " f"Retrying in {retry_delay} seconds..." ) - import time - time.sleep(retry_delay) else: logger.error(f"Failed to load weights after {max_retries} attempts") diff --git a/src/xorl/models/transformers/qwen3/checkpoint_handler.py b/src/xorl/models/transformers/qwen3/checkpoint_handler.py index 53dbb675..d52dc922 100644 --- a/src/xorl/models/transformers/qwen3/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3/checkpoint_handler.py @@ -1,5 +1,6 @@ """Checkpoint handler for dense Qwen3 models.""" +import warnings from typing import Callable, List, Optional, Set, Tuple import torch @@ -158,8 +159,6 @@ def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torc return [(key, tensor)] def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: - import warnings - pending_gu = self._gate_up_buffer.get_pending() if pending_gu: warnings.warn(f"Incomplete gate/up merge pairs after loading: {pending_gu}") diff --git a/src/xorl/models/transformers/qwen3_5/checkpoint_handler.py b/src/xorl/models/transformers/qwen3_5/checkpoint_handler.py index 163201b2..d487c822 100644 --- a/src/xorl/models/transformers/qwen3_5/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3_5/checkpoint_handler.py @@ -1,5 +1,6 @@ """Checkpoint handler for dense Qwen3_5 models.""" +import warnings from typing import Callable, List, Optional, Set, Tuple import torch @@ -154,8 +155,6 @@ def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torc return [(key, tensor)] def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: - import warnings - pending_gu = self._gate_up_buffer.get_pending() if pending_gu: warnings.warn(f"Incomplete gate/up merge pairs after loading: {pending_gu}") diff --git a/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py b/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py index f70b85f8..0683f36b 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py @@ -1,6 +1,7 @@ """Checkpoint handler for Qwen3_5 MoE models.""" import re +import warnings from typing import Callable, Dict, List, Optional, Set, Tuple import torch @@ -375,8 +376,6 @@ def on_skip_weight(self, key: str) -> List[Tuple[str, torch.Tensor]]: return [] def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: - import warnings - if self._expert_buffer is not None: pending = self._expert_buffer.get_pending_counts() if pending: diff --git a/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py b/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py index 63ae7320..0acf7a17 100644 --- a/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py @@ -1,6 +1,7 @@ """Checkpoint handler for Qwen3 MoE models.""" import re +import warnings from typing import Callable, Dict, List, Optional, Set, Tuple import torch @@ -398,8 +399,6 @@ def on_skip_weight(self, key: str) -> List[Tuple[str, torch.Tensor]]: return [] def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: - import warnings - if self._expert_buffer is not None: pending = self._expert_buffer.get_pending_counts() if pending: diff --git a/src/xorl/ops/group_gemm/utils/device.py b/src/xorl/ops/group_gemm/utils/device.py index f304bcd3..21788352 100644 --- a/src/xorl/ops/group_gemm/utils/device.py +++ b/src/xorl/ops/group_gemm/utils/device.py @@ -1,12 +1,12 @@ from functools import lru_cache +import torch + from ....utils.device import get_device_name @lru_cache def get_device_key() -> str: - import torch - if torch.cuda.get_device_capability() == (8, 0): return "A100" # A30 is treated the same way as A100 for the moment. diff --git a/src/xorl/ops/moe/lora.py b/src/xorl/ops/moe/lora.py index 0af7b710..99a50ea3 100644 --- a/src/xorl/ops/moe/lora.py +++ b/src/xorl/ops/moe/lora.py @@ -446,7 +446,7 @@ def backward(ctx, grad_output): grad_permute_tokens = grad_permute_tokens_1 + grad_permute_tokens_2 # === Reduce gradients for shared weights === - from xorl.distributed.parallel_state import get_parallel_state + from xorl.distributed.parallel_state import get_parallel_state # noqa: PLC0415 ep_group = get_parallel_state().ep_group diff --git a/src/xorl/ops/quack/cute_dsl_ptxas.py b/src/xorl/ops/quack/cute_dsl_ptxas.py index 6091aa49..639a7b79 100644 --- a/src/xorl/ops/quack/cute_dsl_ptxas.py +++ b/src/xorl/ops/quack/cute_dsl_ptxas.py @@ -95,7 +95,7 @@ def _patched_load_cuda_library(self): return _original_load_cuda_library(self) # Load cubin - import cuda.bindings.runtime as cuda_runtime + import cuda.bindings.runtime as cuda_runtime # noqa: PLC0415 err, library = cuda_runtime.cudaLibraryLoadData(cubin, None, None, 0, None, None, 0) if err != cuda_runtime.cudaError_t.cudaSuccess: diff --git a/src/xorl/ops/quack/gemm_sm100.py b/src/xorl/ops/quack/gemm_sm100.py index dac96acc..3a73cbcf 100644 --- a/src/xorl/ops/quack/gemm_sm100.py +++ b/src/xorl/ops/quack/gemm_sm100.py @@ -2,6 +2,7 @@ # https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py import argparse +import time from functools import partial from typing import Callable, Literal, Optional, Tuple, Type, Union @@ -24,6 +25,7 @@ from cutlass.cute.runtime import from_dlpack, make_ptr from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.utils import LayoutEnum +from triton.testing import do_bench from . import copy_utils from . import sm100_utils as quack_sm100_utils @@ -2542,8 +2544,6 @@ def create_and_permute_tensor(l, mode0, mode1, is_mode0_major, dtype, is_dynamic # Reference checking ref_d and gpu_d torch.testing.assert_close(gpu_d, ref_d, atol=tolerance, rtol=1e-05) - from triton.testing import do_bench - current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) flops = 2 * m * n * k * l @@ -2551,8 +2551,6 @@ def create_and_permute_tensor(l, mode0, mode1, is_mode0_major, dtype, is_dynamic repeats = iterations warmup = warmup_iterations - import time - time.sleep(0.5) if ab_dtype.width == 8: assert l == 1 diff --git a/src/xorl/ops/quack/utils.py b/src/xorl/ops/quack/utils.py index 1dab5c94..f3c9dd70 100644 --- a/src/xorl/ops/quack/utils.py +++ b/src/xorl/ops/quack/utils.py @@ -208,7 +208,7 @@ def warp_prefix_sum(val: Int32, lane: Optional[Int32] = None) -> Int32: @dsl_user_op def atomic_inc_i32(a: int | Int32, gmem_ptr: cute.Pointer, *, loc=None, ip=None) -> Int32: - from cutlass import CUDA_VERSION + from cutlass import CUDA_VERSION # noqa: PLC0415 # * NVVM call based on nvvm version if CUDA_VERSION.major == 12 and CUDA_VERSION.minor == 9: @@ -221,7 +221,7 @@ def atomic_inc_i32(a: int | Int32, gmem_ptr: cute.Pointer, *, loc=None, ip=None) @dsl_user_op def atomic_add_i32(a: int | Int32, gmem_ptr: cute.Pointer, *, loc=None, ip=None) -> Int32: - from cutlass import CUDA_VERSION + from cutlass import CUDA_VERSION # noqa: PLC0415 # * NVVM call based on nvvm version if CUDA_VERSION.major == 12 and CUDA_VERSION.minor == 9: diff --git a/src/xorl/qlora/modules/linear.py b/src/xorl/qlora/modules/linear.py index 489a246d..00daed2f 100644 --- a/src/xorl/qlora/modules/linear.py +++ b/src/xorl/qlora/modules/linear.py @@ -13,12 +13,17 @@ import math from typing import Dict, Optional, Tuple +import safetensors.torch import torch import torch.nn as nn import torch.nn.functional as F from torch import Tensor +from torch.distributed._tensor import DTensor +from torch.distributed._tensor import Shard as _Shard +from transformers.utils import cached_file from xorl.lora.modules.base import LoraModule +from xorl.qlora.modules.moe_experts import QLoRAMoeExperts class QLoRALinear(LoraModule, nn.Module): @@ -111,9 +116,9 @@ def from_module(cls, module: nn.Module, r: int = 16, lora_alpha: int = 16, **kwa based on ``quant_format`` kwarg (default: "nvfp4"). """ if cls is QLoRALinear: - from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear - from xorl.qlora.modules.nf4_linear import NF4QLoRALinear - from xorl.qlora.modules.nvfp4_linear import NvFP4QLoRALinear + from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear # noqa: PLC0415 + from xorl.qlora.modules.nf4_linear import NF4QLoRALinear # noqa: PLC0415 + from xorl.qlora.modules.nvfp4_linear import NvFP4QLoRALinear # noqa: PLC0415 fmt = kwargs.pop("quant_format", "nvfp4") kwargs.pop("quant_group_size", None) # subclass sets this @@ -145,9 +150,9 @@ def from_quantized( ) -> "QLoRALinear": """Create QLoRALinear from pre-quantized weights (skip quantization).""" if cls is QLoRALinear: - from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear - from xorl.qlora.modules.nf4_linear import NF4QLoRALinear - from xorl.qlora.modules.nvfp4_linear import NvFP4QLoRALinear + from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear # noqa: PLC0415 + from xorl.qlora.modules.nf4_linear import NF4QLoRALinear # noqa: PLC0415 + from xorl.qlora.modules.nvfp4_linear import NvFP4QLoRALinear # noqa: PLC0415 if quant_format == "block_fp8": cls = BlockFP8QLoRALinear @@ -219,7 +224,6 @@ def _write_packed_weight(self, uint8_data: Tensor) -> None: local = self.packed_weight_f32._local_tensor full_shape = self.packed_weight_f32.shape f32_data = f32_data.reshape(full_shape) - from torch.distributed._tensor import Shard as _Shard placements = self.packed_weight_f32.placements shard_dim = 0 @@ -267,8 +271,6 @@ def _read_packed_weight_uint8(self) -> Tensor: def load_prequantized_weights(self, weight_map: dict, shard_cache: dict, weights_path: str) -> None: """Load pre-quantized weights from checkpoint.""" - import safetensors.torch - from transformers.utils import cached_file def _load_tensor(ckpt_key: str) -> Tensor: shard_file = weight_map.get(ckpt_key) @@ -294,7 +296,6 @@ def quantize_weight(self) -> None: if self.weight is None: return w = self.weight.data - from torch.distributed._tensor import DTensor if isinstance(w, DTensor): w = w.full_tensor() @@ -423,7 +424,6 @@ def prefetch_aqn_noise( Returns: Number of modules prefetched. """ - from xorl.qlora.modules.moe_experts import QLoRAMoeExperts count = 0 for module in model.modules(): diff --git a/src/xorl/qlora/modules/moe_experts.py b/src/xorl/qlora/modules/moe_experts.py index e9cd6b28..2edf0e14 100644 --- a/src/xorl/qlora/modules/moe_experts.py +++ b/src/xorl/qlora/modules/moe_experts.py @@ -25,9 +25,11 @@ import math from typing import Optional +import safetensors.torch import torch import torch.nn as nn from torch import Tensor +from transformers.utils import cached_file from xorl.lora.modules.base import LoraModule from xorl.ops.group_gemm.kernel.lora_utils import compute_lora_scaling @@ -390,9 +392,6 @@ def load_and_quantize_weights( if self._weights_loaded or self._source_fqn is None: return - import safetensors.torch - from transformers.utils import cached_file - _shard_cache = shard_cache if shard_cache is not None else {} def _load_tensor(ckpt_key: str) -> Tensor: @@ -456,7 +455,7 @@ def forward( return self._eager_lora_forward(hidden_states, expert_idx) # Check EP -- use unified dispatch/compute/combine path - from xorl.distributed.parallel_state import get_parallel_state + from xorl.distributed.parallel_state import get_parallel_state # noqa: PLC0415 parallel_state = get_parallel_state() @@ -464,7 +463,7 @@ def forward( return self._ep_forward(hidden_states, routing_weights, selected_experts, parallel_state) # Local path -- registry-based - from xorl.models.layers.moe.backend import MOE_EXPERT_BACKENDS_LORA + from xorl.models.layers.moe.backend import MOE_EXPERT_BACKENDS_LORA # noqa: PLC0415 compute_dtype = hidden_states.dtype fn = MOE_EXPERT_BACKENDS_LORA[self.moe_implementation] @@ -498,7 +497,7 @@ def _ep_forward( Uses the same dispatch/combine as MoEExperts._ep_forward() but routes to the LoRA-aware EP compute registry. """ - from xorl.models.layers.moe.backend import EP_COMBINE, EP_DISPATCH, EP_EXPERT_COMPUTE_LORA + from xorl.models.layers.moe.backend import EP_COMBINE, EP_DISPATCH, EP_EXPERT_COMPUTE_LORA # noqa: PLC0415 if self.moe_implementation not in EP_EXPERT_COMPUTE_LORA: raise ValueError( @@ -550,7 +549,7 @@ def _build_dispatch_kwargs(self, hidden_states, routing_weights, selected_expert if self.ep_dispatch == "alltoall": kwargs["ep_group"] = parallel_state.ep_group elif self.ep_dispatch == "deepep": - from xorl.distributed.moe.deepep import get_default_buffer + from xorl.distributed.moe.deepep import get_default_buffer # noqa: PLC0415 kwargs["buffer"] = get_default_buffer( ep_group=parallel_state.ep_group, diff --git a/src/xorl/qlora/utils.py b/src/xorl/qlora/utils.py index aa8778a7..2e141317 100644 --- a/src/xorl/qlora/utils.py +++ b/src/xorl/qlora/utils.py @@ -6,13 +6,25 @@ import json import logging +import math +import math as _math import os +from concurrent.futures import ThreadPoolExecutor from typing import Collection, List, Optional, Tuple +import safetensors.torch import torch +import torch.distributed as dist import torch.nn as nn from safetensors.torch import save_file - +from torch.distributed._tensor import DTensor +from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, cached_file + +from xorl.distributed.parallel_state import get_parallel_state +from xorl.models.checkpoint_handlers.buffers import ( + detect_prequantized_block_fp8_checkpoint, + detect_prequantized_checkpoint, +) from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear from xorl.qlora.modules.linear import QLoRALinear from xorl.qlora.modules.moe_experts import QLoRAMoeExperts @@ -210,8 +222,6 @@ def inject_qlora_into_model( if not any(t in ("gate_proj", "up_proj", "down_proj") for t in target_modules): return model - from xorl.distributed.parallel_state import get_parallel_state - try: parallel_state = get_parallel_state() ep_size = parallel_state.ep_size if parallel_state.ep_enabled else 1 @@ -386,7 +396,6 @@ def maybe_load_and_quantize_moe_qlora( Returns: Number of MoE modules loaded. """ - import json needs_moe = any(isinstance(m, QLoRAMoeExperts) and not m._weights_loaded for m in model.modules()) if not needs_moe: @@ -396,8 +405,6 @@ def maybe_load_and_quantize_moe_qlora( weight_map = None if weights_path: try: - from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, cached_file - index_path = cached_file(weights_path, SAFE_WEIGHTS_INDEX_NAME) if index_path: with open(index_path) as f: @@ -423,7 +430,6 @@ def maybe_load_and_quantize_moe_qlora( shard_cache=shard_cache, ) # Materialize LoRA params from meta device to GPU - from torch.distributed._tensor import DTensor for name, param in module.named_parameters(): if param.device.type == "meta": @@ -446,8 +452,6 @@ def maybe_load_and_quantize_moe_qlora( requires_grad=param.requires_grad, ) else: - import math as _math - materialized = nn.Parameter( torch.zeros( param.shape, @@ -539,7 +543,6 @@ def detect_prequantized_nvfp4(weights_path: str) -> bool: Returns: True if the checkpoint is pre-quantized NVFP4. """ - from xorl.models.checkpoint_handlers.buffers import detect_prequantized_checkpoint return detect_prequantized_checkpoint(weights_path) @@ -555,7 +558,6 @@ def detect_prequantized_block_fp8(weights_path: str) -> bool: Returns: True if the checkpoint is pre-quantized block FP8. """ - from xorl.models.checkpoint_handlers.buffers import detect_prequantized_block_fp8_checkpoint return detect_prequantized_block_fp8_checkpoint(weights_path) @@ -573,11 +575,6 @@ def _broadcast_shard_cache( Returns: shard_cache: dict mapping shard filename → {tensor_name: tensor (CPU)} """ - from concurrent.futures import ThreadPoolExecutor - - import safetensors.torch - import torch.distributed as dist - from transformers.utils import cached_file rank = dist.get_rank() device = torch.device("cuda") @@ -661,13 +658,10 @@ def maybe_load_prequantized_qlora(model: nn.Module, weights_path: str, load_mode Returns: Number of modules loaded. """ - import torch.distributed as dist # Load weight map from index file weight_map = None try: - from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, cached_file - index_path = cached_file(weights_path, SAFE_WEIGHTS_INDEX_NAME) if index_path: with open(index_path) as f: @@ -719,7 +713,6 @@ def maybe_load_prequantized_qlora(model: nn.Module, weights_path: str, load_mode # Since lora_B=0 → initial delta = A @ 0 = 0 regardless of A, # we initialize ALL DTensor params to zeros and let kaiming_uniform # only apply to non-DTensor (non-EP) params. - from torch.distributed._tensor import DTensor _rank = dist.get_rank() if (dist.is_available() and dist.is_initialized()) else 0 for name, param in module.named_parameters(): @@ -754,8 +747,6 @@ def maybe_load_prequantized_qlora(model: nn.Module, weights_path: str, load_mode ) parts = name.split("_") if parts[-1] == "A": - import math - for i in range(materialized.shape[0]): nn.init.kaiming_uniform_(materialized.data[i], a=math.sqrt(5)) setattr(module, name, materialized) diff --git a/src/xorl/server/api_server/inference_endpoints.py b/src/xorl/server/api_server/inference_endpoints.py index 1afbea0b..1acca05b 100644 --- a/src/xorl/server/api_server/inference_endpoints.py +++ b/src/xorl/server/api_server/inference_endpoints.py @@ -3,13 +3,16 @@ from __future__ import annotations import asyncio +import json import logging import os import socket +from pathlib import Path from typing import Any, Dict, List import httpx from fastapi import HTTPException, status +from huggingface_hub import hf_hub_download from xorl.server.api_server.api_types import ( AddInferenceEndpointRequest, @@ -47,8 +50,6 @@ def _detect_quantization_from_hf_config(model_path: str) -> dict | None: Returns the full HF quantization_config dict, or None. """ - import json - from pathlib import Path # Try local path first config_path = Path(model_path) / "config.json" @@ -63,8 +64,6 @@ def _detect_quantization_from_hf_config(model_path: str) -> dict | None: # Try HuggingFace hub (for repo IDs like "Qwen/Qwen3-8B-FP8") if config_dict is None: try: - from huggingface_hub import hf_hub_download - cached_path = hf_hub_download(model_path, "config.json") with open(cached_path) as f: config_dict = json.load(f) diff --git a/src/xorl/server/launcher.py b/src/xorl/server/launcher.py index 8bbfc3a2..e56d85a1 100644 --- a/src/xorl/server/launcher.py +++ b/src/xorl/server/launcher.py @@ -25,12 +25,14 @@ import logging import multiprocessing as mp import os +import random import signal import socket import subprocess import sys import time -from contextlib import closing +from contextlib import asynccontextmanager, closing +from dataclasses import fields from pathlib import Path from typing import Dict, List, Optional, Tuple @@ -41,6 +43,7 @@ from xorl.server.api_server.server import APIServer from xorl.server.orchestrator.orchestrator import Orchestrator from xorl.server.server_arguments import ServerArguments +from xorl.server.utils.network import read_address_file # Setup logging @@ -97,7 +100,6 @@ def find_free_port(start_port: int = 50000, max_attempts: int = 10000) -> int: Raises: RuntimeError: If no free port found """ - import random end_port = min(start_port + max_attempts, 60000) ports = list(range(start_port, end_port)) @@ -265,7 +267,6 @@ def run_api_server( skip_initial_checkpoint: Skip auto-saving initial checkpoint on first create_model. sync_inference_method: Method for syncing weights to inference endpoints. Default: 'nccl_broadcast'. """ - from contextlib import asynccontextmanager # Setup logging for this process logging.basicConfig( @@ -287,8 +288,8 @@ def run_api_server( try: # Import the FastAPI app from api_server module # Update the global state shared between api_server.py and endpoints.py - import xorl.server.api_server._state as _state_module - from xorl.server.api_server.server import app + import xorl.server.api_server._state as _state_module # noqa: PLC0415 + from xorl.server.api_server.server import app # noqa: PLC0415 # Override the lifespan to use our addresses @asynccontextmanager @@ -365,8 +366,6 @@ def load_server_arguments(config_path: str, overrides: Optional[Dict[str, any]] if not config: raise ValueError(f"Empty config file: {config_path}") - from dataclasses import fields - valid_fields = {f.name for f in fields(ServerArguments)} # Check if this is a nested config (has model/train/worker sections) @@ -720,8 +719,6 @@ def _get_rank0_worker_address(self) -> str: # Priority 2: File-based discovery (for multi-node) if self.nnodes > 1: - from xorl.server.utils.network import read_address_file - logger.info(f"Multi-node setup (nnodes={self.nnodes}), waiting for rank 0 address file...") # Wait for address file with extended timeout for multi-node @@ -858,7 +855,6 @@ def _save_initial_checkpoint(self, max_retries: int = 3, retry_delay: float = 5. @staticmethod def _find_free_port() -> int: """Find and return a free TCP port.""" - import socket with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind(("", 0)) diff --git a/src/xorl/server/orchestrator/packing.py b/src/xorl/server/orchestrator/packing.py index 148f8733..1fa1791c 100644 --- a/src/xorl/server/orchestrator/packing.py +++ b/src/xorl/server/orchestrator/packing.py @@ -179,7 +179,6 @@ def unpack_per_token_outputs( >>> result = unpack_per_token_outputs(packed_output, position_ids) >>> # result = [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6]] """ - from xorl.utils.seqlen_pos_transform_utils import pos2culen # Convert to tensors if needed if isinstance(packed_output, list): diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index ab2c916d..4862c913 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -7,9 +7,12 @@ parameters like batch size, epochs, and optimizer settings. """ +import sys from dataclasses import dataclass, field from typing import Any, Dict, List, Literal, Optional +import yaml + @dataclass class ServerArguments: @@ -654,9 +657,6 @@ def parse_server_args() -> ServerArguments: Returns: ServerArguments with all fields populated from YAML and CLI """ - import sys - - import yaml # Read YAML directly to get flat structure config_path = None diff --git a/src/xorl/server/utils/network.py b/src/xorl/server/utils/network.py index 927562d2..bfb6b266 100644 --- a/src/xorl/server/utils/network.py +++ b/src/xorl/server/utils/network.py @@ -8,6 +8,8 @@ import logging import os import socket +import time +from pathlib import Path from typing import Optional @@ -149,7 +151,6 @@ def write_address_file(address: str, output_dir: str, filename: str = ".rank0_ad >>> write_address_file("tcp://10.0.0.5:5556", "/shared/outputs") '/shared/outputs/.rank0_address' """ - from pathlib import Path # Ensure output directory exists Path(output_dir).mkdir(parents=True, exist_ok=True) @@ -187,8 +188,6 @@ def read_address_file( >>> read_address_file("/shared/outputs", timeout=60.0) 'tcp://10.0.0.5:5556' """ - import time - from pathlib import Path address_file = Path(output_dir) / filename start_time = time.time() diff --git a/src/xorl/server/weight_sync/backends/__init__.py b/src/xorl/server/weight_sync/backends/__init__.py index a94addf1..dcb8df5c 100644 --- a/src/xorl/server/weight_sync/backends/__init__.py +++ b/src/xorl/server/weight_sync/backends/__init__.py @@ -26,7 +26,7 @@ def create_backend( connected — call :meth:`initialize` before transferring). """ if method == "nccl_broadcast": - from .nccl_broadcast import NCCLBroadcastBackend + from .nccl_broadcast import NCCLBroadcastBackend # noqa: PLC0415 return NCCLBroadcastBackend(config, **kwargs) raise ValueError(f"Unknown weight sync backend: {method!r}. Supported: 'nccl_broadcast'.") diff --git a/src/xorl/server/weight_sync/backends/nccl_broadcast.py b/src/xorl/server/weight_sync/backends/nccl_broadcast.py index f64471ae..31f4c697 100644 --- a/src/xorl/server/weight_sync/backends/nccl_broadcast.py +++ b/src/xorl/server/weight_sync/backends/nccl_broadcast.py @@ -11,6 +11,7 @@ the pluggable backend interface. """ +import inspect import logging import os import time @@ -22,6 +23,13 @@ import requests import torch import torch.distributed as dist +from torch.distributed import PrefixStore, TCPStore +from torch.distributed.distributed_c10d import ( + Backend, + _new_process_group_helper, + _world, + default_pg_timeout, +) logger = logging.getLogger(__name__) @@ -135,13 +143,6 @@ def _init_training_process_group(self) -> dist.ProcessGroup: Returns: The initialized process group """ - from torch.distributed import PrefixStore, TCPStore - from torch.distributed.distributed_c10d import ( - Backend, - _new_process_group_helper, - _world, - default_pg_timeout, - ) # Important: Only set NCCL_CUMEM_ENABLE=0 os.environ["NCCL_CUMEM_ENABLE"] = "0" @@ -183,7 +184,6 @@ def _init_training_process_group(self) -> dist.ProcessGroup: logger.info("[Training] TCPStore created, creating process group...") # Handle different PyTorch versions by inspecting the actual signature - import inspect _pg_params = inspect.signature(_new_process_group_helper).parameters if "backend_options" in _pg_params: diff --git a/src/xorl/server/weight_sync/handler.py b/src/xorl/server/weight_sync/handler.py index 071f68b5..3b1befea 100644 --- a/src/xorl/server/weight_sync/handler.py +++ b/src/xorl/server/weight_sync/handler.py @@ -32,6 +32,24 @@ import torch import torch.distributed as dist +from torch.distributed._tensor import DTensor +from torch.distributed.fsdp._fully_shard import FSDPModule + +from xorl.distributed.parallel_state import get_parallel_state +from xorl.lora.modules.base import LoraModule +from xorl.lora.modules.linear import LoraLinear +from xorl.models.layers.moe.experts import MoEExperts +from xorl.models.layers.moe.lora import MoEExpertsLoRA +from xorl.models.transformers.qwen3_5_shared import ( + has_linear_attention_layers, + remap_linear_attention_params_for_inference, +) +from xorl.qlora.modules.linear import QLoRALinear +from xorl.qlora.modules.moe_experts import QLoRAMoeExperts +from xorl.server.protocol.operations import SyncWeightsData +from xorl.server.weight_sync.backends import TransportConfig, create_backend +from xorl.server.weight_sync.backends.base import EndpointConfig +from xorl.server.weight_sync.endpoint_manager import EndpointManager logger = logging.getLogger(__name__) @@ -71,8 +89,6 @@ async def handle_sync_inference_weights(self, command_dict: Dict[str, Any]) -> D """ logger.info(f"Rank {self.rank}: [WeightSync] Starting sync_inference_weights") - from xorl.server.protocol.operations import SyncWeightsData - p: SyncWeightsData = command_dict.get("payload", SyncWeightsData()) endpoints = p.endpoints @@ -139,7 +155,6 @@ def _sync_weights( 4. Sends the buffer via the transport backend 5. Pipelines: unshard(N+1) overlaps with transfer(N) """ - from torch.distributed.fsdp._fully_shard import FSDPModule model = self.trainer.model device = f"cuda:{self.trainer.local_rank}" @@ -172,9 +187,6 @@ def _sync_weights( # ------------------------------------------------------------------ # Step 3: Create backend + endpoint manager, initialize on sender ranks # ------------------------------------------------------------------ - from xorl.server.weight_sync.backends import TransportConfig, create_backend - from xorl.server.weight_sync.backends.base import EndpointConfig - from xorl.server.weight_sync.endpoint_manager import EndpointManager transport_cfg = TransportConfig( endpoints=[ @@ -244,8 +256,6 @@ def _sync_weights( modules_to_sync.extend(layer_modules) # Detect EP mode - from xorl.distributed.parallel_state import get_parallel_state - _ps = get_parallel_state() _ep_enabled = _ps.ep_enabled and _ps.ep_size > 1 @@ -328,8 +338,6 @@ def _sync_weights( logger.info( f"Rank {self.rank}: [WeightSync] ep_moe_prefixes={ep_moe_prefixes} for {mod_name}" ) - from torch.distributed._tensor import DTensor - current_buffer = self._extract_params_for_sync( fsdp_mod, mod_name, @@ -521,10 +529,6 @@ def _qlora_collective_ops( collect_results = self.rank == 0 # Discover QLoRA modules (only those directly owned by this FSDP module, # not by child FSDP modules — child modules will be processed separately) - from torch.distributed.fsdp._fully_shard import FSDPModule - - from xorl.qlora.modules.linear import QLoRALinear - from xorl.qlora.modules.moe_experts import QLoRAMoeExperts child_fsdp_prefixes = set() for mname, mod in fsdp_mod.named_modules(): @@ -564,8 +568,6 @@ def _qlora_collective_ops( del base_w, delta # --- Phase 2: QLoRAMoeExperts (gather lora params, defer heavy work) --- - from xorl.distributed.parallel_state import get_parallel_state - ps = get_parallel_state() ep_enabled = ps.ep_enabled and ps.ep_size > 1 @@ -630,14 +632,8 @@ def _collect_ep_moe_data( Must be called between unshard() and reshard(). Returns list of context dicts for _gather_and_broadcast_ep_moe_experts. """ - from torch.distributed._tensor import DTensor # Skip modules owned by child FSDP modules (processed separately) - from torch.distributed.fsdp._fully_shard import FSDPModule - - from xorl.models.layers.moe.experts import MoEExperts - from xorl.models.layers.moe.lora import MoEExpertsLoRA - from xorl.qlora.modules.moe_experts import QLoRAMoeExperts child_fsdp_prefixes = set() for mname, mod in fsdp_mod.named_modules(): @@ -1234,19 +1230,11 @@ def _extract_params_for_sync( Supports both LoraLinear (dense LoRA) and MoEExpertsLoRA (MoE LoRA). """ - from xorl.lora.modules.base import LoraModule - from xorl.lora.modules.linear import LoraLinear - from xorl.models.layers.moe.experts import MoEExperts - from xorl.models.layers.moe.lora import MoEExpertsLoRA - from xorl.qlora.modules.linear import QLoRALinear - from xorl.qlora.modules.moe_experts import QLoRAMoeExperts - buffer = [] # Identify child FSDP module prefixes — their params will be processed # when that child module is unsharded separately. Parent unshard may # expose child params as plain tensors, so we can't rely on DTensor check. - from torch.distributed.fsdp._fully_shard import FSDPModule child_fsdp_prefixes = set() for mname, mod in fsdp_mod.named_modules(): @@ -1386,10 +1374,6 @@ def _unfuse_for_inference( - Qwen3.5 linear attention: remap split GatedDeltaNet params back to HF fused names (q_proj/k_proj/v_proj → in_proj_qkv, etc.) """ - from xorl.models.transformers.qwen3_5_shared import ( - has_linear_attention_layers, - remap_linear_attention_params_for_inference, - ) This splits fused weights back to individual projections before sending. """ @@ -1476,7 +1460,6 @@ def _get_fsdp_modules(model) -> Tuple[Optional[Any], List[Tuple[str, Any]]]: Root module holds non-layer params (embed, norm, head). Layer modules hold per-layer params. """ - from torch.distributed.fsdp._fully_shard import FSDPModule root_module = None layer_modules = [] diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index 3dd54f39..75ed57ea 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -11,6 +11,23 @@ import torch import torch.nn as nn +from xorl.distributed.torch_parallelize import build_parallelize_model as _parallelize +from xorl.lora import freeze_base_parameters +from xorl.lora.utils import inject_lora_into_model, inject_lora_into_model_with_moe +from xorl.models import build_foundation_model +from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules +from xorl.models.layers.rope import set_rope_native +from xorl.qlora import ( + detect_prequantized_block_fp8, + detect_prequantized_nvfp4, + inject_qlora_into_model, + maybe_load_and_quantize_moe_qlora, + maybe_load_prequantized_qlora, + maybe_quantize_qlora, +) +from xorl.qlora.modules.linear import QLoRALinear +from xorl.qlora.modules.moe_experts import QLoRAMoeExperts +from xorl.qlora.utils import _deregister_qlora_weights_from_fsdp from xorl.utils import helper @@ -144,8 +161,6 @@ def build_training_model( Returns a :class:`TrainingModelResult` with model, config, PP state, etc. """ - from xorl.distributed.torch_parallelize import build_parallelize_model as _parallelize - from xorl.models import build_foundation_model # ------------------------------------------------------------------ # 1. Build foundation model @@ -173,8 +188,6 @@ def build_training_model( # Set module-level flags for rope and activation if rope_native: - from xorl.models.layers.rope import set_rope_native - set_rope_native(True) logger.info_rank0("Using native RoPE (flash_attn fused kernel disabled)") if activation_native: @@ -289,8 +302,6 @@ def build_training_model( param.requires_grad = False helper.print_device_mem_info("VRAM usage after QLoRA quantization") elif enable_lora: - from xorl.lora import freeze_base_parameters - freeze_base_parameters(model) logger.info_rank0("Base model parameters frozen, only LoRA parameters trainable") else: @@ -350,11 +361,6 @@ def _inject_qlora( Returns (is_prequantized, checkpoint_quant_format, exclude_modules). """ - from xorl.qlora import ( - detect_prequantized_block_fp8, - detect_prequantized_nvfp4, - inject_qlora_into_model, - ) is_prequantized = False checkpoint_quant_format = None @@ -373,8 +379,6 @@ def _inject_qlora( exclude_modules = set(qlora_exclude_modules) logger.info_rank0(f"Using user-specified exclude_modules: {exclude_modules}") elif is_prequantized: - from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules - exclude_modules = get_prequantized_exclude_modules(weights_path) if exclude_modules: logger.info_rank0( @@ -437,8 +441,6 @@ def _inject_lora( is_moe_model = getattr(model.config, "num_experts", 0) > 0 if is_moe_model and (moe_shared_lora or moe_hybrid_shared_lora): - from xorl.lora.utils import inject_lora_into_model_with_moe - logger.info_rank0( f"MoE-aware LoRA injection (shared={moe_shared_lora}, hybrid_shared={moe_hybrid_shared_lora})" ) @@ -451,8 +453,6 @@ def _inject_lora( moe_hybrid_shared_lora=moe_hybrid_shared_lora, ) else: - from xorl.lora.utils import inject_lora_into_model - inject_lora_into_model( model, r=lora_rank, @@ -475,9 +475,6 @@ def _deferred_qlora_quantize( 2. NF4 linear: bf16 weight already loaded by FSDP → quantize in-place 3. NF4 MoE: load bf16 experts from checkpoint → quantize """ - from xorl.qlora.modules.linear import QLoRALinear - from xorl.qlora.modules.moe_experts import QLoRAMoeExperts - from xorl.qlora.utils import _deregister_qlora_weights_from_fsdp # 1. Pre-quantized linear/MoE loading (nvfp4/block_fp8) needs_prequant_linear = any( @@ -489,8 +486,6 @@ def _deferred_qlora_quantize( ) if needs_prequant_linear or needs_prequant_moe: - from xorl.qlora import maybe_load_prequantized_qlora - logger.info(f"Starting pre-quantized weight loading (mode={load_weights_mode}) …") helper.print_device_mem_info("VRAM before pre-quantized loading") maybe_load_prequantized_qlora(model, weights_path, load_mode=load_weights_mode) @@ -499,8 +494,6 @@ def _deferred_qlora_quantize( # 2. NF4 linear: FSDP loaded bf16 into weight param → quantize to NF4 needs_bf16_quantize = any(isinstance(m, QLoRALinear) and m.weight is not None for m in model.modules()) if needs_bf16_quantize: - from xorl.qlora import maybe_quantize_qlora - logger.info("Quantizing bf16 linear weights to NF4 …") maybe_quantize_qlora(model) @@ -512,8 +505,6 @@ def _deferred_qlora_quantize( for m in model.modules() ) if needs_bf16_moe: - from xorl.qlora import maybe_load_and_quantize_moe_qlora - logger.info("Loading and quantizing bf16 MoE expert weights …") maybe_load_and_quantize_moe_qlora( model, diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index 730f1b2d..b766a0fc 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -17,9 +17,11 @@ import torch import torch.distributed as dist +import torch.distributed.tensor._random +from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict from tqdm import trange -from xorl.arguments import Arguments +from xorl.arguments import Arguments, save_args from xorl.checkpoint import build_checkpointer from xorl.data.constants import IGNORE_INDEX from xorl.data.data_loader import DataLoaderBuilder @@ -27,13 +29,25 @@ from xorl.distributed.gradient_accumulate_loss import gradient_accumulate_loss from xorl.distributed.offloading import build_activation_offloading_context from xorl.distributed.parallel_state import get_parallel_state, init_parallel_state +from xorl.distributed.pipeline_parallel import build_pipeline_schedule, build_pp_stage from xorl.distributed.sync_padding import synchronize_micro_batch_padding from xorl.distributed.torch_parallelize import build_parallelize_model +from xorl.lora.utils import inject_lora_into_model, save_lora_checkpoint from xorl.models import build_foundation_model, build_tokenizer, save_model_assets, save_model_weights +from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules from xorl.models.layers.moe.aux_loss import global_load_balancing_loss_func from xorl.models.layers.moe.routing_replay import RoutingReplay, set_replay_stage from xorl.models.module_utils import compute_loss from xorl.optim import build_lr_scheduler, build_optimizer +from xorl.qlora import ( + detect_prequantized_block_fp8, + detect_prequantized_nvfp4, + inject_qlora_into_model, + maybe_load_and_quantize_moe_qlora, + maybe_load_prequantized_qlora, + maybe_quantize_qlora, +) +from xorl.qlora.utils import _deregister_qlora_weights_from_fsdp from xorl.trainers.model_builder import ( maybe_upcast_trainable_adapter_params, resolve_training_model_dtype, @@ -160,11 +174,9 @@ def _bootstrap(self) -> None: helper.enable_third_party_logging() if args.train.global_rank == 0: - from xorl.arguments import save_args - save_args(args, args.train.output_dir) if args.train.use_wandb: - import wandb + import wandb # noqa: PLC0415 wandb.init( project=args.train.wandb_project, @@ -214,8 +226,6 @@ def _bootstrap(self) -> None: # DTensor RNG tracker (run_state_sync=False to avoid PP deadlock) self.ps = get_parallel_state() if self.ps.device_mesh is not None: - import torch.distributed.tensor._random - torch.distributed.tensor._random.manual_seed(args.train.seed, self.ps.device_mesh) # Routing replay is only needed with EP when MoE forward is recomputed @@ -229,7 +239,7 @@ def _maybe_log_startup_metrics(self, metrics: Dict[str, Any], commit: bool = Fal self._startup_metrics.update(metrics) if self.args.train.global_rank != 0 or not self.args.train.use_wandb or not self._wandb_initialized: return - import wandb + import wandb # noqa: PLC0415 wandb.log(metrics, step=0, commit=commit) @@ -293,7 +303,7 @@ def _log_host_inventory(self) -> None: ) self._maybe_log_startup_metrics({"startup/node_count": len(unique_hostnames)}, commit=False) if self.args.train.use_wandb and self._wandb_initialized: - import wandb + import wandb # noqa: PLC0415 wandb.config.update( { @@ -420,7 +430,6 @@ def _build_model(self) -> None: def _inject_qlora(self) -> None: """QLoRA injection with pre-quantized checkpoint detection.""" args = self.args - from xorl.qlora import detect_prequantized_block_fp8, detect_prequantized_nvfp4, inject_qlora_into_model if detect_prequantized_nvfp4(args.model.model_path): self.is_prequantized = True @@ -435,8 +444,6 @@ def _inject_qlora(self) -> None: self.exclude_modules = set(args.lora.exclude_modules) logger.info_rank0(f"Using user-specified exclude_modules: {self.exclude_modules}") elif self.is_prequantized: - from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules - self.exclude_modules = get_prequantized_exclude_modules(args.model.model_path) if self.exclude_modules: logger.info_rank0( @@ -470,7 +477,6 @@ def _inject_qlora(self) -> None: def _inject_lora(self) -> None: """Plain LoRA injection.""" args = self.args - from xorl.lora.utils import inject_lora_into_model inject_lora_into_model( self.model, @@ -541,16 +547,11 @@ def _deferred_qlora_quantize(self) -> None: """After FSDP loads weights, quantize them into uint8 buffers.""" args = self.args if self.is_prequantized: - from xorl.qlora import maybe_load_prequantized_qlora - logger.info("Starting pre-quantized weight loading...") helper.print_device_mem_info("VRAM before pre-quantized loading") maybe_load_prequantized_qlora(self.model, args.model.model_path) logger.info("Done pre-quantized weight loading, freezing non-LoRA params...") else: - from xorl.qlora import maybe_load_and_quantize_moe_qlora, maybe_quantize_qlora - from xorl.qlora.utils import _deregister_qlora_weights_from_fsdp - logger.info("Starting maybe_quantize_qlora...") helper.print_device_mem_info("VRAM before QLoRA quantization") maybe_quantize_qlora(self.model) @@ -678,8 +679,6 @@ def _get_pp_schedule(self, seq_len: int): With static padding, seq_len is always the same so only one entry is cached. """ if seq_len not in self._pp_schedule_cache: - from xorl.distributed.pipeline_parallel import build_pipeline_schedule, build_pp_stage - stage = build_pp_stage( self.model_parts[0], pp_rank=self.ps.pp_rank, @@ -1045,7 +1044,7 @@ def _maybe_log( and args.train.use_wandb and self.state.global_step % args.train.wandb_log_interval == 0 ): - import wandb + import wandb # noqa: PLC0415 train_metrics.update( { @@ -1126,8 +1125,6 @@ def _finalize(self, total_loss: float, grad_norm: float, lr: float) -> None: hf_model_state_dict = None if args.train.save_hf_weights and not save_peft_adapter: - from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict - logger.info_rank0("Gathering full model state dict for HF checkpoint via NCCL with CPU offload...") hf_model_state_dict = get_model_state_dict( self.model, options=StateDictOptions(full_state_dict=True, cpu_offload=True) @@ -1140,8 +1137,6 @@ def _finalize(self, total_loss: float, grad_norm: float, lr: float) -> None: if args.train.global_rank == 0 and args.train.save_hf_weights: hf_weights_path = os.path.join(args.train.output_dir, f"global_step_{state.global_step}", "hf_ckpt") if save_peft_adapter: - from xorl.lora.utils import save_lora_checkpoint - save_lora_checkpoint( self.model, hf_weights_path, diff --git a/src/xorl/trainers/training_utils.py b/src/xorl/trainers/training_utils.py index 0453c050..2a85cc50 100644 --- a/src/xorl/trainers/training_utils.py +++ b/src/xorl/trainers/training_utils.py @@ -4,6 +4,7 @@ counting, LoRA merge, PP forward-backward) into reusable free functions. """ +import logging from collections import deque from typing import Any, Dict, List, Optional @@ -12,6 +13,9 @@ import torch.nn.functional as F from xorl.data.constants import IGNORE_INDEX +from xorl.lora.modules.base import LoraModule +from xorl.lora.utils import maybe_merge_lora as _merge +from xorl.qlora.utils import maybe_requant_qlora from xorl.utils.device import get_device_type @@ -102,7 +106,6 @@ def reset_lora_optimizer_states( Returns: Number of parameters whose optimizer states were cleared. """ - from xorl.lora.modules.base import LoraModule # Collect LoRA parameter ids lora_param_ids = set() @@ -142,19 +145,13 @@ def maybe_merge_lora( if merge_interval <= 0 or global_step % merge_interval != 0: return if enable_qlora: - from xorl.qlora.utils import maybe_requant_qlora - maybe_requant_qlora(model) elif enable_lora: - from xorl.lora.utils import maybe_merge_lora as _merge - _merge(model) if reset_optimizer and optimizer is not None: count = reset_lora_optimizer_states(model, optimizer) if count > 0: - import logging - logging.getLogger(__name__).info(f"ReLoRA optimizer reset: pruned states for {count} LoRA parameters") @@ -204,7 +201,6 @@ def pad_micro_batches_for_pp( adding a separate all-zero "padding document") to avoid FA3 varlen backward NaN from degenerate inputs and stale max_length_q/k. """ - import torch.nn.functional as F if sample_packing_sequence_len <= 0: return diff --git a/src/xorl/utils/helper.py b/src/xorl/utils/helper.py index f0df2ef3..32bf93c5 100644 --- a/src/xorl/utils/helper.py +++ b/src/xorl/utils/helper.py @@ -20,6 +20,12 @@ from transformers import enable_full_determinism from transformers import set_seed as set_seed_func + +try: + from pyiceberg.metrics import LoggingMetricsReporter +except ImportError: + LoggingMetricsReporter = None + from xorl.distributed.parallel_state import get_parallel_state from xorl.utils import logging from xorl.utils.count_flops import XorlFlopsCounter @@ -28,6 +34,9 @@ get_device_type, get_torch_device, ) +from xorl.utils.device import ( + empty_cache as device_empty_cache, +) from xorl.utils.dist_utils import all_reduce from .multisource_utils import parse_multisource_config @@ -41,6 +50,7 @@ logger = logging.get_logger(__name__) CACHE_DIR = os.path.expanduser(os.getenv("CACHE_DIR", os.path.join("~/.cache", "xorl"))) +IS_NPU_AVAILABLE = hasattr(torch, "npu") and torch.npu.is_available() def _get_global_token_count(micro_batch: Dict[str, "torch.Tensor"]) -> List[int]: @@ -337,10 +347,10 @@ def disable_warning() -> None: """ Enables warning filter. """ - from pyiceberg.metrics import LoggingMetricsReporter - builtin_logging.basicConfig(level=builtin_logging.ERROR) warnings.simplefilter("ignore") + if LoggingMetricsReporter is None: + return LoggingMetricsReporter() LoggingMetricsReporter._logger = builtin_logging.getLogger(LoggingMetricsReporter.__name__) LoggingMetricsReporter._logger.setLevel(builtin_logging.WARNING) @@ -373,10 +383,10 @@ def empty_cache() -> None: """ gc.collect() - if IS_CUDA_AVAILABLE or IS_NPU_AVAILABLE: - from xorl.utils.device import empty_cache - - empty_cache() + if IS_CUDA_AVAILABLE: + device_empty_cache() + elif IS_NPU_AVAILABLE: + torch.npu.empty_cache() def get_cache_dir(path: Optional[str] = None) -> str: diff --git a/src/xorl/utils/import_utils.py b/src/xorl/utils/import_utils.py index e37e3bc8..677abe73 100644 --- a/src/xorl/utils/import_utils.py +++ b/src/xorl/utils/import_utils.py @@ -5,6 +5,7 @@ from functools import lru_cache from typing import TYPE_CHECKING, Dict +import torch from packaging import version @@ -59,8 +60,6 @@ def is_flash_attn_available() -> bool: def is_fused_moe_available() -> bool: - import torch - return torch.cuda.is_available() and _PACKAGE_FLAGS["triton"] diff --git a/src/xorl/utils/recompute_utils.py b/src/xorl/utils/recompute_utils.py index d9d16626..5fdee02b 100644 --- a/src/xorl/utils/recompute_utils.py +++ b/src/xorl/utils/recompute_utils.py @@ -35,18 +35,6 @@ def string_to_op(op_string: str) -> Any: current = torch.ops - # Special handling: ensure accessing aten operations by first trying to trigger registration - if parts[0] == "aten": - try: - # Try to access a basic aten operation to trigger module loading - _ = torch.ops.aten.add - except AttributeError: - # If cannot access aten, may need to import related modules - try: - import torch._C._dispatch - except ImportError: - pass - for i, part in enumerate(parts): if hasattr(current, part): current = getattr(current, part) @@ -113,10 +101,3 @@ def _check_torch_ops_availability(): logger.info_rank0("✓ torch.ops.aten available") except AttributeError as e: logger.info_rank0(f"✗ torch.ops.aten not available: {e}") - logger.info_rank0("Trying to import necessary modules...") - try: - import torch._C._dispatch - - logger.info_rank0("✓ Successfully imported torch._C._dispatch") - except ImportError as e: - logger.info_rank0(f"✗ Cannot import torch._C._dispatch: {e}") diff --git a/tests/data/prepare/test_shared.py b/tests/data/prepare/test_shared.py index b9676908..f163411c 100644 --- a/tests/data/prepare/test_shared.py +++ b/tests/data/prepare/test_shared.py @@ -4,6 +4,7 @@ import pytest from datasets import Dataset as HFDataset +from huggingface_hub.errors import RepositoryNotFoundError from xorl.arguments import DatasetConfig from xorl.data.prepare.shared import ( @@ -13,7 +14,9 @@ datasets_with_name_generator, get_dataset_type, load_dataset_with_config, + load_preprocessed_dataset, merge_datasets, + save_preprocessed_dataset, ) @@ -90,7 +93,6 @@ def test_hub_detection_and_remote_filesystem(self, mock_snapshot_download): assert _check_if_hub_dataset(_make_config(path="user/ds"), use_auth_token=False) is True # Invalid hub dataset - from huggingface_hub.errors import RepositoryNotFoundError mock_response = Mock() mock_response.status_code = 404 @@ -238,7 +240,6 @@ class TestSaveAndLoadPreprocessedDataset: def test_save_load_and_missing(self, tmp_path): """Covers save+load round-trip and load returning None when not found.""" - from xorl.data.prepare.shared import load_preprocessed_dataset, save_preprocessed_dataset args = Mock() args.data.dataset_prepared_path = str(tmp_path) diff --git a/tests/data/test_data_loader.py b/tests/data/test_data_loader.py index 52af64b6..6a6d74d1 100644 --- a/tests/data/test_data_loader.py +++ b/tests/data/test_data_loader.py @@ -5,7 +5,7 @@ import torch from torch.utils.data import Dataset -from xorl.data.collators import DataCollator, PackingConcatCollator +from xorl.data.collators import CollatePipeline, DataCollator, FlattenCollator, PackingConcatCollator from xorl.data.data_loader import ( DataLoaderBuilder, DistributedDataloader, @@ -128,7 +128,6 @@ def test_builder_batch_size_sampler_sp_and_custom_collators( ) collate_fn = mock_dataloader_cls.call_args[1]["collate_fn"] assert isinstance(collate_fn, MicroBatchCollator) - from xorl.data.collators import CollatePipeline assert isinstance(collate_fn.internal_collator, CollatePipeline) assert len(collate_fn.internal_collator.data_collators) == 5 @@ -344,7 +343,6 @@ def __getitem__(self, idx): assert mb["input_ids"].shape == (1, 16) # Flatten + Packing collator - from xorl.data.collators import FlattenCollator flatten_collator = FlattenCollator() packing_collator = PackingConcatCollator(pad_to_multiple_of=1) diff --git a/tests/data/test_data_loader_distributed.py b/tests/data/test_data_loader_distributed.py index 515483d2..cfa67ef3 100644 --- a/tests/data/test_data_loader_distributed.py +++ b/tests/data/test_data_loader_distributed.py @@ -9,8 +9,11 @@ import pytest import torch +from torch.utils.data import Dataset from tests.conftest import FakeTextDataset, SimpleCollator +from xorl.data.collators import CollatePipeline, TextSequenceShardCollator +from xorl.data.constants import IGNORE_INDEX from xorl.data.data_loader import ( DataLoaderBuilder, MicroBatchCollator, @@ -86,7 +89,6 @@ def test_partitioning_micro_batches_sp_batch_size_and_drop_last( mock_sampler_cls.return_value = mock_sampler dataloader_sp = builder_sp.build(verbose=False) - from xorl.data.collators import CollatePipeline, TextSequenceShardCollator internal = builder_sp.collate_fn.internal_collator assert isinstance(internal, CollatePipeline) @@ -157,8 +159,6 @@ def test_order_preservation_ga_alignment_and_error(self): @patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") def test_sequence_sharding_and_padding(self, mock_parallel_state): """Covers correct chunk sizes across SP ranks and padding for non-divisible lengths.""" - from xorl.data.collators import TextSequenceShardCollator - from xorl.data.constants import IGNORE_INDEX cp_size = 4 seq_len = 128 @@ -271,7 +271,6 @@ def test_pipeline_epoch_consistency_and_packed_sequences( assert len(epoch1) == len(epoch2) # --- Packed sequence tests --- - from torch.utils.data import Dataset class PackedDataset(Dataset): def __len__(self): diff --git a/tests/distributed/test_deepep_correctness.py b/tests/distributed/test_deepep_correctness.py index a015493b..91c88fe7 100644 --- a/tests/distributed/test_deepep_correctness.py +++ b/tests/distributed/test_deepep_correctness.py @@ -83,7 +83,7 @@ def run_one_pass( buffer=None, ): """One forward + backward pass with a given dispatch backend.""" - from xorl.models.layers.moe.backend import EP_COMBINE, EP_DISPATCH, EP_EXPERT_COMPUTE + from xorl.models.layers.moe.backend import EP_COMBINE, EP_DISPATCH, EP_EXPERT_COMPUTE # noqa: PLC0415 dispatch_fn = EP_DISPATCH[ep_dispatch] combine_fn = EP_COMBINE[ep_dispatch] @@ -181,7 +181,7 @@ def run_test( down.grad = None # ── DeepEP ─────────────────────────────────────────────────────────────── - from xorl.distributed.moe.deepep import DeepEPBuffer + from xorl.distributed.moe.deepep import DeepEPBuffer # noqa: PLC0415 buffer = DeepEPBuffer(ep_group=ep_group, buffer_size_gb=1.0) out_dep, grad_dep = run_one_pass( diff --git a/tests/distributed/test_ep_lora_weight_slicing.py b/tests/distributed/test_ep_lora_weight_slicing.py index a70696e8..8dcef965 100644 --- a/tests/distributed/test_ep_lora_weight_slicing.py +++ b/tests/distributed/test_ep_lora_weight_slicing.py @@ -15,6 +15,7 @@ from torch.distributed._tensor import Shard from xorl.models.layers.moe import MoEExpertsLoRA, MoELoRAConfig +from xorl.models.transformers.qwen3_moe.parallelize import get_ep_plan pytestmark = [pytest.mark.distributed] @@ -75,7 +76,6 @@ class TestParallelPlanLoRASlicing: def test_ep_plan_and_shard_tensor(self): """EP plan includes all LoRA patterns with Shard(0); shard_tensor slices by ep_rank.""" - from xorl.models.transformers.qwen3_moe.parallelize import get_ep_plan plan = get_ep_plan() @@ -120,7 +120,7 @@ class TestEPLoRAForwardAndGradients: @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_ep_and_non_ep_forward_with_gradients(self): """Both EP and non-EP LoRA forward produce correct output shapes with gradient flow.""" - from xorl.ops.moe.triton_lora import TritonEPGroupGemmWithLoRA, TritonMoeExpertsLoRAFunction + from xorl.ops.moe.triton_lora import TritonEPGroupGemmWithLoRA, TritonMoeExpertsLoRAFunction # noqa: PLC0415 device = "cuda" dtype = torch.bfloat16 diff --git a/tests/distributed/test_parallel_state.py b/tests/distributed/test_parallel_state.py index 1442df0d..cbcab429 100644 --- a/tests/distributed/test_parallel_state.py +++ b/tests/distributed/test_parallel_state.py @@ -92,12 +92,12 @@ class TestGetAndInitParallelState: """Test get_parallel_state and init_parallel_state functions.""" def setup_method(self): - import xorl.distributed.parallel_state as ps_module + import xorl.distributed.parallel_state as ps_module # noqa: PLC0415 ps_module._PARALLEL_STATE = None def teardown_method(self): - import xorl.distributed.parallel_state as ps_module + import xorl.distributed.parallel_state as ps_module # noqa: PLC0415 ps_module._PARALLEL_STATE = None diff --git a/tests/distributed/test_tensor_parallel.py b/tests/distributed/test_tensor_parallel.py index bc82a893..ca310b7c 100644 --- a/tests/distributed/test_tensor_parallel.py +++ b/tests/distributed/test_tensor_parallel.py @@ -8,6 +8,9 @@ import pytest import torch +from xorl.models.transformers.qwen3.configuration_qwen3 import Qwen3Config +from xorl.models.transformers.qwen3.modeling_qwen3 import Qwen3Attention, Qwen3ForCausalLM, Qwen3MLP + pytestmark = [pytest.mark.cpu, pytest.mark.distributed] @@ -17,8 +20,6 @@ class TestUnfuseForTP: def test_attention_and_mlp_unfuse(self): """Attention unfuse creates separate q/k/v; MLP unfuse creates separate gate/up.""" - from xorl.models.transformers.qwen3.configuration_qwen3 import Qwen3Config - from xorl.models.transformers.qwen3.modeling_qwen3 import Qwen3Attention, Qwen3MLP config = Qwen3Config( hidden_size=256, @@ -51,8 +52,6 @@ def test_attention_and_mlp_unfuse(self): def test_unfused_forward_shape_and_model_level(self): """Unfused MLP forward produces same shape; model-level unfuse covers all layers.""" - from xorl.models.transformers.qwen3.configuration_qwen3 import Qwen3Config - from xorl.models.transformers.qwen3.modeling_qwen3 import Qwen3ForCausalLM, Qwen3MLP config = Qwen3Config( hidden_size=256, diff --git a/tests/e2e/e2e_utils.py b/tests/e2e/e2e_utils.py index c27247bf..d45f7289 100644 --- a/tests/e2e/e2e_utils.py +++ b/tests/e2e/e2e_utils.py @@ -21,6 +21,16 @@ import pytest import torch import yaml +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from transformers import AutoConfig, AutoModelForCausalLM + + +try: + import xorl_client +except ModuleNotFoundError: + xorl_client = None # --------------------------------------------------------------------------- @@ -240,18 +250,12 @@ def create_tiny_model_dir( def _save_random_weights(model_dir: str, model_config: dict): - from transformers import AutoConfig, AutoModelForCausalLM - config = AutoConfig.from_pretrained(model_dir) model = AutoModelForCausalLM.from_config(config) model.save_pretrained(model_dir, safe_serialization=True) def _create_tokenizer_files(model_dir: str, vocab_size: int = 1024): - from tokenizers import Tokenizer - from tokenizers.models import WordLevel - from tokenizers.pre_tokenizers import Whitespace - vocab = {f"t{i}": i for i in range(vocab_size)} vocab["[UNK]"] = 0 vocab["[PAD]"] = vocab_size - 1 if vocab_size > 1 else 0 @@ -534,7 +538,8 @@ def samples_to_xorl_datums( samples: List[Tuple[List[int], List[int]]], ) -> list: """Convert raw (input_ids, labels) pairs to xorl_client.Datum objects.""" - import xorl_client + if xorl_client is None: + pytest.skip("xorl_client not installed") datums = [] for input_ids, labels in samples: diff --git a/tests/e2e/qwen3_8b/test_tflops_threshold.py b/tests/e2e/qwen3_8b/test_tflops_threshold.py index 94ec09b8..10856469 100644 --- a/tests/e2e/qwen3_8b/test_tflops_threshold.py +++ b/tests/e2e/qwen3_8b/test_tflops_threshold.py @@ -17,6 +17,13 @@ import time import pytest +from transformers import AutoConfig + + +try: + import xorl_client +except ModuleNotFoundError: + xorl_client = None from tests.e2e.e2e_utils import skip_if_gpu_count_less_than from tests.e2e.server_utils import ( @@ -28,6 +35,7 @@ generate_random_sft_data, generate_server_config, ) +from xorl.utils.count_flops import XorlFlopsCounter, get_device_flops pytestmark = [pytest.mark.e2e, pytest.mark.gpu, pytest.mark.server, pytest.mark.benchmark] @@ -69,10 +77,8 @@ def _skip_if_no_qwen3_8b(): def _run_tflops_benchmark(num_gpus: int, dp_shard_size: int, tmp_path: str): """Run a Qwen3-8B LoRA SFT benchmark and return TFLOPS.""" - import xorl_client - from transformers import AutoConfig - - from xorl.utils.count_flops import XorlFlopsCounter, get_device_flops + if xorl_client is None: + pytest.skip("xorl_client not installed") output_dir = os.path.join(tmp_path, f"bench_{num_gpus}gpu") api_port = _get_free_port() diff --git a/tests/e2e/server_utils.py b/tests/e2e/server_utils.py index 6a515fec..7d57ebb6 100644 --- a/tests/e2e/server_utils.py +++ b/tests/e2e/server_utils.py @@ -15,6 +15,18 @@ import yaml +try: + import xorl_client +except ModuleNotFoundError: + xorl_client = None + + +def _require_xorl_client(): + if xorl_client is None: + pytest.skip("xorl_client not installed") + return xorl_client + + def generate_server_config( model_dir: str, output_dir: str, @@ -296,7 +308,7 @@ def generate_random_sft_data( Returns list of Datum objects ready for the training client. Uses causallm_loss format: input_ids + labels (shifted tokens). """ - import xorl_client + xorl_client = _require_xorl_client() rng = random.Random(seed) datums = [] @@ -357,7 +369,7 @@ def extract_loss(fwd_bwd_result) -> float: def run_sft_steps(training_client, data, num_steps=5, lr=1e-3) -> list: """Run SFT training steps and return loss history.""" - import xorl_client + xorl_client = _require_xorl_client() adam_params = xorl_client.AdamParams(learning_rate=lr, beta1=0.9, beta2=0.95, eps=1e-8) losses = [] @@ -396,7 +408,7 @@ def _start_server_or_fail(server, timeout=180.0): def _create_lora_client(base_url, model_dir, model_id="test", rank=8): """Create a xorl_client LoRA training client.""" - import xorl_client + xorl_client = _require_xorl_client() service_client = xorl_client.ServiceClient(base_url=base_url) training_client = service_client.create_lora_training_client( @@ -413,7 +425,7 @@ def _create_full_weight_client(base_url, model_dir): The server auto-registers model_id="default" on startup, so no create_model call is needed for full-weight training. """ - import xorl_client + xorl_client = _require_xorl_client() service_client = xorl_client.ServiceClient(base_url=base_url) training_client = xorl_client.TrainingClient( diff --git a/tests/models/test_moe_experts_lora.py b/tests/models/test_moe_experts_lora.py index a84d17fb..f5a1af71 100644 --- a/tests/models/test_moe_experts_lora.py +++ b/tests/models/test_moe_experts_lora.py @@ -2,6 +2,8 @@ import pytest +from xorl.lora import LoraLinear, inject_lora_into_model + pytestmark = [pytest.mark.cpu, pytest.mark.gpu] import torch @@ -560,7 +562,6 @@ def test_from_module_and_inject_lora(self, backend): assert torch.allclose(lora_exp.gate_proj, base.gate_proj) # inject_lora_into_model - from xorl.lora import inject_lora_into_model class SimpleModel(nn.Module): def __init__(self, config, backend): @@ -608,7 +609,6 @@ def test_from_module_with_qwen3_subclass(self): def test_injection_error_handling(self): """Test error cases and valid injection for inject_lora_into_model.""" - from xorl.lora import LoraLinear, inject_lora_into_model # No matching modules class ModelA(nn.Module): diff --git a/tests/models/test_moe_routing_cache.py b/tests/models/test_moe_routing_cache.py index 4c7085fe..336f66cc 100644 --- a/tests/models/test_moe_routing_cache.py +++ b/tests/models/test_moe_routing_cache.py @@ -12,6 +12,8 @@ import torch.nn as nn from torch.utils.checkpoint import checkpoint +from xorl.models.base import XorlPreTrainedModel + pytestmark = [pytest.mark.gpu] @@ -87,7 +89,7 @@ class TestContextFnMechanism: def test_forward_sets_mode(self): """During checkpoint forward, mode is 'forward'.""" - import xorl.models.layers.moe.moe_block as mb + import xorl.models.layers.moe.moe_block as mb # noqa: PLC0415 observed = [] @@ -111,7 +113,7 @@ def forward(self, x): def test_mode_none_without_context_fn(self): """Without context_fn, mode stays None — no caching.""" - import xorl.models.layers.moe.moe_block as mb + import xorl.models.layers.moe.moe_block as mb # noqa: PLC0415 observed = [] @@ -389,8 +391,6 @@ def __init__(self): self.layer = _SimpleDecoderLayer() self.gradient_checkpointing = False - from xorl.models.base import XorlPreTrainedModel - # Monkey-patch a minimal model model = XorlPreTrainedModel.__new__(XorlPreTrainedModel) nn.Module.__init__(model) diff --git a/tests/models/test_moe_routing_replay.py b/tests/models/test_moe_routing_replay.py index 0b8ba6ee..a5fa9c74 100644 --- a/tests/models/test_moe_routing_replay.py +++ b/tests/models/test_moe_routing_replay.py @@ -12,6 +12,8 @@ import torch.nn as nn from torch.utils.checkpoint import checkpoint +from xorl.models.base import XorlPreTrainedModel + pytestmark = [pytest.mark.gpu] @@ -455,7 +457,6 @@ class TestBaseModelIntegrationAndR3Preload: def test_enable_routing_replay_checkpoint_e2e_and_r3_preload(self): """Test enable_routing_replay, gradient_checkpointing_enable, full e2e, and R3 preload.""" - from xorl.models.base import XorlPreTrainedModel class _FakeConfig: pass diff --git a/tests/models/test_moe_weight_loading_integration.py b/tests/models/test_moe_weight_loading_integration.py index a2328943..6b00a675 100644 --- a/tests/models/test_moe_weight_loading_integration.py +++ b/tests/models/test_moe_weight_loading_integration.py @@ -5,6 +5,7 @@ into a model that expects fused (stacked) expert format. """ +import random from typing import Dict, Iterator, Tuple import pytest @@ -269,8 +270,6 @@ def test_load_and_verify_shapes_and_order(self): state_dict2[f"model.layers.0.mlp.experts.{expert_idx}.up_proj.weight"] = torch.randn(16, 8) state_dict2[f"model.layers.0.mlp.experts.{expert_idx}.down_proj.weight"] = torch.randn(8, 16) - import random - items = list(state_dict2.items()) random.seed(42) random.shuffle(items) @@ -313,7 +312,6 @@ def test_streaming_and_edge_cases(self): assert loaded_single["model.layers.0.mlp.experts.gate_proj"].shape == (1, 8, 16) # Random order loading - import random model_rand = MockMoeModel(2, 4, 8, 16) state_dict_rand = create_per_expert_state_dict(2, 4, 8, 16) diff --git a/tests/models/test_qwen3_moe_fused_lora.py b/tests/models/test_qwen3_moe_fused_lora.py index 53c9c3e3..f8e109f4 100644 --- a/tests/models/test_qwen3_moe_fused_lora.py +++ b/tests/models/test_qwen3_moe_fused_lora.py @@ -4,6 +4,7 @@ import torch import torch.nn as nn +from xorl.lora import LoraLinear, inject_lora_into_model from xorl.lora.mapping import can_apply_lora, get_lora_class_for_module from xorl.models.layers.moe import MOE_EXPERT_BACKENDS, MoEBlock, MoEExperts, MoEExpertsLoRA, MoELoRAConfig from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import ( @@ -516,7 +517,6 @@ def test_from_module_inject_lora_and_errors(self): assert torch.allclose(lora_exp.gate_proj, base.gate_proj) # inject_lora_into_model - from xorl.lora import inject_lora_into_model class SimpleModel(nn.Module): def __init__(self, config, backend): @@ -561,7 +561,6 @@ def __init__(self, config, backend): assert torch.allclose(lora_exp.gate_proj, base.gate_proj) # Error handling - from xorl.lora import LoraLinear, inject_lora_into_model class ModelA(nn.Module): def __init__(self): diff --git a/tests/ops/test_attention.py b/tests/ops/test_attention.py index 967f338b..7de1b021 100644 --- a/tests/ops/test_attention.py +++ b/tests/ops/test_attention.py @@ -182,7 +182,7 @@ def test_flash_attention_on_gpu(self): if not torch.cuda.is_available(): pytest.skip("CUDA not available") try: - from flash_attn import flash_attn_func # noqa: F401 + from flash_attn import flash_attn_func # noqa: F401, PLC0415 except ImportError: pytest.skip("flash-attn not installed") diff --git a/tests/ops/test_eager_vs_native_moe.py b/tests/ops/test_eager_vs_native_moe.py index 46faee23..758aea12 100644 --- a/tests/ops/test_eager_vs_native_moe.py +++ b/tests/ops/test_eager_vs_native_moe.py @@ -22,8 +22,8 @@ def _import_moe(): """Import MoE layers; returns (MoEBlock, MoEExperts) or skips.""" try: - from xorl.models.layers.moe.experts import MoEExperts - from xorl.models.layers.moe.moe_block import MoEBlock + from xorl.models.layers.moe.experts import MoEExperts # noqa: PLC0415 + from xorl.models.layers.moe.moe_block import MoEBlock # noqa: PLC0415 return MoEBlock, MoEExperts except Exception as e: @@ -146,7 +146,7 @@ def test_determinism_and_edge_cases(): assert torch.equal(out1, out2), f"{backend} is not deterministic" # --- All tokens to same expert --- - from xorl.models.layers.moe.backend.native import native_expert_forward + from xorl.models.layers.moe.backend.native import native_expert_forward # noqa: PLC0415 ne, hd, inter = 4, 64, 128 torch.manual_seed(42) diff --git a/tests/ops/test_group_gemm.py b/tests/ops/test_group_gemm.py index 7772b630..f8b3fcde 100644 --- a/tests/ops/test_group_gemm.py +++ b/tests/ops/test_group_gemm.py @@ -92,7 +92,7 @@ def test_same_nk_comprehensive(self): if not torch.cuda.is_available(): pytest.skip("CUDA not available") try: - from xorl.ops.group_gemm.kernel.group_gemm import group_gemm_same_nk + from xorl.ops.group_gemm.kernel.group_gemm import group_gemm_same_nk # noqa: PLC0415 except ImportError: pytest.skip("group_gemm not available") @@ -158,7 +158,7 @@ def test_same_mn_comprehensive(self): if not torch.cuda.is_available(): pytest.skip("CUDA not available") try: - from xorl.ops.group_gemm.kernel.group_gemm import group_gemm_same_mn + from xorl.ops.group_gemm.kernel.group_gemm import group_gemm_same_mn # noqa: PLC0415 except ImportError: pytest.skip("group_gemm not available") @@ -220,7 +220,7 @@ def test_properties(self): if not torch.cuda.is_available(): pytest.skip("CUDA not available") try: - from xorl.ops.group_gemm.kernel.group_gemm import group_gemm_same_nk + from xorl.ops.group_gemm.kernel.group_gemm import group_gemm_same_nk # noqa: PLC0415 except ImportError: pytest.skip("group_gemm not available") diff --git a/tests/ops/test_moe_act.py b/tests/ops/test_moe_act.py index 1c0a3e58..ab4ba023 100644 --- a/tests/ops/test_moe_act.py +++ b/tests/ops/test_moe_act.py @@ -17,6 +17,9 @@ import torch import torch.nn as nn +from xorl.models.transformers.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig +from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import Qwen3MoeForCausalLM + DEVICE = "cuda" DTYPE = torch.bfloat16 @@ -29,7 +32,7 @@ def _available_moe_act_backends(): """Return backends that have moe_act local variants registered.""" - from xorl.models.layers.moe.backend import MOE_EXPERT_BACKENDS_MOE_ACT + from xorl.models.layers.moe.backend import MOE_EXPERT_BACKENDS_MOE_ACT # noqa: PLC0415 return list(MOE_EXPERT_BACKENDS_MOE_ACT.keys()) @@ -42,7 +45,7 @@ def _make_block_pair(ne, hd, inter, topk, backend, seed=42): Returns (std_block, act_block) where act_block.experts._moe_act = True. """ - from xorl.models.layers.moe.moe_block import MoEBlock + from xorl.models.layers.moe.moe_block import MoEBlock # noqa: PLC0415 torch.manual_seed(seed) std = MoEBlock(hd, ne, topk, inter, moe_implementation=backend) @@ -305,9 +308,7 @@ def bench_moe_act_tflops(seq_len): @pytest.mark.parametrize("backend", AVAILABLE_BACKENDS) def test_moe_act_via_gradient_checkpointing_enable(backend): """gradient_checkpointing_enable(moe_checkpoint_method='moe_act') sets _moe_act correctly.""" - from xorl.models.layers.moe.experts import MoEExperts - from xorl.models.transformers.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig - from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import Qwen3MoeForCausalLM + from xorl.models.layers.moe.experts import MoEExperts # noqa: PLC0415 config = Qwen3MoeConfig( vocab_size=1000, diff --git a/tests/ops/test_moe_gkn_format.py b/tests/ops/test_moe_gkn_format.py index ae678955..8e4e547d 100644 --- a/tests/ops/test_moe_gkn_format.py +++ b/tests/ops/test_moe_gkn_format.py @@ -138,7 +138,7 @@ class TestCheckpointLoadingGKN: @staticmethod def _get_buffer_class(): try: - from xorl.models.checkpoint_handlers.buffers import ExpertWeightBuffer + from xorl.models.checkpoint_handlers.buffers import ExpertWeightBuffer # noqa: PLC0415 return ExpertWeightBuffer except (ImportError, ModuleNotFoundError): @@ -210,7 +210,7 @@ def test_backends_match_reference_and_agree(self): """Eager, native, triton backends match reference; all MoEBlock backends agree.""" # --- Eager backend --- try: - from xorl.models.layers.moe.backend.eager import eager_expert_forward + from xorl.models.layers.moe.backend.eager import eager_expert_forward # noqa: PLC0415 except (ImportError, ModuleNotFoundError): pytest.skip("eager backend import requires transformers") @@ -230,7 +230,7 @@ def test_backends_match_reference_and_agree(self): torch.testing.assert_close(eager_out, ref_out, atol=1e-5, rtol=1e-5) # --- Native backend --- - from xorl.models.layers.moe.backend.native import native_expert_forward + from xorl.models.layers.moe.backend.native import native_expert_forward # noqa: PLC0415 device, dtype = "cuda", torch.bfloat16 gate_cuda = gate_gkn.to(device).to(dtype) @@ -254,11 +254,11 @@ def test_backends_match_reference_and_agree(self): # --- Triton backend --- try: - from xorl.utils.import_utils import is_fused_moe_available + from xorl.utils.import_utils import is_fused_moe_available # noqa: PLC0415 if not is_fused_moe_available(): raise ImportError - from xorl.ops.moe.triton import TritonMoeExpertsFunction + from xorl.ops.moe.triton import TritonMoeExpertsFunction # noqa: PLC0415 triton_out = TritonMoeExpertsFunction.apply( num_experts, @@ -291,7 +291,7 @@ def test_backends_match_reference_and_agree(self): # --- MoEBlock all backends agree --- try: - from xorl.models.layers.moe import MOE_EXPERT_BACKENDS, MoEBlock + from xorl.models.layers.moe import MOE_EXPERT_BACKENDS, MoEBlock # noqa: PLC0415 except (ImportError, ModuleNotFoundError): return # skip if not importable diff --git a/tests/ops/test_moe_ops.py b/tests/ops/test_moe_ops.py index 0ba1caf5..aa5986db 100644 --- a/tests/ops/test_moe_ops.py +++ b/tests/ops/test_moe_ops.py @@ -54,7 +54,7 @@ def test_histogram_and_index_compute(self): if not torch.cuda.is_available(): pytest.skip("CUDA not available") try: - from xorl.ops.group_gemm.kernel.moe import expert_histogram, moe_index_compute + from xorl.ops.group_gemm.kernel.moe import expert_histogram, moe_index_compute # noqa: PLC0415 except ImportError: pytest.skip("moe ops not available") @@ -111,7 +111,7 @@ def test_gather_scatter_add_gather(self): if not torch.cuda.is_available(): pytest.skip("CUDA not available") try: - from xorl.ops.group_gemm.kernel.moe import moe_add_gather, moe_gather, moe_scatter + from xorl.ops.group_gemm.kernel.moe import moe_add_gather, moe_gather, moe_scatter # noqa: PLC0415 except ImportError: pytest.skip("moe ops not available") @@ -190,7 +190,7 @@ def test_full_moe_pipeline(self): if not torch.cuda.is_available(): pytest.skip("CUDA not available") try: - from xorl.ops.group_gemm.kernel.moe import ( + from xorl.ops.group_gemm.kernel.moe import ( # noqa: PLC0415 expert_histogram, moe_gather, moe_index_compute, diff --git a/tests/ops/test_moe_torch_compile.py b/tests/ops/test_moe_torch_compile.py index a22b310e..df0f88e8 100644 --- a/tests/ops/test_moe_torch_compile.py +++ b/tests/ops/test_moe_torch_compile.py @@ -16,6 +16,13 @@ import torch import torch.nn as nn +from xorl.models.layers.rope import RotaryEmbedding +from xorl.models.transformers.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig +from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import ( + Qwen3MoeDecoderLayer, + Qwen3MoeForCausalLM, +) + DEVICE = "cuda" DTYPE = torch.bfloat16 @@ -23,7 +30,6 @@ def _tiny_moe_config(**overrides): """Create a minimal Qwen3MoeConfig for fast testing.""" - from xorl.models.transformers.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig defaults = dict( vocab_size=1000, @@ -49,7 +55,6 @@ def _tiny_moe_config(**overrides): def _make_position_embeddings(config, seq_len, device, dtype): """Create position_embeddings (cos, sin) for decoder layer tests.""" - from xorl.models.layers.rope import RotaryEmbedding rotary = RotaryEmbedding(config=config).to(device) dummy_hidden = torch.randn(1, seq_len, config.hidden_size, device=device, dtype=dtype) @@ -60,7 +65,7 @@ def _make_position_embeddings(config, seq_len, device, dtype): def _make_moe_block(moe_backend, hidden_size=128, num_experts=4, top_k=2, intermediate=128): """Create an MoEBlock with xavier init for numerical stability.""" - from xorl.models.layers.moe.moe_block import MoEBlock + from xorl.models.layers.moe.moe_block import MoEBlock # noqa: PLC0415 block = MoEBlock( hidden_size=hidden_size, @@ -80,7 +85,7 @@ def _available_backends(): """Return list of available MoE backends on this system.""" backends = ["native", "eager"] try: - from xorl.utils.import_utils import is_fused_moe_available + from xorl.utils.import_utils import is_fused_moe_available # noqa: PLC0415 if is_fused_moe_available(): backends.append("triton") @@ -156,7 +161,6 @@ def test_moe_block_compile(moe_backend): @pytest.mark.parametrize("moe_backend", AVAILABLE_BACKENDS) def test_decoder_layer_compile(moe_backend): """Decoder layer compile: aot_eager and inductor, forward + backward.""" - from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import Qwen3MoeDecoderLayer seq_len = 8 for compile_backend in ["aot_eager", "inductor"]: @@ -200,10 +204,6 @@ def test_decoder_layer_compile(moe_backend): ) def test_full_model_per_layer_compile(moe_backend, compile_backend): """Apply torch.compile to each decoder layer, run forward + backward.""" - from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import ( - Qwen3MoeDecoderLayer, - Qwen3MoeForCausalLM, - ) config = _tiny_moe_config(_moe_implementation=moe_backend) model = Qwen3MoeForCausalLM(config).to(DEVICE, DTYPE) @@ -520,7 +520,6 @@ def _measure_decoder_layer_peak_memory(layer, x, position_ids, position_embeddin @pytest.mark.parametrize("seq_len", [1024, 4096]) def test_decoder_layer_benchmark(seq_len): """Benchmark full Qwen3MoeDecoderLayer: TFLOPS + peak memory, compiled vs uncompiled.""" - from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import Qwen3MoeDecoderLayer hidden = 1024 intermediate = 512 diff --git a/tests/qlora/test_qlora.py b/tests/qlora/test_qlora.py index 7749a4d4..2808571e 100644 --- a/tests/qlora/test_qlora.py +++ b/tests/qlora/test_qlora.py @@ -82,7 +82,6 @@ def _quantize_injected_model(model): inject_qlora_into_model creates QLoRA modules with empty packed_weight_f32. This helper fills them with properly quantized random weights for testing. """ - from xorl.qlora.modules.linear import QLoRALinear for m in model.modules(): if isinstance(m, QLoRALinear) and m._is_prequantized: diff --git a/tests/server/api_server/test_checkpoint_paths.py b/tests/server/api_server/test_checkpoint_paths.py index 597ed2dc..96ca3a5a 100644 --- a/tests/server/api_server/test_checkpoint_paths.py +++ b/tests/server/api_server/test_checkpoint_paths.py @@ -12,6 +12,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException +from fastapi.exceptions import HTTPException from xorl.server.api_server.api_types import ( DeleteCheckpointRequest, @@ -86,7 +88,6 @@ def teardown_method(self): def test_save_load_all_formats_errors_and_auto_checkpoint(self): """Test save creates path, rejects duplicates, load works for all formats, 404s, isolation, and 000000 auto-save.""" - from fastapi.exceptions import HTTPException mock_output = MagicMock() @@ -210,7 +211,6 @@ def teardown_method(self): def test_list_delete_and_isolation(self): """Test list scoped by model, delete success/errors, and cross-model isolation.""" - from fastapi import HTTPException # --- List checkpoints scoped by model --- model_id = "user_list_test" @@ -276,7 +276,6 @@ class TestModelIdValidation: def test_valid_invalid_defaults_and_traversal(self): """Test valid IDs, invalid IDs, defaults, and path traversal blocking.""" - from fastapi import HTTPException # Valid for mid in ["default", "user_123", "user-123", "User123", "a", "A1_b2-C3"]: @@ -324,7 +323,6 @@ def teardown_method(self): def test_sampler_listing_deletion_and_adapter_tracking(self): """Test sampler in listings, shared across models, deletion, path resolution, and adapter tracking.""" - from fastapi import HTTPException # --- Sampler in listings --- os.makedirs(os.path.join(self.temp_dir, "weights", "user_test", "ckpt-001"), exist_ok=True) diff --git a/tests/server/orchestrator/test_cu_seqlens_alignment.py b/tests/server/orchestrator/test_cu_seqlens_alignment.py index c0928555..92769e4f 100644 --- a/tests/server/orchestrator/test_cu_seqlens_alignment.py +++ b/tests/server/orchestrator/test_cu_seqlens_alignment.py @@ -23,7 +23,10 @@ from xorl.data.collators.packing_concat_collator import ( PackingConcatCollator, ) -from xorl.server.orchestrator.packing import SequentialPacker +from xorl.data.collators.sequence_shard_collator import TextSequenceShardCollator +from xorl.server.backend import DummyBackend +from xorl.server.orchestrator.packing import SequentialPacker, unpack_per_token_outputs +from xorl.server.orchestrator.request_processor import RequestProcessor from xorl.utils.seqlen_pos_transform_utils import prepare_fa_kwargs_from_position_ids @@ -160,7 +163,6 @@ class TestOriginalPositionIdsAndPadding: def test_sequence_shard_collator_preservation_and_stale_overwrite(self): """Test _original_position_ids preservation, no-pad, and stale cu_seq_lens overwrite.""" - from xorl.data.collators.sequence_shard_collator import TextSequenceShardCollator # Preservation (SP size 2, 6 tokens) with ( @@ -269,8 +271,6 @@ def test_lcm_padding_and_collator_divisibility(self, mock_ps): assert len(batches[0]["input_ids"][0]) == expected # RequestProcessor computes lcm correctly - from xorl.server.backend import DummyBackend - from xorl.server.orchestrator.request_processor import RequestProcessor backend = DummyBackend() assert RequestProcessor(backend=backend, pad_to_multiple_of=128, cp_size=3).pad_to_multiple_of == 384 @@ -278,7 +278,6 @@ def test_lcm_padding_and_collator_divisibility(self, mock_ps): assert RequestProcessor(backend=backend, pad_to_multiple_of=128, cp_size=6).pad_to_multiple_of == 384 # After TextSequenceShardCollator, lengths divisible by cp_size - from xorl.data.collators.sequence_shard_collator import TextSequenceShardCollator cp_size = 3 mock_ps.return_value = Mock(cp_enabled=True, cp_size=cp_size, cp_rank=0, ringattn_size=1) @@ -316,7 +315,6 @@ class TestUnpackingWithPadding: @patch("xorl.server.orchestrator.packing.get_parallel_state") def test_padding_boundary_handling(self, mock_ps): """Padding creates extra boundary; no padding has correct count.""" - from xorl.server.orchestrator.packing import unpack_per_token_outputs mock_ps.return_value = Mock(cp_enabled=False, cp_size=1, cp_rank=0, ringattn_size=1) diff --git a/tests/server/orchestrator/test_orchestrator.py b/tests/server/orchestrator/test_orchestrator.py index 122d7f2f..f2f5ddb3 100644 --- a/tests/server/orchestrator/test_orchestrator.py +++ b/tests/server/orchestrator/test_orchestrator.py @@ -13,6 +13,7 @@ - Tests focus on Orchestrator's scheduling, routing, and output formatting """ +import msgpack import pytest @@ -107,7 +108,6 @@ def output_socket(addresses): def receive_outputs(output_socket, timeout_ms=1000, max_outputs=10): """Receive outputs from output socket.""" - import msgpack outputs = [] diff --git a/tests/server/orchestrator/test_orchestrator_client_communication.py b/tests/server/orchestrator/test_orchestrator_client_communication.py index 62691ae4..29b71413 100644 --- a/tests/server/orchestrator/test_orchestrator_client_communication.py +++ b/tests/server/orchestrator/test_orchestrator_client_communication.py @@ -9,6 +9,7 @@ import asyncio import logging +import socket as sock import pytest import pytest_asyncio @@ -155,8 +156,6 @@ async def _handle_request(self, request): @pytest.fixture def zmq_addresses(): - import socket as sock - def find_free_port(): with sock.socket(sock.AF_INET, sock.SOCK_STREAM) as s: s.bind(("127.0.0.1", 0)) @@ -315,7 +314,6 @@ async def test_edge_cases_and_lifecycle(zmq_addresses, mock_engine, orchestrator await client.get_output(timeout=1.0) # Client start/stop - use fresh ports to avoid address conflicts - import socket as sock def _find_free_port(): with sock.socket(sock.AF_INET, sock.SOCK_STREAM) as s: diff --git a/tests/server/weight_sync/test_pp_nccl_transfer.py b/tests/server/weight_sync/test_pp_nccl_transfer.py index 1af81cb3..b757592a 100644 --- a/tests/server/weight_sync/test_pp_nccl_transfer.py +++ b/tests/server/weight_sync/test_pp_nccl_transfer.py @@ -9,7 +9,7 @@ import torch -from xorl.server.weight_sync.handler import _prod +from xorl.server.weight_sync.handler import WeightSyncHandler, _prod class TestProd: @@ -37,7 +37,6 @@ def _make_handler(self, rank): handler = MagicMock() handler.rank = rank # Bind the real method - from xorl.server.weight_sync.handler import WeightSyncHandler handler._pp_nccl_transfer_buffer = WeightSyncHandler._pp_nccl_transfer_buffer.__get__(handler) return handler From 0f005bb4bd99bb522fb79943486a45052ceb976b Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Wed, 8 Apr 2026 12:09:58 -0700 Subject: [PATCH 19/41] Add hybrid_shared LoRA args to offline Trainer (#63) * Add moe_shared_lora and moe_hybrid_shared_lora to offline Trainer The MoE-aware LoRA injection (shared and hybrid_shared modes) was only available through the server's model_builder path. This adds the same support to the offline Trainer and direct_train so YAML configs can use: lora: moe_shared_lora: true # all experts share both lora_A and lora_B moe_hybrid_shared_lora: true # gate/up share lora_A, down shares lora_B Changes: - Add moe_shared_lora / moe_hybrid_shared_lora fields to LoRAArguments - Trainer._inject_lora: route to inject_lora_into_model_with_moe when flags are set - direct_train: same injection logic - Both save paths pass moe_hybrid_shared_lora to save_lora_checkpoint * Remove unsupported moe_shared_lora mode * Address PR #63 review: docs sync and load_lora_checkpoint docstring - Remove stale moe_shared_lora references from server.md and lora.mdx - Document that load_lora_checkpoint does not support hybrid-shared MoE LoRA checkpoints (key remapping + re-stacking not implemented) * Fix ruff-format: line wrap and blank line * Remove unused imports flagged by ruff --------- Co-authored-by: kiddyboots216 (cherry picked from commit c4829b24d574cc314670057a9699e822f6c80c63) --- .../content/docs/config-reference/server.md | 1 - docs/src/content/docs/moe/lora.mdx | 1 - src/xorl/arguments.py | 6 ++++ src/xorl/cli/direct_train.py | 31 ++++++++++++++---- src/xorl/lora/utils.py | 21 +++++++----- src/xorl/models/layers/moe/moe_block.py | 2 -- src/xorl/server/runner/model_runner.py | 1 - src/xorl/server/server_arguments.py | 5 --- src/xorl/trainers/model_builder.py | 14 +++----- src/xorl/trainers/trainer.py | 32 ++++++++++++++----- tests/models/test_moe_experts_lora.py | 19 +++++++++++ 11 files changed, 91 insertions(+), 42 deletions(-) diff --git a/docs/src/content/docs/config-reference/server.md b/docs/src/content/docs/config-reference/server.md index c4d41f7d..21eccae1 100644 --- a/docs/src/content/docs/config-reference/server.md +++ b/docs/src/content/docs/config-reference/server.md @@ -163,7 +163,6 @@ ZMQ communication between the launcher, workers, and API server. | `lora_rank` | `32` | LoRA rank (`r`). Default is 32 for server (vs 16 for local). | | `lora_alpha` | `16` | LoRA scaling factor. | | `lora_target_modules` | `null` | Module names to inject LoRA into. `null` = default for architecture. | -| `moe_shared_lora` | `false` | Share LoRA weights across all MoE experts. | | `moe_hybrid_shared_lora` | `false` | Share `lora_A` for gate/up projections and `lora_B` for down projections across experts. | | `enable_qlora` | `false` | Quantize base weights and train LoRA adapters on top. | | `quant_format` | `nvfp4` | Quantization format: `nvfp4`, `block_fp8`. | diff --git a/docs/src/content/docs/moe/lora.mdx b/docs/src/content/docs/moe/lora.mdx index b777dfbe..73541d56 100644 --- a/docs/src/content/docs/moe/lora.mdx +++ b/docs/src/content/docs/moe/lora.mdx @@ -46,7 +46,6 @@ This halves the LoRA parameter count while keeping per-expert expressiveness whe ```yaml # Server training config -moe_shared_lora: false # no sharing moe_hybrid_shared_lora: true # hybrid shared (recommended for large E) ``` diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index d3d44933..717d37f5 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -1037,6 +1037,12 @@ class LoRAArguments: default=False, metadata={"help": "Only save LoRA weights (not full model) in HF checkpoints"}, ) + moe_hybrid_shared_lora: bool = field( + default=False, + metadata={ + "help": "Enable hybrid shared LoRA for MoE: share lora_A for gate/up_proj, lora_B for down_proj across experts" + }, + ) # QLoRA: quantize base weights for memory savings enable_qlora: bool = field( default=False, diff --git a/src/xorl/cli/direct_train.py b/src/xorl/cli/direct_train.py index 292c391d..2ea3d810 100644 --- a/src/xorl/cli/direct_train.py +++ b/src/xorl/cli/direct_train.py @@ -8,7 +8,7 @@ from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict from xorl.distributed.pipeline_parallel import build_pipeline_schedule -from xorl.lora.utils import inject_lora_into_model, save_lora_checkpoint +from xorl.lora.utils import save_lora_checkpoint from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules from xorl.models.layers.moe.routing_replay import RoutingReplay, set_replay_stage from xorl.qlora import ( @@ -285,12 +285,28 @@ def main(): model._qlora_exclude_modules = exclude_modules helper.print_device_mem_info("VRAM usage after QLoRA injection") elif args.lora.enable_lora: - inject_lora_into_model( - model, - r=args.lora.lora_rank, - lora_alpha=args.lora.lora_alpha, - target_modules=args.lora.lora_target_modules, - ) + is_moe_model = getattr(model.config, "num_experts", 0) > 0 + + if is_moe_model and args.lora.moe_hybrid_shared_lora: + from xorl.lora.utils import inject_lora_into_model_with_moe + + logger.info_rank0(f"MoE-aware LoRA injection (hybrid_shared={args.lora.moe_hybrid_shared_lora})") + inject_lora_into_model_with_moe( + model, + r=args.lora.lora_rank, + lora_alpha=args.lora.lora_alpha, + target_modules=args.lora.lora_target_modules, + moe_hybrid_shared_lora=args.lora.moe_hybrid_shared_lora, + ) + else: + from xorl.lora.utils import inject_lora_into_model + + inject_lora_into_model( + model, + r=args.lora.lora_rank, + lora_alpha=args.lora.lora_alpha, + target_modules=args.lora.lora_target_modules, + ) helper.print_device_mem_info("VRAM usage after LoRA injection") maybe_upcast_trainable_adapter_params( @@ -811,6 +827,7 @@ def pp_loss_fn(pred, labels): target_modules=args.lora.lora_target_modules, r=args.lora.lora_rank, lora_alpha=args.lora.lora_alpha, + moe_hybrid_shared_lora=args.lora.moe_hybrid_shared_lora, ) logger.info_rank0(f"PEFT adapter checkpoint saved at {hf_weights_path} successfully!") elif hf_model_state_dict is not None: diff --git a/src/xorl/lora/utils.py b/src/xorl/lora/utils.py index 78f1df65..7270b7d7 100644 --- a/src/xorl/lora/utils.py +++ b/src/xorl/lora/utils.py @@ -569,7 +569,18 @@ def load_lora_checkpoint( """ Load LoRA weights from checkpoint. - Supports both xorl format and PEFT format checkpoints. + Supports both xorl format and PEFT format checkpoints for dense + (nn.Linear) and per-expert MoE LoRA weights. + + .. note:: + Hybrid-shared MoE LoRA checkpoints are **not** supported. + ``save_lora_checkpoint`` exports shared expert weights under + ``.mlp.experts.shared.{proj}.lora_{A|B}.weight``, but this + function cannot map those keys back to xorl's internal stacked + MoE parameter layout (``{proj}_lora_{A|B}`` with shape + ``[1, ...]``). Loading such a checkpoint would require + re-stacking the per-expert and shared tensors, which is not + yet implemented. Args: model: Model with LoRA layers already injected @@ -662,7 +673,6 @@ def inject_lora_into_moe_blocks( model: nn.Module, r: int = 16, lora_alpha: int = 16, - shared_lora: bool = False, target_modules: Optional[List[str]] = None, hybrid_shared: bool = False, ) -> int: @@ -682,7 +692,6 @@ def inject_lora_into_moe_blocks( model: Model containing MoE blocks r: LoRA rank lora_alpha: LoRA alpha for scaling - shared_lora: If True, share LoRA across all experts (more parameter efficient) target_modules: Which expert projections to apply LoRA to. Options: ["gate_proj", "up_proj", "down_proj"] Default: all three projections @@ -713,7 +722,6 @@ def inject_lora_into_moe_blocks( module.inject_lora( r=r, lora_alpha=lora_alpha, - shared_lora=shared_lora, target_modules=target_modules, hybrid_shared=hybrid_shared, ) @@ -723,7 +731,7 @@ def inject_lora_into_moe_blocks( if injected_count > 0: logger.info( f"Injected LoRA into {injected_count} MoE blocks with r={r}, " - f"alpha={lora_alpha}, shared={shared_lora}, hybrid_shared={hybrid_shared}" + f"alpha={lora_alpha}, hybrid_shared={hybrid_shared}" ) else: logger.warning( @@ -738,7 +746,6 @@ def inject_lora_into_model_with_moe( r: int = 16, lora_alpha: int = 16, target_modules: Optional[List[str]] = None, - moe_shared_lora: bool = False, moe_hybrid_shared_lora: bool = False, ) -> nn.Module: """ @@ -760,7 +767,6 @@ def inject_lora_into_model_with_moe( target_modules: List of module names to target. For attention: ["q_proj", "k_proj", "v_proj", "o_proj"] For MLP/experts: ["gate_proj", "up_proj", "down_proj"] - moe_shared_lora: If True, MoE experts share LoRA adapters (all shared) moe_hybrid_shared_lora: If True, use hybrid sharing (lora_A shared for gate/up, lora_B shared for down) Returns: @@ -802,7 +808,6 @@ def inject_lora_into_model_with_moe( model, r=r, lora_alpha=lora_alpha, - shared_lora=moe_shared_lora, target_modules=expert_modules, hybrid_shared=moe_hybrid_shared_lora, ) diff --git a/src/xorl/models/layers/moe/moe_block.py b/src/xorl/models/layers/moe/moe_block.py index 7255a45b..0e2ffb2a 100644 --- a/src/xorl/models/layers/moe/moe_block.py +++ b/src/xorl/models/layers/moe/moe_block.py @@ -86,7 +86,6 @@ def inject_lora( self, r: int = 16, lora_alpha: int = 16, - shared_lora: bool = False, target_modules: list = None, hybrid_shared: bool = False, ) -> None: @@ -99,7 +98,6 @@ def inject_lora( Args: r: LoRA rank. lora_alpha: LoRA alpha for scaling. - shared_lora: Unused (kept for API compat). target_modules: Which projections to apply LoRA to. Options: ``["gate_proj", "up_proj", "down_proj"]``. Default: all three. diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index bf9f6797..477d1a0b 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -488,7 +488,6 @@ def _initialize_model(self): lora_rank=self.lora_config.get("lora_rank", 32), lora_alpha=self.lora_config.get("lora_alpha", 16), lora_target_modules=target_modules, - moe_shared_lora=self.lora_config.get("moe_shared_lora", False), moe_hybrid_shared_lora=self.lora_config.get("moe_hybrid_shared_lora", False), enable_qlora=enable_qlora, quant_format=self.lora_config.get("quant_format", "nvfp4"), diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index 4862c913..173ca147 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -388,10 +388,6 @@ class ServerArguments: }, ) - moe_shared_lora: bool = field( - default=False, metadata={"help": "Enable shared LoRA for MoE: share LoRA weights across experts"} - ) - moe_hybrid_shared_lora: bool = field( default=False, metadata={ @@ -544,7 +540,6 @@ def to_config_dict(self) -> Dict[str, Any]: "lora_rank": self.lora_rank, "lora_alpha": self.lora_alpha, "lora_target_modules": self.lora_target_modules, - "moe_shared_lora": self.moe_shared_lora, "moe_hybrid_shared_lora": self.moe_hybrid_shared_lora, "enable_qlora": self.enable_qlora, "quant_format": self.quant_format, diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index 75ed57ea..39644534 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -13,7 +13,7 @@ from xorl.distributed.torch_parallelize import build_parallelize_model as _parallelize from xorl.lora import freeze_base_parameters -from xorl.lora.utils import inject_lora_into_model, inject_lora_into_model_with_moe +from xorl.lora.utils import inject_lora_into_model from xorl.models import build_foundation_model from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules from xorl.models.layers.rope import set_rope_native @@ -116,7 +116,6 @@ def build_training_model( lora_rank: int = 32, lora_alpha: int = 16, lora_target_modules: Optional[List[str]] = None, - moe_shared_lora: bool = False, moe_hybrid_shared_lora: bool = False, # --- QLoRA --- enable_qlora: bool = False, @@ -229,7 +228,6 @@ def build_training_model( lora_rank=lora_rank, lora_alpha=lora_alpha, lora_target_modules=lora_target_modules, - moe_shared_lora=moe_shared_lora, moe_hybrid_shared_lora=moe_hybrid_shared_lora, ) @@ -434,22 +432,20 @@ def _inject_lora( lora_rank: int, lora_alpha: int, lora_target_modules: Optional[List[str]], - moe_shared_lora: bool = False, moe_hybrid_shared_lora: bool = False, ) -> None: """Plain LoRA injection (dense + optional MoE-aware).""" is_moe_model = getattr(model.config, "num_experts", 0) > 0 - if is_moe_model and (moe_shared_lora or moe_hybrid_shared_lora): - logger.info_rank0( - f"MoE-aware LoRA injection (shared={moe_shared_lora}, hybrid_shared={moe_hybrid_shared_lora})" - ) + if is_moe_model and moe_hybrid_shared_lora: + from xorl.lora.utils import inject_lora_into_model_with_moe + + logger.info_rank0(f"MoE-aware LoRA injection (hybrid_shared={moe_hybrid_shared_lora})") inject_lora_into_model_with_moe( model, r=lora_rank, lora_alpha=lora_alpha, target_modules=lora_target_modules, - moe_shared_lora=moe_shared_lora, moe_hybrid_shared_lora=moe_hybrid_shared_lora, ) else: diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index b766a0fc..b19013c9 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -32,7 +32,7 @@ from xorl.distributed.pipeline_parallel import build_pipeline_schedule, build_pp_stage from xorl.distributed.sync_padding import synchronize_micro_batch_padding from xorl.distributed.torch_parallelize import build_parallelize_model -from xorl.lora.utils import inject_lora_into_model, save_lora_checkpoint +from xorl.lora.utils import save_lora_checkpoint from xorl.models import build_foundation_model, build_tokenizer, save_model_assets, save_model_weights from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules from xorl.models.layers.moe.aux_loss import global_load_balancing_loss_func @@ -475,15 +475,30 @@ def _inject_qlora(self) -> None: helper.print_device_mem_info("VRAM usage after QLoRA injection") def _inject_lora(self) -> None: - """Plain LoRA injection.""" + """Plain LoRA injection (dense + optional MoE-aware).""" args = self.args + is_moe_model = getattr(self.model.config, "num_experts", 0) > 0 - inject_lora_into_model( - self.model, - r=args.lora.lora_rank, - lora_alpha=args.lora.lora_alpha, - target_modules=args.lora.lora_target_modules, - ) + if is_moe_model and args.lora.moe_hybrid_shared_lora: + from xorl.lora.utils import inject_lora_into_model_with_moe + + logger.info_rank0(f"MoE-aware LoRA injection (hybrid_shared={args.lora.moe_hybrid_shared_lora})") + inject_lora_into_model_with_moe( + self.model, + r=args.lora.lora_rank, + lora_alpha=args.lora.lora_alpha, + target_modules=args.lora.lora_target_modules, + moe_hybrid_shared_lora=args.lora.moe_hybrid_shared_lora, + ) + else: + from xorl.lora.utils import inject_lora_into_model + + inject_lora_into_model( + self.model, + r=args.lora.lora_rank, + lora_alpha=args.lora.lora_alpha, + target_modules=args.lora.lora_target_modules, + ) helper.print_device_mem_info("VRAM usage after LoRA injection") def _parallelize(self) -> None: @@ -1144,6 +1159,7 @@ def _finalize(self, total_loss: float, grad_norm: float, lr: float) -> None: target_modules=args.lora.lora_target_modules, r=args.lora.lora_rank, lora_alpha=args.lora.lora_alpha, + moe_hybrid_shared_lora=args.lora.moe_hybrid_shared_lora, ) logger.info_rank0(f"PEFT adapter saved at {hf_weights_path}") elif hf_model_state_dict is not None: diff --git a/tests/models/test_moe_experts_lora.py b/tests/models/test_moe_experts_lora.py index f5a1af71..d99cbedd 100644 --- a/tests/models/test_moe_experts_lora.py +++ b/tests/models/test_moe_experts_lora.py @@ -221,6 +221,25 @@ def test_eager_forward_backward_and_moe_block(self): output.sum().backward() assert block.experts.gate_proj_lora_A.grad is not None + def test_hybrid_shared_shapes(self): + """Hybrid-shared injection keeps the supported shared tensor layout.""" + block = MoEBlock( + hidden_size=32, + num_experts=4, + top_k=2, + intermediate_size=64, + moe_implementation="eager", + ) + + block.inject_lora(r=4, lora_alpha=8, hybrid_shared=True) + + assert block.experts.gate_proj_lora_A.shape == (1, 32, 4) + assert block.experts.gate_proj_lora_B.shape == (4, 4, 64) + assert block.experts.up_proj_lora_A.shape == (1, 32, 4) + assert block.experts.up_proj_lora_B.shape == (4, 4, 64) + assert block.experts.down_proj_lora_A.shape == (4, 64, 4) + assert block.experts.down_proj_lora_B.shape == (1, 4, 32) + # --------------------------------------------------------------------------- # 4. GPU LoRA forward/backward (triton + native) From d669b47b395736f69b9f96d3643f3f94298f53ad Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 9 Apr 2026 10:04:44 -0700 Subject: [PATCH 20/41] Fix checkpoint utils import in arguments (#96) (cherry picked from commit df96b96182d022fe0ccbcf2bba39ef85c0596bb6) --- src/xorl/arguments.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 717d37f5..4afb0306 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -25,8 +25,8 @@ import torch import yaml -from .checkpoint_utils import get_checkpoint_path from .utils import logging +from .utils.checkpoint_utils import get_checkpoint_path T = TypeVar("T") From 5e2c9c1b5649b409c2ac05b71731190dc70c1d8f Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 9 Apr 2026 19:22:03 -0700 Subject: [PATCH 21/41] Make example configs path-generic (#99) (cherry picked from commit 51a2879488913b87dce85b15009e35f2af590b13) --- tests/test_example_assets.py | 50 ++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) create mode 100644 tests/test_example_assets.py diff --git a/tests/test_example_assets.py b/tests/test_example_assets.py new file mode 100644 index 00000000..a7137dd6 --- /dev/null +++ b/tests/test_example_assets.py @@ -0,0 +1,50 @@ +import re +import subprocess +from pathlib import Path + +import pytest + + +pytestmark = [pytest.mark.cpu] + + +TRACKED_TEXT_SUFFIXES = {".md", ".mdx", ".sh", ".yaml", ".yml"} +FORBIDDEN_PATTERNS = { + "home_dir": re.compile(r"(? list[Path]: + try: + result = subprocess.run( + ["git", "-C", str(repo_root), "ls-files", "examples", "experiments"], + check=True, + capture_output=True, + text=True, + ) + except (FileNotFoundError, subprocess.CalledProcessError) as exc: + pytest.skip(f"git ls-files unavailable: {exc}") + + tracked_files = [] + for relative_path in result.stdout.splitlines(): + path = repo_root / relative_path + if path.suffix in TRACKED_TEXT_SUFFIXES: + tracked_files.append(path) + return tracked_files + + +def test_examples_and_experiments_avoid_personal_absolute_paths(): + repo_root = Path(__file__).resolve().parents[1] + violations: list[str] = [] + + for path in _tracked_example_and_experiment_files(repo_root): + text = path.read_text(encoding="utf-8") + for pattern_name, pattern in FORBIDDEN_PATTERNS.items(): + if match := pattern.search(text): + line_number = text.count("\n", 0, match.start()) + 1 + violations.append(f"{path.relative_to(repo_root)}:{line_number}: {pattern_name}: {match.group(0)}") + + assert not violations, "Tracked example and experiment assets contain personal absolute paths:\n" + "\n".join( + violations + ) From 0082dc3971bdb030618ac7fd04d7ac207c07cc08 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 9 Apr 2026 21:52:54 -0700 Subject: [PATCH 22/41] Fix package import cycles for LoRA and runner modules (#101) * Fix package import cycles for LoRA and runner modules * Fix missing Optional import in weight sync backend (cherry picked from commit 1b654738a6cf229aabe28e2b40523d91465f3f12) --- src/xorl/lora/__init__.py | 132 +++++++++++-------- src/xorl/models/layers/moe/__init__.py | 40 ++++-- src/xorl/models/layers/moe/moe_block.py | 2 +- src/xorl/qlora/__init__.py | 96 +++++++++++--- src/xorl/server/runner/__init__.py | 18 ++- src/xorl/server/runner/adapters/__init__.py | 18 ++- src/xorl/server/weight_sync/backends/base.py | 2 +- 7 files changed, 222 insertions(+), 86 deletions(-) diff --git a/src/xorl/lora/__init__.py b/src/xorl/lora/__init__.py index 74d1e2ff..1a698376 100644 --- a/src/xorl/lora/__init__.py +++ b/src/xorl/lora/__init__.py @@ -1,42 +1,11 @@ """ Xorl LoRA (Low-Rank Adaptation) module. -This module provides a simple, FSDP-compatible LoRA implementation -with flat parameter structure for easy checkpoint management. - -Supports both dense layers (nn.Linear) and MoE expert blocks -with group GEMM kernels for efficient forward/backward passes. +This package uses lazy exports to avoid circular imports between: +`xorl.models.layers.moe.*`, `xorl.lora.mapping`, and `xorl.lora.utils`. """ -# Base class and implementations -# Mapping -from xorl.lora.mapping import ( - LORA_MAPPING, - can_apply_lora, - get_lora_class_for_module, - register_lora_mapping, -) from xorl.lora.modules import LoraLinear, LoraModule - -# Utility functions -from xorl.lora.utils import ( - count_lora_parameters, - freeze_base_parameters, - get_all_lora_state_dict, - get_lora_parameters, - get_lora_state_dict, - get_moe_lora_state_dict, - inject_lora_into_model, - inject_lora_into_model_with_moe, - inject_lora_into_moe_blocks, - load_lora_checkpoint, - load_lora_state_dict, - save_lora_checkpoint, -) - -# MoE LoRA — deferred to avoid circular import with xorl.models.layers.moe. -# Access via xorl.lora.MoEExpertsLoRA etc. triggers __getattr__ below. -# Group GEMM ops for direct access from xorl.ops.group_gemm.kernel import ( compute_lora_scaling, get_lora_delta_weight_stacked, @@ -46,39 +15,61 @@ ) +_MAPPING_ATTRS = { + "LORA_MAPPING", + "can_apply_lora", + "get_lora_class_for_module", + "register_lora_mapping", +} + +_UTIL_ATTRS = { + "count_lora_parameters", + "freeze_base_parameters", + "get_all_lora_state_dict", + "get_lora_parameters", + "get_lora_state_dict", + "get_moe_lora_state_dict", + "inject_lora_into_model", + "inject_lora_into_model_with_moe", + "inject_lora_into_moe_blocks", + "load_lora_checkpoint", + "load_lora_state_dict", + "save_lora_checkpoint", +} + +_MOE_LAZY_ATTRS = { + "MoELoRAConfig", + "MoEExpertsLoRA", + "copy_weights_to_lora_experts", + "mark_only_lora_as_trainable", + "lora_state_dict", +} + + __all__ = [ - # Base class "LoraModule", - # Dense LoRA "LoraLinear", - # Mapping "LORA_MAPPING", "can_apply_lora", "get_lora_class_for_module", "register_lora_mapping", - # Injection utilities "inject_lora_into_model", "inject_lora_into_model_with_moe", "inject_lora_into_moe_blocks", - # State dict utilities "get_lora_state_dict", "get_moe_lora_state_dict", "get_all_lora_state_dict", "load_lora_state_dict", - # Parameter utilities "freeze_base_parameters", "get_lora_parameters", "count_lora_parameters", - # Checkpoint utilities "save_lora_checkpoint", "load_lora_checkpoint", - # MoE LoRA (lazy) "MoELoRAConfig", "MoEExpertsLoRA", "copy_weights_to_lora_experts", "mark_only_lora_as_trainable", "lora_state_dict", - # Group GEMM ops "init_lora_weights_stacked", "compute_lora_scaling", "merge_lora_weights_stacked", @@ -87,18 +78,53 @@ ] -# Lazy imports for MoE LoRA to break circular import: -# moe/lora.py → xorl.lora.modules.base → xorl.lora → xorl.models.layers.moe (cycle) -_MOE_LAZY_ATTRS = { - "MoELoRAConfig", - "MoEExpertsLoRA", - "copy_weights_to_lora_experts", - "mark_only_lora_as_trainable", - "lora_state_dict", -} +def __getattr__(name): + if name in _MAPPING_ATTRS: + from xorl.lora.mapping import ( # noqa: PLC0415 + LORA_MAPPING, + can_apply_lora, + get_lora_class_for_module, + register_lora_mapping, + ) + g = globals() + g["LORA_MAPPING"] = LORA_MAPPING + g["can_apply_lora"] = can_apply_lora + g["get_lora_class_for_module"] = get_lora_class_for_module + g["register_lora_mapping"] = register_lora_mapping + return g[name] + + if name in _UTIL_ATTRS: + from xorl.lora.utils import ( # noqa: PLC0415 + count_lora_parameters, + freeze_base_parameters, + get_all_lora_state_dict, + get_lora_parameters, + get_lora_state_dict, + get_moe_lora_state_dict, + inject_lora_into_model, + inject_lora_into_model_with_moe, + inject_lora_into_moe_blocks, + load_lora_checkpoint, + load_lora_state_dict, + save_lora_checkpoint, + ) + + g = globals() + g["count_lora_parameters"] = count_lora_parameters + g["freeze_base_parameters"] = freeze_base_parameters + g["get_all_lora_state_dict"] = get_all_lora_state_dict + g["get_lora_parameters"] = get_lora_parameters + g["get_lora_state_dict"] = get_lora_state_dict + g["get_moe_lora_state_dict"] = get_moe_lora_state_dict + g["inject_lora_into_model"] = inject_lora_into_model + g["inject_lora_into_model_with_moe"] = inject_lora_into_model_with_moe + g["inject_lora_into_moe_blocks"] = inject_lora_into_moe_blocks + g["load_lora_checkpoint"] = load_lora_checkpoint + g["load_lora_state_dict"] = load_lora_state_dict + g["save_lora_checkpoint"] = save_lora_checkpoint + return g[name] -def __getattr__(name): if name in _MOE_LAZY_ATTRS: from xorl.models.layers.moe.lora import ( # noqa: PLC0415 MoEExpertsLoRA, @@ -108,7 +134,6 @@ def __getattr__(name): mark_only_lora_as_trainable, ) - # Cache all at once to avoid repeated imports g = globals() g["MoELoRAConfig"] = MoELoRAConfig g["MoEExpertsLoRA"] = MoEExpertsLoRA @@ -116,4 +141,5 @@ def __getattr__(name): g["mark_only_lora_as_trainable"] = mark_only_lora_as_trainable g["lora_state_dict"] = lora_state_dict return g[name] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/xorl/models/layers/moe/__init__.py b/src/xorl/models/layers/moe/__init__.py index e9f600dc..715812d8 100644 --- a/src/xorl/models/layers/moe/__init__.py +++ b/src/xorl/models/layers/moe/__init__.py @@ -13,14 +13,6 @@ from .aux_loss import global_load_balancing_loss_func from .backend import MOE_EXPERT_BACKENDS from .experts import MoEExperts -from .lora import ( - MoEExpertsLoRA, - MoELoRAConfig, - copy_weights_to_lora_experts, - inject_lora_into_experts, - lora_state_dict, - mark_only_lora_as_trainable, -) from .moe_block import MoEBlock from .router import TopKRouter @@ -38,3 +30,35 @@ "MOE_EXPERT_BACKENDS", "global_load_balancing_loss_func", ] + + +_LAZY_ATTRS = { + "MoEExpertsLoRA", + "MoELoRAConfig", + "copy_weights_to_lora_experts", + "inject_lora_into_experts", + "lora_state_dict", + "mark_only_lora_as_trainable", +} + + +def __getattr__(name): + if name in _LAZY_ATTRS: + from .lora import ( # noqa: PLC0415 + MoEExpertsLoRA, + MoELoRAConfig, + copy_weights_to_lora_experts, + inject_lora_into_experts, + lora_state_dict, + mark_only_lora_as_trainable, + ) + + g = globals() + g["MoEExpertsLoRA"] = MoEExpertsLoRA + g["MoELoRAConfig"] = MoELoRAConfig + g["copy_weights_to_lora_experts"] = copy_weights_to_lora_experts + g["inject_lora_into_experts"] = inject_lora_into_experts + g["lora_state_dict"] = lora_state_dict + g["mark_only_lora_as_trainable"] = mark_only_lora_as_trainable + return g[name] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/xorl/models/layers/moe/moe_block.py b/src/xorl/models/layers/moe/moe_block.py index 0e2ffb2a..13abe2af 100644 --- a/src/xorl/models/layers/moe/moe_block.py +++ b/src/xorl/models/layers/moe/moe_block.py @@ -7,7 +7,6 @@ import torch.nn.functional as F from .experts import MoEExperts -from .lora import inject_lora_into_experts from .router import TopKRouter from .routing_replay import RoutingReplay, get_replay_stage @@ -104,6 +103,7 @@ def inject_lora( hybrid_shared: If True, share ``lora_A`` for gate/up and ``lora_B`` for down across experts. """ + from .lora import inject_lora_into_experts # noqa: PLC0415 inject_lora_into_experts( self, diff --git a/src/xorl/qlora/__init__.py b/src/xorl/qlora/__init__.py index cf22245a..0d273463 100644 --- a/src/xorl/qlora/__init__.py +++ b/src/xorl/qlora/__init__.py @@ -1,20 +1,25 @@ -from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules - -from .modules.block_fp8_linear import BlockFP8QLoRALinear -from .modules.linear import QLoRALinear, prefetch_aqn_noise -from .modules.moe_experts import QLoRAMoeExperts -from .modules.nf4_linear import NF4QLoRALinear -from .modules.nvfp4_linear import NvFP4QLoRALinear -from .utils import ( - detect_prequantized_block_fp8, - detect_prequantized_nvfp4, - inject_qlora_into_model, - maybe_load_and_quantize_moe_qlora, - maybe_load_prequantized_qlora, - maybe_quantize_qlora, - maybe_requant_qlora, - save_qlora_checkpoint, -) +"""Lazy exports for QLoRA to avoid package-level circular imports.""" + +_MODULE_ATTRS = { + "BlockFP8QLoRALinear", + "NF4QLoRALinear", + "NvFP4QLoRALinear", + "QLoRALinear", + "QLoRAMoeExperts", + "prefetch_aqn_noise", +} + +_UTIL_ATTRS = { + "detect_prequantized_block_fp8", + "detect_prequantized_nvfp4", + "get_prequantized_exclude_modules", + "inject_qlora_into_model", + "maybe_load_and_quantize_moe_qlora", + "maybe_load_prequantized_qlora", + "maybe_quantize_qlora", + "maybe_requant_qlora", + "save_qlora_checkpoint", +} __all__ = [ @@ -34,3 +39,60 @@ "maybe_load_prequantized_qlora", "get_prequantized_exclude_modules", ] + + +def __getattr__(name): + if name in _MODULE_ATTRS: + from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear # noqa: PLC0415 + from xorl.qlora.modules.linear import QLoRALinear, prefetch_aqn_noise # noqa: PLC0415 + from xorl.qlora.modules.moe_experts import QLoRAMoeExperts # noqa: PLC0415 + from xorl.qlora.modules.nf4_linear import NF4QLoRALinear # noqa: PLC0415 + from xorl.qlora.modules.nvfp4_linear import NvFP4QLoRALinear # noqa: PLC0415 + + globals().update( + { + "QLoRALinear": QLoRALinear, + "NvFP4QLoRALinear": NvFP4QLoRALinear, + "BlockFP8QLoRALinear": BlockFP8QLoRALinear, + "NF4QLoRALinear": NF4QLoRALinear, + "QLoRAMoeExperts": QLoRAMoeExperts, + "prefetch_aqn_noise": prefetch_aqn_noise, + } + ) + return globals()[name] + + if name in _UTIL_ATTRS: + from xorl.qlora.utils import ( # noqa: PLC0415 + detect_prequantized_block_fp8, + detect_prequantized_nvfp4, + inject_qlora_into_model, + maybe_load_and_quantize_moe_qlora, + maybe_load_prequantized_qlora, + maybe_quantize_qlora, + maybe_requant_qlora, + save_qlora_checkpoint, + ) + + globals().update( + { + "detect_prequantized_block_fp8": detect_prequantized_block_fp8, + "detect_prequantized_nvfp4": detect_prequantized_nvfp4, + "inject_qlora_into_model": inject_qlora_into_model, + "maybe_load_and_quantize_moe_qlora": maybe_load_and_quantize_moe_qlora, + "maybe_load_prequantized_qlora": maybe_load_prequantized_qlora, + "maybe_quantize_qlora": maybe_quantize_qlora, + "maybe_requant_qlora": maybe_requant_qlora, + "save_qlora_checkpoint": save_qlora_checkpoint, + } + ) + + if "get_prequantized_exclude_modules" not in globals(): + from xorl.models.checkpoint_handlers.buffers import ( # noqa: PLC0415 + get_prequantized_exclude_modules, + ) + + globals()["get_prequantized_exclude_modules"] = get_prequantized_exclude_modules + + return globals()[name] + + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/xorl/server/runner/__init__.py b/src/xorl/server/runner/__init__.py index 563195c8..bb13a89f 100644 --- a/src/xorl/server/runner/__init__.py +++ b/src/xorl/server/runner/__init__.py @@ -1,5 +1,17 @@ -from xorl.server.runner.model_runner import ModelRunner -from xorl.server.runner.runner_dispatcher import RunnerDispatcher - +"""Lazy exports for server runner modules.""" __all__ = ["RunnerDispatcher", "ModelRunner"] + + +def __getattr__(name): + if name == "ModelRunner": + from xorl.server.runner.model_runner import ModelRunner # noqa: PLC0415 + + globals()["ModelRunner"] = ModelRunner + return ModelRunner + if name == "RunnerDispatcher": + from xorl.server.runner.runner_dispatcher import RunnerDispatcher # noqa: PLC0415 + + globals()["RunnerDispatcher"] = RunnerDispatcher + return RunnerDispatcher + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/xorl/server/runner/adapters/__init__.py b/src/xorl/server/runner/adapters/__init__.py index 924cbc10..c2b14a56 100644 --- a/src/xorl/server/runner/adapters/__init__.py +++ b/src/xorl/server/runner/adapters/__init__.py @@ -1,5 +1,17 @@ -from xorl.server.runner.adapters.adapter_coordinator import AdapterCoordinator -from xorl.server.runner.adapters.manager import LoRAAdapterManager - +"""Lazy exports for runner adapter modules.""" __all__ = ["LoRAAdapterManager", "AdapterCoordinator"] + + +def __getattr__(name): + if name == "LoRAAdapterManager": + from xorl.server.runner.adapters.manager import LoRAAdapterManager # noqa: PLC0415 + + globals()["LoRAAdapterManager"] = LoRAAdapterManager + return LoRAAdapterManager + if name == "AdapterCoordinator": + from xorl.server.runner.adapters.adapter_coordinator import AdapterCoordinator # noqa: PLC0415 + + globals()["AdapterCoordinator"] = AdapterCoordinator + return AdapterCoordinator + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/xorl/server/weight_sync/backends/base.py b/src/xorl/server/weight_sync/backends/base.py index f348f600..e6f65a53 100644 --- a/src/xorl/server/weight_sync/backends/base.py +++ b/src/xorl/server/weight_sync/backends/base.py @@ -40,7 +40,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import Any, Dict, FrozenSet, List, Tuple +from typing import Any, Dict, FrozenSet, List, Optional, Tuple import torch From 80b03fced1098117f5c179a44b2803f9df90f41e Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 10 Apr 2026 00:24:11 -0700 Subject: [PATCH 23/41] Improve documentation: installation, server training, RL guide, and README (#57) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * docs: add DeepEP multi-node prerequisites and cluster setup - Document nvidia_peermem requirement for GPUDirect RDMA — without it NVSHMEM cannot register GPU buffers with IB HCAs and DeepEP crashes with SIGABRT at the first dispatch - Document IBGDA driver settings (NVreg_EnableStreamMemOPs, PeerMappingOverride) with initramfs rebuild and reboot steps - Add troubleshooting entries for SIGABRT/num_recv_tokens=-1 and IBGDA init failures - Add 2-node EP16 dummy benchmark configs (deepep and alltoall) - Add slurm_train_30b_2node.sh launch script with corrected NVSHMEM lib path detection (use __path__ instead of __file__ for namespace packages) - Add enable_ibgda_reboot.sh for per-node IBGDA kernel module setup * docs: remove DeepEP install section from README * docs: consolidate DeepEP install and multi-node prerequisites under single section * docs: point DeepEP install to upstream repo instead of internal wheel * docs: remove multi-node setup section from installation page * docs: add link to installation guide in quickstart * docs: simplify README, link to docs for features and details * revert: remove examples and scripts from PR * docs: use standard GitHub Pages URL (org/repo format) * docs: update README — XoRL logo, remove tagline * docs: add overview section and improve README layout * docs: add emojis to README section headings * docs: fix clone URL to togethercomputer/xorl-internal * docs: rename DEVELOPMENT.md to CONTRIBUTING.md, add contributing section to README * docs: add supported models section to README * docs: add emojis to header nav links * add xorl-client and xorl-sglang as git submodules * docs: improve installation guide, add xorl-client and xorl-sglang documentation - Add xorl stack overview table (xorl, xorl-client, xorl-sglang) to README - Add conda and uv install options for both base and sglang environments - Add pyproject.sglang.toml for unified install with PyTorch 2.9.1 - Add xorl-client as a git dependency in pyproject.toml - Expand xorl-client docs: client classes, key features, quick example - Expand xorl-sglang docs: detailed diff from upstream SGLang covering weight sync endpoints, MoE routing export, numerical alignment, batch-invariant mode, and bug fixes - Update submodule install instructions to reference submodules/ paths * docs: fix repo table column spacing in README * docs: add server training examples (no_robot_sft, password_memorization) * docs: add full stack launch guide for training + inference servers * docs: add RL loop illustration and rename to "Server Training for RL" * docs: replace ASCII diagram with HTML illustration for RL architecture * docs: reorganize server training section into sub-sections - Add RL architecture SVG diagram - Reorganize sidebar: Training Server, Client SDK, Examples as sub-sections - Move architecture.mdx -> training-server/launching.mdx - Move api-reference.md -> training-server/api-reference.md - Split client-sdk into overview, training-loop, loss-functions pages - Split examples into sft-no-robots, password-memorization pages - Merge rl-training.mdx content into client-sdk pages - Shorten overview to concise entry point with links to sub-pages * docs: add recommended reading section with RL and system design papers * Fix doc hallucinations: nonexistent config, invalid loss_fn, wrong API pattern, port inconsistency * Fix server training docs examples * Fix remaining server training docs issues --------- Co-authored-by: zzz0906 Co-authored-by: Ashwinee Panda (cherry picked from commit 6e0d8d440a3b5cb975dc8a732d855e68c7e39b39) --- .gitmodules | 6 + README.md | 66 ++- docs/astro.config.mjs | 30 +- docs/src/assets/rl-architecture.svg | 119 ++++ .../docs/getting-started/installation.md | 57 +- .../docs/getting-started/quickstart.md | 48 +- .../client-sdk/loss-functions.md | 149 +++++ .../server-training/client-sdk/overview.md | 190 +++++++ .../client-sdk/training-loop.md | 197 +++++++ .../examples/password-memorization.md | 73 +++ .../server-training/examples/sft-no-robots.md | 45 ++ .../content/docs/server-training/overview.md | 538 ++++-------------- .../docs/server-training/rl-training.mdx | 329 ----------- .../content/docs/server-training/sglang.mdx | 168 ++++-- .../{ => training-server}/api-reference.md | 8 +- .../launching.mdx} | 2 +- .../server-training/weight-sync/overview.mdx | 20 +- examples/server/no_robot_sft/README.md | 42 +- pyproject.sglang.toml | 123 ++++ pyproject.toml | 1 + submodules/xorl-client | 1 + submodules/xorl-sglang | 1 + 22 files changed, 1356 insertions(+), 857 deletions(-) create mode 100644 .gitmodules create mode 100644 docs/src/assets/rl-architecture.svg create mode 100644 docs/src/content/docs/server-training/client-sdk/loss-functions.md create mode 100644 docs/src/content/docs/server-training/client-sdk/overview.md create mode 100644 docs/src/content/docs/server-training/client-sdk/training-loop.md create mode 100644 docs/src/content/docs/server-training/examples/password-memorization.md create mode 100644 docs/src/content/docs/server-training/examples/sft-no-robots.md delete mode 100644 docs/src/content/docs/server-training/rl-training.mdx rename docs/src/content/docs/server-training/{ => training-server}/api-reference.md (86%) rename docs/src/content/docs/server-training/{architecture.mdx => training-server/launching.mdx} (99%) create mode 100644 pyproject.sglang.toml create mode 160000 submodules/xorl-client create mode 160000 submodules/xorl-sglang diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 00000000..8695bdcd --- /dev/null +++ b/.gitmodules @@ -0,0 +1,6 @@ +[submodule "submodules/xorl-client"] + path = submodules/xorl-client + url = https://github.com/togethercomputer/xorl-client +[submodule "submodules/xorl-sglang"] + path = submodules/xorl-sglang + url = https://github.com/togethercomputer/xorl-sglang diff --git a/README.md b/README.md index b663c727..6103e43b 100644 --- a/README.md +++ b/README.md @@ -18,10 +18,18 @@ XoRL is a distributed training framework designed for large language models with composable parallelism and flexible training modes. +**The XoRL stack consists of three repos:** + +| Repo                        | Description | +|---|---| +| **[xorl](https://github.com/togethercomputer/xorl-internal)** | Distributed training framework — local SFT/pretraining and server-mode RL training | +| **[xorl-client](https://github.com/togethercomputer/xorl-client)** | Lightweight Python SDK for driving the xorl training server (forward/backward, optimizer steps, checkpointing, sampling) | +| **[xorl-sglang](https://github.com/togethercomputer/xorl-sglang)** | Fork of [SGLang](https://github.com/sgl-project/sglang) with weight-sync APIs, MoE routing export, and numerical alignment for online RL | + **Two training modes:** - **Local** — `torchrun`-based training for offline SFT and pretraining -- **Server** — REST API-driven training for online RL loops where an external orchestrator (e.g. [xorl_client](https://github.com/xorl-org/xorl_client)) controls the training loop +- **Server** — REST API-driven training for online RL loops where [xorl-client](https://github.com/togethercomputer/xorl-client) drives the training loop and [xorl-sglang](https://github.com/togethercomputer/xorl-sglang) serves inference **Parallelism strategies** — mix and match freely: @@ -40,14 +48,61 @@ XoRL is a distributed training framework designed for large language models with ## 🚀 Installation ```bash -git clone --recurse-submodules git@github.com:togethercomputer/xorl.git -cd xorl -uv sync +git clone --recurse-submodules git@github.com:togethercomputer/xorl-internal.git +cd xorl-internal ``` > Already cloned without `--recurse-submodules`? Run `git submodule update --init --recursive` -See the [installation guide](https://togethercomputer.github.io/xorl/getting-started/installation/) for full setup including optional dependencies (DeepEP, Flash Attention). +### Option A: uv (recommended) + +```bash +uv sync +source .venv/bin/activate +``` + +### Option B: conda + +```bash +conda create -n xorl python=3.12 +conda activate xorl +pip install -e . +``` + +### Submodules + +The repo includes two git submodules under `submodules/` (needed for server / online RL training): + +- **[xorl-client](https://github.com/togethercomputer/xorl-client)** — Lightweight Python SDK (no PyTorch dependency) for driving the xorl training server. Provides `ServiceClient`, `TrainingClient`, `SamplingClient`, and `RestClient` with async-first `APIFuture` semantics, automatic request ordering, and Tinker API compatibility. +- **[xorl-sglang](https://github.com/togethercomputer/xorl-sglang)** — XoRL's fork of [SGLang](https://github.com/sgl-project/sglang) with NCCL-based weight sync endpoints, MoE routing data export (R3), and numerical alignment flags for online RL. + +Install individually: + +```bash +pip install -e submodules/xorl-client +pip install -e "submodules/xorl-sglang/python[all]" +``` + +Or use the bundled `pyproject.sglang.toml` which pins PyTorch to 2.9.1 (required by sglang) and installs everything together: + +**uv:** +```bash +cp pyproject.sglang.toml pyproject.toml +uv sync +source .venv/bin/activate +``` + +**conda:** +```bash +conda create -n xorl-sglang python=3.12 +conda activate xorl-sglang +cp pyproject.sglang.toml pyproject.toml +pip install -e . +``` + +> **Note:** The default `pyproject.toml` uses PyTorch 2.10.0. sglang requires PyTorch 2.9.1, so the two cannot coexist in the same environment unless you use `pyproject.sglang.toml`. + +See the [installation guide](https://togethercomputer.github.io/xorl-internal/getting-started/installation/) for full setup including optional dependencies (DeepEP, Flash Attention). ## ⚡ Quick Start @@ -87,4 +142,5 @@ Models are loaded directly from HuggingFace checkpoints — no preprocessing nee ## 🤝 Contributing + See [CONTRIBUTING.md](CONTRIBUTING.md) for development setup, coding conventions, and how to run tests. diff --git a/docs/astro.config.mjs b/docs/astro.config.mjs index 0bdc59a5..f0f3c21a 100644 --- a/docs/astro.config.mjs +++ b/docs/astro.config.mjs @@ -37,14 +37,28 @@ export default defineConfig({ ], }, { - label: "Server Training", + label: "Server Training for RL", collapsed: true, items: [ { label: "Overview", slug: "server-training/overview" }, - { label: "Server Architecture", slug: "server-training/architecture" }, - { label: "API Reference", slug: "server-training/api-reference" }, - { label: "RL Training", slug: "server-training/rl-training" }, + { + label: "Training Server (xorl)", + collapsed: true, + items: [ + { label: "Launching & Configuration", slug: "server-training/training-server/launching" }, + { label: "API Reference", slug: "server-training/training-server/api-reference" }, + ], + }, { label: "Inference: xorl-sglang", slug: "server-training/sglang" }, + { + label: "Client SDK (xorl-client)", + collapsed: true, + items: [ + { label: "Overview", slug: "server-training/client-sdk/overview" }, + { label: "Training Loop Patterns", slug: "server-training/client-sdk/training-loop" }, + { label: "Loss Functions", slug: "server-training/client-sdk/loss-functions" }, + ], + }, { label: "Weight Sync", collapsed: true, @@ -53,6 +67,14 @@ export default defineConfig({ { label: "Backend: nccl_broadcast", slug: "server-training/weight-sync/nccl-broadcast" }, ], }, + { + label: "Examples", + collapsed: true, + items: [ + { label: "SFT on No Robots", slug: "server-training/examples/sft-no-robots" }, + { label: "Password Memorization", slug: "server-training/examples/password-memorization" }, + ], + }, ], }, { diff --git a/docs/src/assets/rl-architecture.svg b/docs/src/assets/rl-architecture.svg new file mode 100644 index 00000000..11038f2a --- /dev/null +++ b/docs/src/assets/rl-architecture.svg @@ -0,0 +1,119 @@ + + + + + + + + + + THE RL LOOP + + + + Your RL Orchestrator + Reward model · Code sandbox · Math verifier · Rule-based scorer + + YOU BUILD THIS + + + + + + + 1 + scored rollouts + + + 2 + prompts + + + + xorl_client + .forward_backward() + .optim_step() + .save_state() + + + xorl_client + .sample() + .batch_sample() + logprobs + routing data + + + + + + + 3 + HTTP :5555 + + + 4 + HTTP :30000 + + + + xorl + Training Server + FastAPI → ZMQ → NCCL + FSDP2 · TP · EP · PP + + GPU 0,1,2,3 + + + xorl-sglang + Inference Server + SGLang fork + FP8 · TP · Continuous batching + + GPU 4,5,6,7 + + + + + + + + + NCCL Weight Sync + + + + 5 + + + + + + + + + + + rollouts + logprobs ▲ + + ▼ train with updated policy + diff --git a/docs/src/content/docs/getting-started/installation.md b/docs/src/content/docs/getting-started/installation.md index 1cea3633..7b70716b 100644 --- a/docs/src/content/docs/getting-started/installation.md +++ b/docs/src/content/docs/getting-started/installation.md @@ -10,6 +10,15 @@ title: "Installation" - PyTorch 2.10+ - NVIDIA Hopper GPU (H100/H800) or newer recommended for NVFP4 and DeepEP +## Clone the repo + +```bash +git clone --recurse-submodules https://github.com/togethercomputer/xorl-internal +cd xorl-internal +``` + +> Already cloned without `--recurse-submodules`? Run `git submodule update --init --recursive` + ## Install with uv (recommended) [uv](https://github.com/astral-sh/uv) is the recommended package manager for reproducible installs. @@ -18,21 +27,59 @@ title: "Installation" # Install uv if not already installed curl -LsSf https://astral.sh/uv/install.sh | sh -# Clone and install -git clone https://github.com/togethercomputer/xorl -cd xorl +# Install and activate uv sync source .venv/bin/activate ``` -`uv sync` reads `pyproject.toml` and installs all pinned dependencies. +`uv sync` reads `pyproject.toml` and installs all pinned dependencies into a `.venv` virtual environment. -## Install with pip +## Install with conda ```bash +conda create -n xorl python=3.12 +conda activate xorl pip install -e . ``` +## Install Submodules + +The repo ships two git submodules under `submodules/`: + +| Submodule | Description | +|---|---| +| [xorl-client](https://github.com/togethercomputer/xorl-client) | Lightweight Python client for the XoRL training service. Required for server/RL training mode. | +| [xorl-sglang](https://github.com/togethercomputer/xorl-sglang) | XoRL's fork of [SGLang](https://github.com/sgl-project/sglang). Used as the inference engine in online RL loops. | + +Install individually: + +```bash +pip install -e submodules/xorl-client +pip install -e "submodules/xorl-sglang/python[all]" +``` + +Alternatively, use the bundled `pyproject.sglang.toml` which pins PyTorch to 2.9.1 (required by sglang) and installs xorl, xorl-client, and xorl-sglang together: + +**uv:** +```bash +cp pyproject.sglang.toml pyproject.toml +uv sync +source .venv/bin/activate +``` + +**conda:** +```bash +conda create -n xorl-sglang python=3.12 +conda activate xorl-sglang +cp pyproject.sglang.toml pyproject.toml +pip install -e . +``` + +> **Note:** The default `pyproject.toml` uses PyTorch 2.10.0. sglang requires PyTorch 2.9.1, so the two cannot coexist in the same environment unless you use `pyproject.sglang.toml`. + +> These submodules are only needed for **server training / online RL**. If you are only running local SFT or pretraining, you can skip this step. + + ## Key Dependencies | Package | Version | Notes | diff --git a/docs/src/content/docs/getting-started/quickstart.md b/docs/src/content/docs/getting-started/quickstart.md index 8d4af41b..e53bc385 100644 --- a/docs/src/content/docs/getting-started/quickstart.md +++ b/docs/src/content/docs/getting-started/quickstart.md @@ -68,7 +68,7 @@ Start the training server: python -m xorl.server.launcher \ --mode auto \ --config examples/server/configs/full/qwen3_8b_full.yaml \ - --api-port 5555 + --api-port 6000 ``` Then drive training from a Python client. All training endpoints use a **two-phase async pattern**: the POST returns a `request_id` immediately, and you poll `/api/v1/retrieve_future` to get the actual result. @@ -77,7 +77,7 @@ Then drive training from a Python client. All training endpoints use a **two-pha import requests import time -base_url = "http://localhost:5555" +base_url = "http://localhost:6000" # Check health requests.get(f"{base_url}/health").json() @@ -107,6 +107,50 @@ future = requests.post(f"{base_url}/api/v1/optim_step", json={ }).json() ``` +### Example: SFT on No Robots + +[`examples/server/no_robot_sft/`](https://github.com/togethercomputer/xorl-internal/tree/main/examples/server/no_robot_sft) — Supervised fine-tuning on the [No Robots](https://huggingface.co/datasets/HuggingFaceH4/no_robots) dataset using `xorl_client`. + +```bash +# 1. Start the training server +python -m xorl.server.launcher \ + --mode auto \ + --config examples/server/configs/lora/qwen3_8b_lora.yaml \ + --api-port 6000 + +# 2. Run the SFT script (in another terminal) +pip install xorl-client tinker-cookbook +python examples/server/no_robot_sft/run_sft.py \ + --config.base_url http://localhost:6000 \ + --config.model_name Qwen/Qwen3-8B \ + --config.lora_rank 32 +``` + +The checked-in server config above uses `Qwen/Qwen3-8B` with LoRA rank 32, so the example overrides the script's older 4B defaults. +The script uses `xorl_client.TrainingClient` to drive a LoRA SFT loop with online tokenization, linear LR decay, periodic checkpointing, and per-token NLL validation. + +### Example: Password Memorization (end-to-end weight sync) + +[`examples/server/password_memorization/`](https://github.com/togethercomputer/xorl-internal/tree/main/examples/server/password_memorization) — End-to-end test for the training → weight sync → inference pipeline. Trains a model to memorize 3 secret codes via SFT, syncs weights to a running xorl-sglang instance, and queries inference to verify recall. + +```bash +# 1. Start the training server +python -m xorl.server.launcher \ + --mode auto \ + --config examples/server/configs/full/qwen3_8b_full.yaml \ + --api-port 6000 + +# 2. Start xorl-sglang inference (in another terminal) +python -m sglang.launch_server \ + --model-path Qwen/Qwen3-8B-FP8 --port 30000 + +# 3. Run the test (in another terminal) +python examples/server/password_memorization/run_password_test.py \ + --model Qwen/Qwen3-8B --steps 16 --lr 1e-5 +``` + +Supports all training modes (full, LoRA, QLoRA nvfp4/block_fp8/nf4), LR schedules (constant, cosine, warmup+cosine), and FP8 weight sync re-quantization. See the [example README](https://github.com/togethercomputer/xorl-internal/tree/main/examples/server/password_memorization/README.md) for the full test matrix across Qwen3-8B, Qwen3-30B, and Qwen3-235B. + ## LoRA Fine-tuning ```yaml diff --git a/docs/src/content/docs/server-training/client-sdk/loss-functions.md b/docs/src/content/docs/server-training/client-sdk/loss-functions.md new file mode 100644 index 00000000..1e714f9c --- /dev/null +++ b/docs/src/content/docs/server-training/client-sdk/loss-functions.md @@ -0,0 +1,149 @@ +--- +title: "Loss Functions" +--- + +xorl supports multiple loss functions for SFT and RL training, configured via the `loss_fn` parameter in `forward_backward()`. + +## Overview + +| Algorithm | `loss_fn` | Key params | When to use | +|---|---|---|---| +| SFT / continued pretraining | `causallm_loss` | — | Standard next-token prediction | +| PPO | `policy_loss` | `eps_clip=0.2` | Clipped policy gradient, most stable for RL | +| GRPO (simpler RL) | `importance_sampling` | — | No clipping, simpler but less stable | +| PPO + stale data correction | `policy_loss` | `use_tis=True` | Multiple epochs over same rollout | + +--- + +## Causal LM Loss (`causallm_loss`) + +Standard cross-entropy loss for SFT and pretraining: + +```python +fwd = client.forward_backward(data, loss_fn="causallm_loss") +``` + +--- + +## PPO Policy Loss (`policy_loss`) + +Full PPO-style clipped policy gradient loss ([source](https://github.com/togethercomputer/xorl-internal/blob/main/src/xorl/ops/loss/policy_loss.py)): + +``` +ratio = exp(new_logprobs - old_logprobs) +pg_loss = max(ratio × A, clip(ratio, 1-ε, 1+ε_high) × A) +``` + +**Parameters:** + +| Parameter | Default | Description | +|---|---|---| +| `eps_clip` | `0.2` | Lower clip ratio | +| `eps_clip_high` | `0.2` | Upper clip ratio (asymmetric clipping) | +| `eps_clip_c` | `null` | Dual-clip for negative advantages (e.g. `3.0`) | +| `compute_kl_stats` | `false` | Return KL statistics in metrics | + +**Metrics returned:** + +| Metric | Description | +|---|---| +| `pg_clipfrac` | Fraction of tokens where clipping was applied | +| `kl_sample_train_k3` | Schulman's K3 KL estimator | +| `entropy_sample` | Mean entropy of old policy | +| `ratio_mean/min/max` | Importance sampling ratio statistics | + +```python +fwd = client.forward_backward(data, loss_fn="policy_loss", loss_fn_params={ + "eps_clip": 0.2, + "eps_clip_high": 0.2, + "eps_clip_c": 3.0, + "compute_kl_stats": True, +}) +``` + +--- + +## GRPO / Importance Sampling (`importance_sampling`) + +Simpler importance-sampling loss ([source](https://github.com/togethercomputer/xorl-internal/blob/main/src/xorl/ops/loss/importance_sampling_loss.py)): + +``` +ratio = exp(new_logprobs - old_logprobs) +loss = -(ratio × advantages).mean() +``` + +No clipping — relies on advantages being bounded. Suitable when the policy doesn't drift far from the rollout policy. + +```python +fwd = client.forward_backward(data, loss_fn="importance_sampling", loss_fn_params={ + "compute_kl_stats": True, +}) +``` + +--- + +## IcePop + +IcePop (from [GLM-5, arXiv:2602.15763](https://arxiv.org/abs/2602.15763)) is a **hard masking** technique that zeros gradients for tokens where the importance sampling ratio falls outside `[1/β, β]`. Complementary to PPO's soft clipping. + +```python +fwd = client.forward_backward(data, loss_fn="policy_loss", loss_fn_params={ + "eps_clip": 0.2, + "icepop_beta": 5.0, # zero gradients when ratio < 0.2 or ratio > 5.0 +}) +``` + +--- + +## TIS — Temporal Importance Sampling + +Corrects for policy drift when running multiple training steps on the same rollout batch: + +``` +tis_weight = clip(exp(train_logprobs - rollout_logprobs), tis_clip_low, tis_clip_high) +loss = (tis_weight × pg_loss).mean() +``` + +Requires passing `rollout_logprobs` separately from `logprobs`: + +```python +datum = xorl_client.Datum( + model_input=xorl_client.ModelInput.from_ints(token_ids), + loss_fn_inputs={ + "labels": token_ids, + "logprobs": logprobs_at_last_train_step, + "rollout_logprobs": logprobs_at_rollout, # fixed from inference + "advantages": advantages, + }, +) + +fwd = client.forward_backward([datum], loss_fn="policy_loss", loss_fn_params={ + "use_tis": True, + "tis_clip_low": 0.1, + "tis_clip_high": 2.0, +}) +``` + +--- + +## R3 — Routing Replay for MoE + +For MoE models, R3 replays expert routing decisions from inference during training to ensure gradient consistency. Pass routing data from xorl-sglang on each `Datum`: + +```python +datum = xorl_client.Datum( + model_input=xorl_client.ModelInput.from_ints(token_ids), + loss_fn_inputs={ + "labels": token_ids, + "logprobs": rollout_logprobs, + "advantages": advantages, + }, + routed_experts=rollout_routing_indices, # [T, L, K] from sglang +) + +fwd = client.forward_backward([datum], loss_fn="policy_loss") +``` + +The current `xorl-client` SDK only exposes `routed_experts`. The server can also consume `routed_expert_logits`, but that field is not yet wired through `Datum` / `TrainingClient`. + +See the [Router page](/moe/router/#routing-replay-r3) for details on how R3 works, and the [xorl-sglang page](/server-training/sglang/) for how routing data is exported from inference. diff --git a/docs/src/content/docs/server-training/client-sdk/overview.md b/docs/src/content/docs/server-training/client-sdk/overview.md new file mode 100644 index 00000000..72234586 --- /dev/null +++ b/docs/src/content/docs/server-training/client-sdk/overview.md @@ -0,0 +1,190 @@ +--- +title: "Client SDK Overview" +--- + +[xorl-client](https://github.com/togethercomputer/xorl-client) is the Python SDK for driving the xorl training server. It is lightweight (no PyTorch dependency), async-first, and Tinker API compatible. + +## Installation + +```bash +pip install xorl-client +``` + +Or from source: + +```bash +pip install git+https://github.com/togethercomputer/xorl-client.git +``` + +xorl-client is also installed automatically with xorl — see the [installation guide](/getting-started/installation/). + +--- + +## Client Classes + +| Class | Purpose | +|---|---| +| [`ServiceClient`](#serviceclient) | Main entry point — connects to the server, creates all other clients | +| [`TrainingClient`](#trainingclient) | Training loop: `forward_backward()`, `optim_step()`, `save_state()`, `load_state()` | +| [`SamplingClient`](#samplingclient) | Rollout generation via xorl-sglang | +| [`RestClient`](#restclient) | Checkpoint management: list, delete, get metadata | + +All methods return `APIFuture` objects — call `.result()` to block or `await` in async code. + +--- + +### ServiceClient + +The main entry point. Connects to the xorl training server and provides factory methods for all other clients. + +```python +import xorl_client + +service = xorl_client.ServiceClient(base_url="http://localhost:6000") +``` + +| Method | Returns | Description | +|---|---|---| +| `create_training_client(base_model)` | `TrainingClient` | Full-weight training (no LoRA) | +| `create_lora_training_client(base_model, rank, ...)` | `TrainingClient` | LoRA training with configurable rank/alpha | +| `create_sampling_client(base_url, model_path)` | `SamplingClient` | Load a saved LoRA adapter on xorl-sglang and return a `SamplingClient` | +| `create_rest_client(model_id)` | `RestClient` | Checkpoint management | +| `create_training_client_from_state(checkpoint_path)` | `TrainingClient` | Resume training from checkpoint (weights only) | +| `create_training_client_from_state_with_optimizer(checkpoint_path)` | `TrainingClient` | Resume training with optimizer state | + +Environment variables: `XORL_BASE_URL` (default server URL), `XORL_API_KEY` (authentication). + +`create_sampling_client(...)` is the LoRA helper path: it calls `/api/v1/create_sampling_session` +on the training server to load a saved LoRA adapter on the registered inference workers, then returns a `SamplingClient`. +Full-weight exports from `save_weights_for_sampler()` are not dynamically loaded by this helper. + +--- + +### TrainingClient + +Drives the training loop. All methods return `APIFuture` — submit multiple calls without blocking, then collect results. + +```python +client = service.create_lora_training_client( + base_model="Qwen/Qwen3-8B", rank=32 +) + +# Non-blocking pipeline: both submitted immediately +fwd = client.forward_backward(data, loss_fn="importance_sampling") +opt = client.optim_step(xorl_client.AdamParams(learning_rate=1e-5)) + +# Collect results +result = fwd.result() +opt.result() +print(f"logprobs={result.loss_fn_outputs[0]['logprobs']}") +``` + +| Method | Description | +|---|---| +| `forward_backward(data, loss_fn, loss_fn_params)` | Compute loss + accumulate gradients | +| `forward(data, loss_fn)` | Forward-only (validation, reference logprobs) | +| `optim_step(adam_params)` | Apply gradients with Adam optimizer | +| `sync_inference_weights(master_address, ...)` | Broadcast current weights to registered inference endpoints | +| `save_state(path)` | Save full checkpoint (model + optimizer) | +| `load_state(path)` | Load checkpoint | +| `save_weights_for_sampler(path)` | Save inference weights under `sampler_weights/` (LoRA adapter in LoRA mode, full HF checkpoint otherwise) | +| `get_tokenizer()` | Get the model's tokenizer | + +Requests are automatically ordered by `seq_id` — `forward_backward` always executes before `optim_step` regardless of when they arrive. + +**Data format:** Training data is a list of `Datum` objects: + +```python +datum = xorl_client.Datum( + model_input=xorl_client.ModelInput.from_ints(input_ids), + loss_fn_inputs={"labels": labels, "logprobs": old_logprobs, "advantages": advantages}, +) +``` + +For RL losses, `model_input` should contain the full prompt + completion sequence. +`SamplingClient` returns generated output tokens only, so reconstruct the full sequence before calling `forward_backward()`, +mask prompt positions with `advantages=0.0`, and make `logprobs` line up with the server's shifted `full_ids[1:]` view. + +--- + +### SamplingClient + +Connects to xorl-sglang for rollout generation. Supports batch sampling with straggler handling. + +```python +sampler = xorl_client.SamplingClient(base_url="http://localhost:30000") + +# Single sample +response = sampler.sample( + "What is 2+2?", + xorl_client.SamplingParams(max_tokens=64), +).result() + +# Batch sampling with timeout-based straggler handling +result = sampler.sample_batch( + prompts, + xorl_client.SamplingParams(max_tokens=64), + timeout=30.0, +).result() +# result.completed, result.failed, result.cancelled +``` + +`response.tokens` and `response.logprobs` contain generated output tokens only, not the prompt tokens. +If you use these results for RL training, keep the prompt token IDs so you can rebuild full prompt + completion sequences. + +--- + +### RestClient + +Checkpoint management operations. + +```python +rest = service.create_rest_client() + +# List checkpoints +checkpoints = rest.list_checkpoints().result() +for cp in checkpoints.checkpoints: + print(f"{cp.checkpoint_id}: {cp.checkpoint_type}") + +# Delete a checkpoint +rest.delete_checkpoint("my_checkpoint").result() + +# Get checkpoint metadata +info = rest.get_weights_info("xorl://default/weights/step_42").result() +``` + +--- + +## Key Concepts + +### APIFuture + +All client methods return `APIFuture` objects for non-blocking operation: + +```python +# Sync usage +result = client.forward_backward(data, "causallm_loss").result() + +# Async usage +result = await client.forward_backward(data, "causallm_loss") + +# Pipeline: submit all, then wait +fwd = client.forward_backward(data, "policy_loss") +opt = client.optim_step(adam_params) +fwd_result = fwd.result() # blocks until forward_backward completes +opt.result() # blocks until optim_step completes +``` + +### Automatic Chunking + +Large batches of `Datum` objects are automatically split into chunks to stay within HTTP payload limits. This is transparent — you pass the full batch and xorl-client handles the rest. + +### Tinker API Compatibility + +xorl-client is wire-compatible with the Tinker training API protocol. Field mappings are handled automatically: + +| Tinker field | xorl field | +|---|---| +| `session_id` | `model_id` | +| `loss_fn_config` | `loss_fn_params` | +| `model_input.chunks[].tokens` | `model_input.input_ids` | diff --git a/docs/src/content/docs/server-training/client-sdk/training-loop.md b/docs/src/content/docs/server-training/client-sdk/training-loop.md new file mode 100644 index 00000000..e9fa120f --- /dev/null +++ b/docs/src/content/docs/server-training/client-sdk/training-loop.md @@ -0,0 +1,197 @@ +--- +title: "Training Loop Patterns" +--- + +Common patterns for using `xorl_client` to build training loops. + +## SFT Training Loop + +Supervised fine-tuning with a fixed dataset: + +```python +import xorl_client + +service = xorl_client.ServiceClient(base_url="http://localhost:6000") +client = service.create_lora_training_client( + base_model="Qwen/Qwen3-8B", rank=32 +) +tokenizer = client.get_tokenizer() +adam = xorl_client.AdamParams(learning_rate=1e-4) + +for step, batch in enumerate(dataloader): + data = [ + xorl_client.Datum( + model_input=xorl_client.ModelInput.from_ints(sample["input_ids"]), + loss_fn_inputs={"labels": sample["labels"]}, + ) + for sample in batch + ] + + fwd = client.forward_backward(data, loss_fn="causallm_loss") + opt = client.optim_step(adam) + result = fwd.result() + opt.result() + + print(f"step={step} logprobs={result.loss_fn_outputs[0]['logprobs']}") +``` + +--- + +## RL Training Loop (PPO / GRPO) + +Online RL loop with rollout generation, reward scoring, and policy updates: + +```python +import xorl_client +import requests + +service = xorl_client.ServiceClient(base_url="http://localhost:6000") +client = service.create_training_client(base_model="Qwen/Qwen3-8B") +tokenizer = client.get_tokenizer() +sampler = xorl_client.SamplingClient(base_url="http://localhost:30000") +adam = xorl_client.AdamParams(learning_rate=1e-6, beta1=0.9, beta2=0.95, eps=1e-8) + +# Register inference endpoint +requests.post("http://localhost:6000/add_inference_endpoint", json={ + "host": "localhost", + "port": 30000, + "worker_port": 30000, + "world_size": 1, +}) + +def build_rl_datum(prompt_ids, response, advantage): + completion_ids = response.tokens + completion_logprobs = response.logprobs or [0.0] * len(completion_ids) + completion_advantages = [advantage] * len(completion_ids) + full_ids = prompt_ids + completion_ids + + return xorl_client.Datum( + model_input=xorl_client.ModelInput.from_ints(full_ids), + loss_fn_inputs={ + "labels": full_ids, + # The server shifts `labels` to `full_ids[1:]`, so old logprobs + # need to line up with that shifted view. + "logprobs": [0.0] * max(len(prompt_ids) - 1, 0) + completion_logprobs, + "advantages": [0.0] * len(prompt_ids) + completion_advantages, + }, + ) + +for step in range(num_steps): + # 1. Generate rollouts + prompt_token_ids = [ + tokenizer.encode(prompt, add_special_tokens=False) + for prompt in prompts + ] + prompt_inputs = [ + xorl_client.ModelInput.from_ints(token_ids) + for token_ids in prompt_token_ids + ] + + batch = sampler.sample_batch( + prompt_inputs, + xorl_client.SamplingParams(max_tokens=512), + return_logprobs=True, + timeout=30.0, + ).result() + completions = [response for _, response in sorted(batch.completed)] + if len(completions) != len(prompts): + continue # retry or drop stragglers in a real loop + + # 2. Score with reward model + rewards = reward_model.score([response.text for response in completions]) + sample_advantages = compute_advantages(rewards) + + # 3. Pack prompt + completion for policy loss + data = [ + build_rl_datum(prompt_token_ids[i], response, sample_advantages[i]) + for i, response in enumerate(completions) + ] + + # 4. Train + fwd = client.forward_backward(data, loss_fn="policy_loss", loss_fn_params={ + "eps_clip": 0.2, "compute_kl_stats": True, + }) + opt = client.optim_step(adam) + result = fwd.result() + opt.result() + + # 5. Sync weights to inference + if step % sync_every == 0: + client.sync_inference_weights( + master_address="localhost", + master_port=29600, + ).result() +``` + +`SamplingClient` returns completion-only tokens and logprobs. +The helper above rebuilds the full prompt + completion sequence, zero-pads prompt `advantages`, +and zero-pads prompt `logprobs` so they line up with the server's shifted `full_ids[1:]` view. + +--- + +## Multi-Adapter (LoRA) + +Run multiple LoRA adapters simultaneously, switchable per request via `model_id`: + +```python +# Create policy and reference adapters +policy = service.create_lora_training_client( + base_model="Qwen/Qwen3-8B", rank=32, model_id="policy" +) +reference = service.create_lora_training_client( + base_model="Qwen/Qwen3-8B", rank=16, model_id="reference" +) + +# Train policy, forward-only on reference for KL +fwd = policy.forward_backward(data, loss_fn="policy_loss") +ref_logprobs = reference.forward(data, loss_fn="cross_entropy") +``` + +--- + +## Checkpoint Resume + +Resume training from a saved checkpoint: + +```python +# Weights only (optimizer state reset) +client = service.create_training_client_from_state( + checkpoint_path="xorl://default/weights/step_1000", + base_model="Qwen/Qwen3-8B", +) + +# With optimizer state (exact resume) +client = service.create_training_client_from_state_with_optimizer( + checkpoint_path="xorl://default/weights/step_1000", + base_model="Qwen/Qwen3-8B", +) +``` + +--- + +## Save LoRA Weights for Inference + +For LoRA training, save an adapter checkpoint and ask the training server to load it on the registered xorl-sglang workers: + +```python +service = xorl_client.ServiceClient(base_url="http://localhost:6000") +client = service.create_lora_training_client( + base_model="Qwen/Qwen3-8B", rank=32 +) + +# Save LoRA adapter weights for inference +client.save_weights_for_sampler("step_100").result() + +# Load that adapter on xorl-sglang and get a SamplingClient +sampler = service.create_sampling_client( + base_url="http://localhost:30000", + model_path="xorl://default/sampler_weights/step_100", +) +response = sampler.sample("Hello", xorl_client.SamplingParams(max_tokens=64)).result() +``` + +`create_sampling_client(...)` currently loads saved LoRA adapters via `/api/v1/create_sampling_session`. +Launch xorl-sglang with `--enable-lora` for this flow, and use `dp_size == 1` on the SGLang side. + +For full-weight training, `save_weights_for_sampler()` exports a full Hugging Face checkpoint instead. +Serve that checkpoint by launching xorl-sglang from the exported path, or use `sync_inference_weights()` for online rollout serving. diff --git a/docs/src/content/docs/server-training/examples/password-memorization.md b/docs/src/content/docs/server-training/examples/password-memorization.md new file mode 100644 index 00000000..e569aa47 --- /dev/null +++ b/docs/src/content/docs/server-training/examples/password-memorization.md @@ -0,0 +1,73 @@ +--- +title: "Password Memorization" +--- + +[`examples/server/password_memorization/`](https://github.com/togethercomputer/xorl-internal/tree/main/examples/server/password_memorization) — End-to-end test for the full **training → weight sync → inference** pipeline. Trains a model to memorize 3 secret project codes via SFT, syncs weights to a running xorl-sglang instance, and queries inference to verify recall. + +**Run:** + +```bash +# 1. Start training server +CUDA_VISIBLE_DEVICES=0,1,2,3 python -m xorl.server.launcher \ + --mode auto \ + --config examples/server/configs/full/qwen3_8b_full.yaml \ + --api-port 6000 + +# 2. Start xorl-sglang inference (in another terminal) +CUDA_VISIBLE_DEVICES=4 python -m sglang.launch_server \ + --model-path Qwen/Qwen3-8B-FP8 --port 30000 + +# 3. Run the test (in another terminal) +python examples/server/password_memorization/run_password_test.py \ + --model Qwen/Qwen3-8B --steps 16 --lr 1e-5 +``` + +**Options:** + +| Flag | Default | Description | +|---|---|---| +| `--model` | (required) | HuggingFace model name | +| `--steps` | 64 | Total training steps | +| `--lr` | 1e-4 | Peak learning rate | +| `--lr-schedule` | constant | `constant`, `cosine`, or `warmup_cosine` | +| `--sync-quant` | fp8 | Sync quantization: `fp8` or `none` | +| `--train-url` | `http://localhost:6000` | Training server URL | +| `--infer-url` | `http://localhost:30000` | Inference endpoint URL | + +## Weight sync pipeline + +| Training mode | Sync path | +|---|---| +| Full-weight bf16 | bf16 → fp8 requant → SGLang | +| LoRA | bf16 base + LoRA merged → fp8 requant → SGLang | +| QLoRA nvfp4 | nvfp4 dequant → bf16 merged → fp8 requant → SGLang | +| QLoRA block_fp8 | fp8 dequant → bf16 merged → fp8 requant → SGLang | +| QLoRA nf4 | nf4 dequant → bf16 merged → fp8 requant → SGLang | + +## Test matrix + +**Qwen3-8B (4x H100):** + +| Mode | Steps | LR | Schedule | Result | +|---|---|---|---|---| +| Full-weight bf16 | 16 | 1e-5 | constant | 3/3 | +| LoRA rank 32 | 32 | 1e-4 | constant | 3/3 | +| QLoRA nvfp4 | 64 | 5e-5 | cosine | 3/3 | +| QLoRA block_fp8 | 64 | 5e-4 | constant | 3/3 | + +**Qwen3-Coder-30B-A3B (4-8x H100):** + +| Mode | Parallelism | Steps | LR | Schedule | Result | +|---|---|---|---|---|---| +| Full-weight bf16 | SP=4, shard=2 | 32 | 1e-5 | constant | 3/3 | +| LoRA rank 32 | SP=4 | 32 | 1e-4 | constant | 3/3 | +| QLoRA nvfp4 | EP=4, SP=4 | 128 | 5e-4 | cosine | 3/3 | + +**Qwen3-235B-A22B (8x H100, cross-node inference):** + +| Mode | Parallelism | Steps | LR | Schedule | Result | +|---|---|---|---|---|---| +| QLoRA nvfp4 | EP=8, SP=8 | 128 | 5e-4 | cosine | 3/3 | +| QLoRA nf4 | EP=8, SP=8 | 128 | 5e-4 | cosine | 3/3 | + +See the [example README](https://github.com/togethercomputer/xorl-internal/tree/main/examples/server/password_memorization/README.md) for the full test matrix and detailed setup instructions. diff --git a/docs/src/content/docs/server-training/examples/sft-no-robots.md b/docs/src/content/docs/server-training/examples/sft-no-robots.md new file mode 100644 index 00000000..9f61aef4 --- /dev/null +++ b/docs/src/content/docs/server-training/examples/sft-no-robots.md @@ -0,0 +1,45 @@ +--- +title: "SFT on No Robots" +--- + +[`examples/server/no_robot_sft/`](https://github.com/togethercomputer/xorl-internal/tree/main/examples/server/no_robot_sft) — Supervised fine-tuning on the [No Robots](https://huggingface.co/datasets/HuggingFaceH4/no_robots) dataset. + +**What it demonstrates:** +- LoRA SFT training loop driven by `xorl_client.TrainingClient` +- Online tokenization using tinker-cookbook renderers +- Initial validation step (forward-only, no gradients) +- Linear learning rate decay +- Periodic checkpoint saving and resume support +- Per-token NLL metrics + +**Run:** + +```bash +# 1. Start the training server (4 GPUs) +CUDA_VISIBLE_DEVICES=0,1,2,3 python -m xorl.server.launcher \ + --mode auto \ + --config examples/server/configs/lora/qwen3_8b_lora.yaml \ + --api-port 6000 + +# 2. Run SFT (in another terminal) +pip install xorl-client tinker-cookbook +python examples/server/no_robot_sft/run_sft.py \ + --config.base_url http://localhost:6000 \ + --config.model_name Qwen/Qwen3-8B \ + --config.lora_rank 32 +``` + +The checked-in server config above loads `Qwen/Qwen3-8B` with LoRA rank 32. +`run_sft.py` still defaults to an older 4B example, so pass the overrides above until those defaults are updated. + +**Config options:** + +| Field | Default | Description | +|---|---|---| +| `base_url` | `http://localhost:6000` | Training server URL | +| `model_name` | `Qwen/Qwen3-4B-Instruct-2507` | Model name (for tokenizer). Override to `Qwen/Qwen3-8B` for the 8B LoRA config above. | +| `batch_size` | 128 | Training batch size | +| `learning_rate` | 1e-4 | Peak learning rate | +| `max_length` | 32768 | Max sequence length | +| `lora_rank` | 64 | LoRA rank. Override to `32` to match `examples/server/configs/lora/qwen3_8b_lora.yaml`. | +| `save_every` | 20 | Checkpoint every N steps (0 = disabled) | diff --git a/docs/src/content/docs/server-training/overview.md b/docs/src/content/docs/server-training/overview.md index fa96a200..1c7c37d6 100644 --- a/docs/src/content/docs/server-training/overview.md +++ b/docs/src/content/docs/server-training/overview.md @@ -1,493 +1,155 @@ --- -title: "Server Training" +title: "Server Training for RL" --- -Server training exposes the training loop as a REST API, enabling external processes (RL frameworks, data pipelines, experiment orchestrators) to drive gradient updates step by step. This is the primary mode for online RL training. +Server training exposes the training loop as a REST API, enabling external processes to drive gradient updates step by step. This is the primary mode for online RL training. -## xorl-client +## How Online RL Works with XoRL -To use server training from Python, install the official client library: +Online RL requires a **training server**, an **inference server**, and an **orchestrator** with a reward signal. XoRL provides the first three — you bring the reward. -```bash -pip install git+https://github.com/togethercomputer/xorl-client.git -``` +![XoRL RL Architecture](../../../assets/rl-architecture.svg) -The client handles HTTP communication, request serialization, async result polling, and provides typed wrappers around all API endpoints. All examples in this page use `xorl_client`. - ---- +**The RL loop:** -## Architecture +1. **Generate rollouts** — send prompts to xorl-sglang via `xorl_client.SamplingClient`. Returns completions + per-token logprobs. +2. **Score** — your reward model / verifier / environment scores the rollouts. +3. **Train** — pack scored rollouts into `Datum` objects, call `forward_backward()` + `optim_step()` on the xorl training server. +4. **Sync weights** — broadcast updated weights to xorl-sglang via NCCL. KV cache is flushed automatically. +5. **Repeat** with the updated policy. -``` -Client (xorl_client / HTTP) - │ - ▼ HTTP (port 5555) -API Server (FastAPI) - │ - ▼ ZMQ + NCCL -Runner Dispatcher (rank 0) - │ - ▼ NCCL -Worker Ranks (1 … N-1) -``` - -The API server receives requests, forwards them to the model runner via ZMQ, and broadcasts to all GPU workers via NCCL collectives. All ranks participate in every forward/backward. +| Component | Provided by | Description | +|---|---|---| +| Training server | **[xorl](/server-training/training-server/launching/)** | Forward/backward, optimizer, checkpointing, parallelism | +| Inference server | **[xorl-sglang](/server-training/sglang/)** | Rollout generation, per-token logprobs, weight sync | +| Client SDK | **[xorl-client](/server-training/client-sdk/overview/)** | Python SDK: `TrainingClient`, `SamplingClient`, `RestClient` | +| Reward / environment | **You** | Reward model, code sandbox, math verifier, rule-based scorer | --- -## Starting the Server +## Quick Start -### Single Node +**1. Start the training server:** ```bash -python -m xorl.server.launcher \ +CUDA_VISIBLE_DEVICES=0,1,2,3 python -m xorl.server.launcher \ --mode auto \ --config examples/server/configs/full/qwen3_8b_full.yaml \ - --api-port 5555 + --api-port 6000 ``` -`--mode auto` launches `torchrun` workers internally and starts the API on the specified port. +**2. Start xorl-sglang inference:** -### Multi-Node - -On worker nodes (one command per node): -```bash -torchrun --nproc_per_node=8 --nnodes=4 --node_rank= \ - --master_addr=HEAD_IP --master_port=29500 \ - -m xorl.server.runner.runner_dispatcher \ - --config config.yaml -``` - -On the head node (connect to already-running workers): ```bash -python -m xorl.server.launcher \ - --mode connect \ - --config config.yaml \ - --api-port 5555 \ - --master-addr HEAD_IP +CUDA_VISIBLE_DEVICES=4 python -m sglang.launch_server \ + --model-path Qwen/Qwen3-8B-FP8 \ + --port 30000 \ + --tp-size 1 \ + --rl-on-policy-target xorl \ + --enable-fp32-router \ + --enable-fp32-lm-head ``` -See `scripts/run_multinode_server.sh` for a full multi-node launch script. - -### Launcher CLI Options - -| Option | Default | Description | -|---|---|---| -| `--mode` | `auto` | `auto` (launch workers) or `connect` (attach to existing) | -| `--config` | required | Path to server YAML config | -| `--api-host` | `0.0.0.0` | Bind address for the REST API | -| `--api-port` | auto | Port for the REST API (auto-finds a free port if not specified) | -| `--nnodes` | `1` | Number of nodes (auto mode only) | -| `--master-addr` | `127.0.0.1` | Head node address | -| `--master-port` | `29500` | torchrun master port | -| `--operation-timeout` | `1800` | Max seconds per operation | -| `--log-level` | `INFO` | `DEBUG`, `INFO`, `WARNING` | - ---- - -## Server Configuration - -Server configs have a flat structure (no `train:` section nesting). See [Server Config Reference](/config-reference/server/) for all fields. - -```yaml -# examples/server/configs/full/qwen3_8b_full.yaml -model_path: Qwen/Qwen3-8B -tokenizer_path: Qwen/Qwen3-8B -attn_implementation: flash_attention_3 - -# Parallelism -data_parallel_mode: fsdp2 -data_parallel_shard_size: 8 -expert_parallel_size: 1 -pipeline_parallel_size: 1 -ulysses_parallel_size: 1 - -# Memory -enable_mixed_precision: true -enable_gradient_checkpointing: true -enable_full_shard: true -init_device: meta - -# Packing -sample_packing_sequence_len: 8192 -enable_packing: true - -# Checkpointing -output_dir: outputs/server_run -ckpt_manager: dcp -``` - ---- - -## Sequence Packing - -The server packs multiple client-submitted sequences into a single training bin to maximize GPU utilization. Packing is controlled by two config fields: - -| Field | Default | Description | -|---|---|---| -| `enable_packing` | `true` | Pack multiple sequences per bin. Set `false` to send one sequence per forward pass. | -| `sample_packing_sequence_len` | `32000` | Bin capacity in tokens. Sequences are concatenated until this limit is reached. | - -Unlike local training where packing happens offline during dataset preparation, the server packs sequences **at request time** inside the `RequestProcessor`. The packing algorithm is **sequential** (greedy, preserves submission order): incoming `Datum` objects are concatenated into a bin until `sample_packing_sequence_len` is reached, then a new bin is started. - -Each packed bin carries `cu_seqlens` (cumulative sequence lengths) and `position_ids` so Flash Attention treats each original sequence independently within the bin — there is no cross-sequence attention leakage. - -For Ring Attention (`ringattn_parallel_size > 1`), each document in the bin must be individually padded to a length divisible by `2 × ringattn_parallel_size` before zigzag sharding. The server handles this padding automatically in `TextSequenceShardCollator` using sequential dummy position IDs for the pad region. - ---- - -## Python Client (`xorl_client`) - -### Connecting +**3. Run training from Python:** ```python import xorl_client -service_client = xorl_client.ServiceClient(base_url="http://localhost:5555") -``` - -`ServiceClient` is the main entry point. It manages the connection, authentication, and provides factory methods for training clients. +service = xorl_client.ServiceClient(base_url="http://localhost:6000") +client = service.create_training_client(base_model="Qwen/Qwen3-8B") -### Creating a training client +# Register inference endpoint for weight sync +import requests +requests.post("http://localhost:6000/add_inference_endpoint", json={ + "host": "localhost", + "port": 30000, + "worker_port": 30000, + "world_size": 1, +}) -**Full-weight training** (no LoRA — server auto-registers `model_id="default"` on startup): +# RL loop +for step in range(num_steps): + # ... generate rollouts from sglang, compute rewards ... + fwd = client.forward_backward(data, loss_fn="importance_sampling") + opt = client.optim_step(xorl_client.AdamParams(learning_rate=1e-5)) + fwd.result(); opt.result() -```python -training_client = xorl_client.TrainingClient( - holder=service_client.holder, - model_id="default", - base_model="Qwen/Qwen3-8B", -) + if step % sync_interval == 0: + client.sync_inference_weights( + master_address="localhost", + master_port=29600, + ).result() ``` -**LoRA training** (creates and registers a new adapter): - -```python -training_client = service_client.create_lora_training_client( - base_model="Qwen/Qwen3-8B", - rank=32, - model_id="my_adapter", -) -``` - -### Preparing data - -Training data is passed as a list of `Datum` objects: - -```python -datum = xorl_client.Datum( - model_input=xorl_client.ModelInput.from_ints(input_ids), - loss_fn_inputs={"labels": labels}, -) -``` +For multi-node setup, server configuration, and launcher CLI options, see [Launching & Configuration](/server-training/training-server/launching/). -`ModelInput.from_ints` wraps a flat token list into the expected `{"input_ids": [...]}` format. +### Available configs -### Training loop +| Config | Model | Mode | GPUs | +|---|---|---|---| +| `full/qwen3_8b_full.yaml` | Qwen3-8B | Full-weight bf16 | 4 | +| `lora/qwen3_8b_lora.yaml` | Qwen3-8B | LoRA rank 32 | 4 | +| `qlora/qwen3_8b_qlora_nvfp4.yaml` | Qwen3-8B | QLoRA nvfp4 | 4 | +| `full/qwen3_coder_30b_a3b_full.yaml` | Qwen3-Coder-30B-A3B | Full bf16 (SP=4) | 8 | +| `qlora/qwen3_coder_30b_a3b_qlora.yaml` | Qwen3-Coder-30B-A3B | QLoRA (EP=4, SP=4) | 4 | +| `full/qwen3_235b_a22b_8node_ep64.yaml` | Qwen3-235B-A22B | Full bf16 (EP=64) | 64 | -`forward_backward` and `optim_step` return futures — call `.result()` to block and retrieve the output: +### Matching inference models -```python -adam_params = xorl_client.AdamParams( - learning_rate=1e-5, - beta1=0.9, - beta2=0.95, - eps=1e-8, -) - -for step, batch in enumerate(dataloader): - data = [ - xorl_client.Datum( - model_input=xorl_client.ModelInput.from_ints(sample["input_ids"]), - loss_fn_inputs={"labels": sample["labels"]}, - ) - for sample in batch - ] - - fwd_bwd = training_client.forward_backward(data, loss_fn="causallm_loss") - optim = training_client.optim_step(adam_params) - - result = fwd_bwd.result() - optim.result() - - print(f"step={step} loss={result.loss_fn_outputs[0]['loss'].data[0]:.4f}") -``` - -### Available loss functions - -| `loss_fn` | Description | -|---|---| -| `causallm_loss` | Standard causal language model cross-entropy | -| `policy_loss` | PPO-style policy gradient loss | -| `importance_sampling` | Importance-sampling weighted loss for off-policy RL | +| Training model | Inference model | TP | +|---|---|---| +| `Qwen/Qwen3-8B` | `Qwen/Qwen3-8B-FP8` | 1 | +| `Qwen/Qwen3-Coder-30B-A3B-Instruct` | `Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8` | 2 | +| `Qwen/Qwen3-235B-A22B-Instruct-2507` | `Qwen/Qwen3-235B-A22B-Instruct-2507-FP8` | 4 | --- -## REST API Reference - -All endpoints are at `http://:/api/v1/`. The xorl_client wraps these — use the Python client for most cases. - -### Health Check - -```http -GET /health -``` -```json -{"status": "healthy", "engine_running": true} -``` - -### Create Model - -Creates and registers a new model/adapter session. - -```http -POST /api/v1/create_model -``` -```json -{ - "model_id": "my_adapter", - "base_model": "Qwen/Qwen3-8B", - "lora_config": {"rank": 32, "alpha": 16} -} -``` - -### Forward + Backward - -Computes loss and accumulates gradients. Call multiple times before `optim_step` for gradient accumulation. - -```http -POST /api/v1/forward_backward -``` -```json -{ - "model_id": "default", - "forward_backward_input": { - "data": [ - { - "model_input": {"input_ids": [1, 2, 3, ...]}, - "loss_fn_inputs": {"labels": [1, 2, 3, ...]} - } - ], - "loss_fn": "causallm_loss" - } -} -``` - -### Optimizer Step - -Applies accumulated gradients, clips, and advances the LR scheduler. - -```http -POST /api/v1/optim_step -``` -```json -{ - "model_id": "default", - "adam_params": { - "learning_rate": 1e-5, - "beta1": 0.9, - "beta2": 0.95, - "eps": 1e-8 - }, - "gradient_clip": 1.0 -} -``` - -### Forward Only - -Same as `forward_backward` but without gradient accumulation. Used for evaluation or reference model log-prob computation. - -```http -POST /api/v1/forward -``` - -### Save Checkpoint - -```http -POST /api/v1/save_weights -``` -```json -{"model_id": "default", "path": null} -``` -`null` path = auto-timestamped path under `output_dir`. - -### Load Checkpoint - -```http -POST /api/v1/load_weights -``` -```json -{"model_id": "default", "path": "outputs/server_run/weights/default/step_42"} -``` - -### Weight Sync to Inference - -Push current training weights to registered SGLang inference endpoints. - -```http -POST /api/v1/sync_inference_weights -``` -```json -{ - "master_address": "HEAD_NODE_IP", - "master_port": 29600, - "group_name": "sync_group_0", - "buffer_size_mb": 512 -} -``` - -### Inference Endpoint Management - -```http -POST /add_inference_endpoint -``` -```json -{"host": "inference-node-01", "port": 30000, "world_size": 8} -``` - -```http -POST /remove_inference_endpoint -``` - -### Sleep / Wake - -Offload model weights to CPU to free GPU memory while inference runs: - -```http -POST /sleep -POST /wake_up -``` - -### Kill Session +## xorl-client -```http -POST /api/v1/kill_session -``` -```json -{"model_id": "default", "reset_weights": false} -``` +The [xorl-client](/server-training/client-sdk/overview/) Python SDK drives the training server — see the [Client SDK page](/server-training/client-sdk/overview/) for installation, client classes, loss functions, and training loop examples. --- -## Tinker API Compatibility +## Recommended Reading -xorl is compatible with the **Tinker training API** — a standardized protocol for driving training servers from RL frameworks and experiment orchestrators. Any client built against the Tinker spec works with xorl without modification. +XoRL's server training mode is designed for online RL with LLMs. The following papers provide background on the algorithms and system designs that XoRL supports: -### Tinker-compatible endpoints +### Foundational RL for LLMs -| Endpoint | Description | +| Paper | Description | |---|---| -| `POST /api/v1/create_model` | Create/register a model session | -| `POST /api/v1/unload_model` | Unload and release a session | -| `POST /api/v1/forward_backward` | Forward + backward pass | -| `POST /api/v1/optim_step` | Optimizer step | -| `POST /api/v1/weights_info` | Checkpoint metadata for model loading | -| `GET /api/v1/training_runs` | List training runs | - -### Field mappings +| [Training language models to follow instructions with human feedback](https://arxiv.org/abs/2203.02155) (Ouyang et al., 2022) | InstructGPT — the original RLHF pipeline: SFT → reward model → PPO fine-tuning. Established the standard three-stage approach. | +| [DeepSeekMath: Pushing the Limits of Mathematical Reasoning](https://arxiv.org/abs/2402.03300) (Shao et al., 2024) | Introduces **GRPO** (Group Relative Policy Optimization) — a simpler alternative to PPO that uses group-level advantages without a value model. XoRL supports GRPO via `loss_fn="importance_sampling"`. | +| [DAPO: An Open-Source LLM Reinforcement Learning System](https://arxiv.org/abs/2503.14476) (Yu et al., 2025) | Decoupled clip ratios, dynamic sampling, token-level policy gradient, and overlong reward shaping. Demonstrates RL scaling without a value model. | +| [ReMax: A Simple, Effective, and Efficient Method for Aligning LLMs](https://arxiv.org/abs/2310.10505) (Li et al., 2023) | REINFORCE with a max-reward baseline — no critic needed. Shows that simpler RL methods can match PPO performance with less compute. | -xorl automatically maps Tinker's field names to its own format, so Tinker clients work without changes: +### System Design: Pipelining and Async RL -| Tinker field | xorl field | Where | -|---|---|---| -| `session_id` | `model_id` | All request types | -| `loss_fn_config` | `loss_fn_params` | `forward_backward` input | -| `lora_config.rank` | `lora_rank` | `create_model` request | -| `model_input.chunks[].tokens` | `model_input.input_ids` | `Datum` model input | - -### TensorData format - -Tinker uses a typed tensor wire format for all numeric outputs. xorl uses the same format: - -```json -{ - "data": [0.423, 0.891, 0.312], - "dtype": "float32", - "shape": [3] -} -``` - -`LossFnOutput` responses always use `TensorData` values: - -```json -{ - "loss": {"data": [2.345], "dtype": "float32", "shape": [1]}, - "logprobs": {"data": [...], "dtype": "float32", "shape": [512]}, - "elementwise_loss": {"data": [...], "dtype": "float32", "shape": [512]} -} -``` - -`logprobs` and `elementwise_loss` are only present when `loss_fn: "per_token_ce"` is used. - ---- - -## RL Training Pattern +| Paper | Description | +|---|---| +| [Asynchronous RLHF: Faster and More Efficient Off-Policy RL for Language Models](https://arxiv.org/abs/2410.18252) (Noukhovitch et al., 2024) | Decouples generation and training into async processes. Shows that off-policy RL (training on stale rollouts) works well with proper importance correction — the motivation behind XoRL's TIS (Temporal Importance Sampling). | +| [INTELLECT-2: A Reasoning Model Trained Through Globally-Distributed Reinforcement Learning](https://arxiv.org/abs/2505.07291) (Primeintellect, 2025) | Distributed async RL training across globally distributed nodes. Demonstrates that RL training can scale across unreliable, heterogeneous clusters. | +| [OpenRLHF: An Easy-to-use, Scalable and High-performance RLHF Framework](https://arxiv.org/abs/2405.11143) (Hu et al., 2024) | Ray-based RLHF framework with vLLM integration. Demonstrates the separation of training and inference into distinct processes — the same architecture XoRL uses. | +| [verl: Volcano Engine Reinforcement Learning for LLMs](https://arxiv.org/abs/2409.19256) (Sheng et al., 2024) | Hybrid pipeline that colocates actor and rollout on the same GPUs with memory sharing. Shows how to minimize GPU idle time in the RL loop. | -```python -import xorl_client +### Techniques Supported by XoRL -service_client = xorl_client.ServiceClient(base_url="http://localhost:5555") -training_client = xorl_client.TrainingClient( - holder=service_client.holder, - model_id="default", - base_model="Qwen/Qwen3-8B", -) -adam_params = xorl_client.AdamParams(learning_rate=1e-5, beta1=0.9, beta2=0.95, eps=1e-8) - -for rl_step in range(num_rl_steps): - # 1. Rollout from inference (SGLang) - samples = sglang_client.generate(prompts, max_tokens=512) - - # 2. Score with reward model - rewards = reward_model.score(samples) - - # 3. Pack into Datum objects - data = [ - xorl_client.Datum( - model_input=xorl_client.ModelInput.from_ints(s.token_ids), - loss_fn_inputs={"labels": s.labels}, - ) - for s in samples - ] - - # 4. Gradient accumulation + optimizer step - fwd_bwd = training_client.forward_backward( - data, - loss_fn="importance_sampling", - loss_fn_params={"eps_clip": 0.2}, - ) - optim = training_client.optim_step(adam_params) - fwd_bwd.result() - optim.result() - - # 5. Sync weights to inference every N steps - if rl_step % sync_interval == 0: - service_client.sync_inference_weights(...) -``` +| Paper | XoRL feature | +|---|---| +| [GLM-5: Open Multilingual Multitask Model](https://arxiv.org/abs/2602.15763) (Team GLM, 2025) | **IcePop** — hard gradient masking for extreme importance ratios. Available via `icepop_beta` parameter in `policy_loss`. | +| [LoRA: Low-Rank Adaptation of LLMs](https://arxiv.org/abs/2106.09685) (Hu et al., 2021) | LoRA and multi-adapter support. XoRL supports multiple concurrent LoRA adapters switchable per request. | +| [QLoRA: Efficient Finetuning of Quantized LLMs](https://arxiv.org/abs/2305.14314) (Dettmers et al., 2023) | QLoRA training with nvfp4, block_fp8, and nf4 quantization formats. | --- -## Multi-Adapter (LoRA) Support - -The server supports multiple named LoRA adapters, switchable per request via `model_id`: - -```python -# Create two adapters -client_a = service_client.create_lora_training_client( - base_model="Qwen/Qwen3-8B", rank=32, model_id="policy" -) -client_b = service_client.create_lora_training_client( - base_model="Qwen/Qwen3-8B", rank=16, model_id="reference" -) - -# Train policy, forward-only on reference -fwd_bwd = client_a.forward_backward(data, loss_fn="policy_loss") -ref_logprobs = client_b.forward(data, loss_fn="per_token_ce") -``` - -LoRA adapters can be saved and loaded via the standard weight endpoints: +## Further Reading -```http -POST /api/v1/save_weights -POST /api/v1/load_weights -``` +| Topic | Page | +|---|---| +| Server architecture, multi-node, launcher CLI | [Launching & Configuration](/server-training/training-server/launching/) | +| REST API endpoints | [API Reference](/server-training/training-server/api-reference/) | +| xorl-sglang: weight sync, numerical alignment | [Inference: xorl-sglang](/server-training/sglang/) | +| Client SDK, loss functions, training patterns | [Client SDK (xorl-client)](/server-training/client-sdk/overview/) | +| NCCL weight sync protocol | [Weight Sync](/server-training/weight-sync/overview/) | +| SFT fine-tuning example | [SFT on No Robots](/server-training/examples/sft-no-robots/) | +| End-to-end weight sync test | [Password Memorization](/server-training/examples/password-memorization/) | diff --git a/docs/src/content/docs/server-training/rl-training.mdx b/docs/src/content/docs/server-training/rl-training.mdx deleted file mode 100644 index e0a50357..00000000 --- a/docs/src/content/docs/server-training/rl-training.mdx +++ /dev/null @@ -1,329 +0,0 @@ ---- -title: "RL Training" ---- - -xorl is designed from the ground up for online reinforcement learning on large language models. The server training API exposes every step of the training loop — forward pass, backward pass, optimizer step, weight sync — as individual HTTP endpoints, letting an external RL orchestrator drive training precisely when and how it needs to. - ---- - -## Choosing a Loss Function - -| Algorithm | `loss_fn` | Key params | When to use | -|---|---|---|---| -| SFT / continued pretraining | `causallm_loss` | — | Standard next-token prediction | -| PPO | `policy_loss` | `eps_clip=0.2` | Clipped policy gradient, most stable for RL | -| GRPO (simpler RL) | `importance_sampling` | — | No clipping, simpler but less stable | -| PPO + stale data correction | `policy_loss` | `use_tis=True, tis_clip_low=0.1` | Multiple epochs over same rollout | - -## Architecture - -A typical xorl RL setup runs three components in parallel: - - - - - - - - - - - - - - - xorl RL Training Architecture - - - RL Orchestrator - xorl-client SDK · HTTP - - - xorl Training Server - forward_backward - optim_step - FSDP2 · EP · PP · LoRA - PPO · GRPO · IcePop · TIS - - - xorl-sglang - Inference server - generate rollouts - R3 routing data · logprobs - weight sync via NCCL - - - forward_backward - optim_step - - - generate(prompts) - logprobs · routing - - - sync_inference_weights (NCCL) - - - Reward model, advantage computation, data packing — handled by the orchestrator - - -**Components:** - -| Component | Role | -|---|---| -| **RL Orchestrator** | Drives the RL loop: prompts → rollouts → rewards → advantages → training steps → weight sync. Built with `xorl-client`. | -| **xorl Training Server** | Exposes training as HTTP API; handles FSDP, EP, PP, all parallelism internally | -| **xorl-sglang** | Modified SGLang inference server; generates rollouts, exports per-token logprobs and MoE routing decisions | - ---- - -## Getting Started - -Install both client libraries: - -```bash -pip install git+https://github.com/togethercomputer/xorl-client.git -``` - -Start the training server: - -```bash -python -m xorl.server.launcher \ - --config server_config.yaml \ - --api-port 5555 -``` - -Start xorl-sglang (see the [xorl-sglang page](/server-training/sglang/) for details): - -```bash -python -m sglang.launch_server \ - --model-path Qwen/Qwen3-8B \ - --port 30000 -``` - -A minimal PPO training loop with `xorl_client`: - -```python -import xorl_client - -service = xorl_client.ServiceClient(base_url="http://localhost:5555") -training = xorl_client.TrainingClient( - holder=service.holder, - model_id="default", - base_model="Qwen/Qwen3-8B", -) -adam = xorl_client.AdamParams(learning_rate=1e-6, beta1=0.9, beta2=0.95, eps=1e-8) - -for step in range(num_steps): - # 1. Generate rollouts from inference - completions = sglang_client.generate(prompts, max_new_tokens=512, return_logprobs=True) - - # 2. Score with reward model - rewards = reward_model.score(completions) - - # 3. Compute per-token advantages (GAE or simple returns) - advantages = compute_advantages(rewards, completions) - - # 4. Pack into Datum objects with RL fields - data = [ - xorl_client.Datum( - model_input=xorl_client.ModelInput.from_ints(c.token_ids), - loss_fn_inputs={ - "labels": c.token_ids, - "logprobs": c.logprobs, # old policy logprobs from rollout - "advantages": advantages[i], - }, - ) - for i, c in enumerate(completions) - ] - - # 5. Training step (PPO loss) - fwd = training.forward_backward( - data, - loss_fn="policy_loss", - loss_fn_params={"eps_clip": 0.2, "compute_kl_stats": True}, - ) - opt = training.optim_step(adam) - result = fwd.result() - opt.result() - - # 6. Sync weights to inference every N steps - if step % sync_every == 0: - service.sync_inference_weights( - master_address=TRAINING_HOST, - master_port=29600, - ) -``` - ---- - -## RL Features Implemented - -### PPO Policy Loss - -Full PPO-style clipped policy gradient loss ([`src/xorl/ops/loss/policy_loss.py`](https://github.com/togethercomputer/xorl/blob/main/src/xorl/ops/loss/policy_loss.py)): - -``` -ratio = exp(new_logprobs - old_logprobs) -pg_loss = max(ratio × A, clip(ratio, 1-ε, 1+ε_high) × A) -``` - -**Loss function parameters:** - -| Parameter | Default | Description | -|---|---|---| -| `eps_clip` | `0.2` | Lower clip ratio for PPO | -| `eps_clip_high` | `0.2` | Upper clip ratio (can differ from lower for asymmetric clipping) | -| `eps_clip_c` | `null` | Dual-clip ratio for negative advantages (set to e.g. `3.0` to prevent large negative updates) | -| `compute_kl_stats` | `false` | Return KL statistics in metrics (K3 estimator, entropy, ratio stats) | - -**Metrics returned:** - -| Metric | Description | -|---|---| -| `pg_clipfrac` | Fraction of tokens where gradient clipping was applied | -| `kl_sample_train_k3` | Schulman's K3 KL estimator: `mean(exp(log_ratio) - log_ratio - 1)` — non-negative, unbiased | -| `entropy_sample` | Mean entropy of old policy: `-mean(old_logprobs)` | -| `ratio_mean/min/max` | Importance sampling ratio statistics | - -```python -fwd = training.forward_backward(data, loss_fn="policy_loss", loss_fn_params={ - "eps_clip": 0.2, - "eps_clip_high": 0.2, - "eps_clip_c": 3.0, # dual-clip negative advantages - "compute_kl_stats": True, -}) -``` - ---- - -### GRPO / Importance Sampling Loss - -Simpler importance-sampling loss for GRPO-style training ([`src/xorl/ops/loss/importance_sampling_loss.py`](https://github.com/togethercomputer/xorl/blob/main/src/xorl/ops/loss/importance_sampling_loss.py)): - -``` -ratio = exp(new_logprobs - old_logprobs) -loss = -(ratio × advantages).mean() -``` - -No clipping — relies on advantages being bounded. Suitable when the policy doesn't drift far from the rollout policy. - -```python -fwd = training.forward_backward(data, loss_fn="importance_sampling", loss_fn_params={ - "compute_kl_stats": True, -}) -``` - ---- - -### IcePop - -IcePop (from [GLM-5, arXiv:2602.15763](https://arxiv.org/abs/2602.15763)) is a **hard masking** technique that zeros gradients for tokens where the importance sampling ratio falls outside the band `[1/β, β]`. This prevents large policy updates on tokens where the current and old policies have already diverged significantly — complementary to PPO's soft clipping. - - - - IcePop: Hard Masking by Importance Sampling Ratio - - - - - - ratio < 1/β - gradient ZEROED - - - 1/β ≤ ratio ≤ β - gradient ACTIVE (PPO applies) - - - ratio > β - gradient ZEROED - - ratio = 1/β - ratio = β - icepop_maskfrac metric: fraction of valid tokens that were zeroed - - -```python -fwd = training.forward_backward(data, loss_fn="policy_loss", loss_fn_params={ - "eps_clip": 0.2, - "icepop_beta": 5.0, # zero gradients when ratio < 0.2 or ratio > 5.0 -}) -``` - -IcePop and PPO clipping are complementary: PPO softly clips the loss value, IcePop hard-zeros the gradient for extreme ratios. - ---- - -### TIS — Temporal Importance Sampling - -When multiple training steps run on the same rollout batch (e.g. multiple epochs over data collected at step T), the policy drifts away from the rollout policy. TIS corrects for this by weighting each token's gradient by `exp(train_logprobs - rollout_logprobs)`: - -``` -tis_weight = clip(exp(train_logprobs - rollout_logprobs), tis_clip_low, tis_clip_high) -loss = (tis_weight × pg_loss).mean() -``` - -Requires passing `rollout_logprobs` separately from `logprobs` (old policy at training time): - -```python -fwd = training.forward_backward(data, loss_fn="policy_loss", loss_fn_params={ - "use_tis": True, - "tis_clip_low": 0.1, - "tis_clip_high": 2.0, -}) -``` - -Pass rollout logprobs as a separate field in `loss_fn_inputs`: - -```python -datum = xorl_client.Datum( - model_input=xorl_client.ModelInput.from_ints(token_ids), - loss_fn_inputs={ - "labels": token_ids, - "logprobs": logprobs_at_last_train_step, # old policy: updated each train step - "rollout_logprobs": logprobs_at_rollout, # reference: fixed from inference - "advantages": advantages, - }, -) -``` - ---- - -### R3 — Routing Replay for MoE - -For MoE models, gradient checkpointing causes a problem: flash attention is non-deterministic on recompute, which produces different hidden states → different top-K routing → AllToAll shape mismatches during backward. R3 (Rollout Routing Replay) solves this by recording expert routing decisions and replaying them verbatim during backward recompute. - -In RL training, R3 goes further: routing decisions from the inference rollout (xorl-sglang) can be pre-loaded into the training server so that the training policy routes tokens to **exactly the same experts as inference did**. This gives gradient-level consistency between rollout and training steps. - -Pass routing data from xorl-sglang in the `forward_backward` request: - -```python -datum = xorl_client.Datum( - model_input=xorl_client.ModelInput.from_ints(token_ids), - loss_fn_inputs={ - "labels": token_ids, - "logprobs": logprobs, - "advantages": advantages, - }, -) - -fwd = training.forward_backward( - [datum], - loss_fn="policy_loss", - routed_experts=rollout_routing_indices, # [T, L, K] from sglang - routed_expert_logits=rollout_routing_weights, # [T, L, K] softmax weights -) -``` - -See the [Router page](/moe/router/#routing-replay-r3) for full R3 details. - ---- - -## In This Section - -| Page | What it covers | -|---|---| -| [Server Training](/server-training/overview/) | REST API reference, xorl-client SDK usage, multi-node launch | -| [Inference: xorl-sglang](/server-training/sglang/) | Modified SGLang backend, what was changed, integration with xorl | -| [Loss Functions](/loss-functions/) | causallm_loss, policy_loss, importance_sampling, per_token_ce | diff --git a/docs/src/content/docs/server-training/sglang.mdx b/docs/src/content/docs/server-training/sglang.mdx index 2992bdbd..81993257 100644 --- a/docs/src/content/docs/server-training/sglang.mdx +++ b/docs/src/content/docs/server-training/sglang.mdx @@ -22,31 +22,34 @@ The core RL training loop requires capabilities that upstream SGLang does not pr ## What Was Modified -### Weight Update API +xorl-sglang adds ~2,500 lines across 32 files on top of upstream SGLang. The changes fall into five categories. The `xorl-rl-target` branch in the repo contains all xorl-specific commits. -xorl-sglang adds three HTTP endpoints that implement the NCCL-based weight synchronization protocol: +### 1. Weight Update Endpoints and Protocol + +The largest change (~1,700 lines in `model_runner.py`, `scheduler_update_weights_mixin.py`, `tokenizer_communicator_mixin.py`, `http_server.py`). xorl-sglang adds several HTTP endpoints for NCCL-based weight synchronization: | Endpoint | Description | |---|---| -| `POST /init_weights_update_group` | Join an NCCL process group for weight sync. Called by the xorl training server before each sync. | +| `POST /init_weights_update_group` | Join an NCCL process group for weight sync. Forces eager NCCL communicator creation with `device_id` and `NCCL_CUMEM_ENABLE=0` to match xorl's training side. | | `POST /update_weights_from_distributed` | Receive broadcasted weight tensors via NCCL `dist.broadcast`. Called per weight bucket. | +| `POST /prepare_weights_update` | Phase 1 of two-phase protocol: starts background recv threads ready to receive NCCL broadcasts. | +| `POST /complete_weights_update` | Phase 2: waits for recv threads to complete and applies the received weights. | +| `POST /receive_weights` | Single-call weight receive via NCCL broadcast (pause/broadcast/resume protocol). | +| `POST /receive_weights_ep_scatter` | Receive expert weights from multiple EP training ranks via NCCL P2P scatter. Receives directly into `param.data` slices (zero-copy). | +| `POST /list_weights` | List all weight names/shapes — used to verify weight name mapping between training and inference. | | `POST /destroy_weights_update_group` | Tear down the NCCL group after sync completes. | -The training server's `nccl_broadcast` backend drives these calls — see [Backend: nccl_broadcast](/server-training/weight-sync/nccl-broadcast/) for the full protocol. - -### Per-Token Logprob Export +Key implementation details: +- **Eager NCCL init**: Both sides must use `device_id` to force eager communicator creation. Without this, sglang uses lazy init and xorl's rank 0 hangs waiting for peers. +- **Two-phase protocol**: `prepare_weights_update` starts background recv threads, then `complete_weights_update` applies them after the training-side broadcast finishes. This avoids blocking the scheduler. +- **EP scatter**: For MoE models with expert parallelism, each EP training rank sends its local experts directly into the corresponding slice of `param.data` via P2P ops — no intermediate buffers. +- **Health check bypass**: During weight updates, the `/health` endpoint returns 200 immediately instead of running a test generation, avoiding timeouts during NCCL operations. -xorl-sglang returns per-token log probabilities for every generated token when `return_logprobs=True`. These are passed directly to the xorl training server as `logprobs` in the `loss_fn_inputs` field of each `Datum`. +The training server's `nccl_broadcast` backend drives these calls — see [Backend: nccl_broadcast](/server-training/weight-sync/nccl-broadcast/) for the full protocol. -```python -completions = sglang_client.generate( - prompts, - sampling_params={"max_new_tokens": 512, "return_logprobs": True}, -) -logprobs = completions[i].meta_info["logprobs"] # list of per-token log probs -``` +### 2. MoE Routing Data Export (R3) -### MoE Routing Data Export (R3) +Modified files: `routed_experts_capturer.py`, `qwen3_moe.py`, `qwen2_moe.py`, `io_struct.py`, `detokenizer_manager.py`, `scheduler_output_processor_mixin.py` For MoE models, xorl-sglang records which experts each token was routed to during generation and returns this alongside completions: @@ -58,6 +61,8 @@ routed_experts = completions[i].meta_info["routed_experts"] The training server decodes and replays these routing decisions via R3, ensuring training uses the same expert assignments as inference did. This is critical for gradient consistency in GRPO/PPO training with MoE models. +**Expert routing weights export** (`--enable-return-expert-logits`): In addition to expert indices, xorl-sglang can also capture and return `topk_weights` (the softmax routing probabilities). The `RoutedExpertsCapturer` allocates a separate `topk_weights_buffer` in float32 to record these per layer. + **Format:** xorl-sglang encodes routing data as base64 int32 arrays with shape metadata to minimize transfer overhead: ```json { @@ -68,9 +73,16 @@ The training server decodes and replays these routing decisions via R3, ensuring The `RoutingReplayHandler` in xorl decodes and distributes this data across context-parallel and packing dimensions automatically. -### Numerical Alignment +### 3. Numerical Alignment + +Modified files: `qwen3_moe.py`, `qwen3.py`, `qwen2_moe.py`, `qwen2.py`, `layernorm.py`, `server_args.py` + +xorl and xorl-sglang must produce identical logits for the same input. Mismatched logits lead to incorrect importance sampling ratios and degraded RL performance. xorl-sglang adds: -xorl and xorl-sglang use the same numerical alignment flags to ensure training and inference produce identical logits for the same input. Mismatched logits lead to incorrect importance sampling ratios and degraded RL performance. +- **`--enable-fp32-router`**: MoE router gate computation in FP32. In `Qwen3MoeSparseMoeBlock`, when `rl_on_policy_target` is set, the router uses `F.linear` in float32 with explicit `F.softmax` + `torch.topk` instead of the fused TopK kernel. +- **`--enable-fp32-lm-head`**: LM head logits in FP32 (already upstream, but used together with the router flag). +- **RMSNorm alignment**: When `rl_on_policy_target` is set, Q/K norms use `cast_x_before_out_mul=True` and `fp32_residual=False` to match xorl's training-side norm behavior. +- **Fused kernel bypass**: Disables `fused_qk_norm_rope` and `fused_kv_buffer` when `rl_on_policy_target` is set, since fused kernels can produce slightly different numerics. Configure both sides identically: @@ -84,46 +96,106 @@ rope_native: false attention_cast_bf16: false ``` -```python +```bash # xorl-sglang launch python -m sglang.launch_server \ --model-path Qwen/Qwen3-8B \ - --router-fp32 \ - --lm-head-fp32 \ + --enable-fp32-router \ + --enable-fp32-lm-head \ + --rl-on-policy-target xorl \ --port 30000 ``` +### 4. Batch-Invariant Mode + +Modified files: `model_runner.py`, `batch_invariant_ops.py` + +When `--rl-on-policy-target xorl-batch-invariant` is set, xorl-sglang enables batch-invariant mode. This ensures that the model produces identical outputs regardless of how requests are batched together — a requirement for correct importance sampling in RL. Without this, different batching of the same request can produce different logits due to padding interactions in attention and MoE kernels. + +### 5. Bug Fixes + +Two additional fixes on top of the upstream merge: + +- **`req_to_token_pool` slot leak** (`schedule_batch.py`, `scheduler.py`): When a `max_new_tokens=0` (prefill-only) request arrives during an idle window, its `ScheduleBatch` gets `is_prefill_only=True`. If normal generation requests are later merged in, `merge_batch()` never cleared this flag, so `get_next_batch_to_run()` skipped the decode path. Requests allocated pool slots during prefill but never decoded, never finished, and never freed their slots — exhausting the pool. Fixed by clearing `is_prefill_only` on merge and recomputing it from actual request state. + ### KV Cache Flush on Weight Update After a weight sync, xorl-sglang flushes its KV cache. The training server sends `flush_cache=True` on the last weight bucket to trigger this. Without flushing, cached key/value tensors from the old policy weights would be reused with new weights, producing incorrect logits. --- -## Launching xorl-sglang +## Installation + +xorl-sglang is included as a git submodule under `submodules/xorl-sglang`. If you cloned with `--recurse-submodules`, it's already checked out. + +```bash +pip install -e "submodules/xorl-sglang/python[all]" +``` + +Or use `pyproject.sglang.toml` to install xorl, xorl-client, and xorl-sglang together (pins PyTorch to 2.9.1): ```bash -# Clone and install -git clone https://github.com/togethercomputer/xorl-sglang -cd xorl-sglang -pip install -e ".[all]" +cp pyproject.sglang.toml pyproject.toml +uv sync # or: pip install -e . +``` + +See the [installation guide](/getting-started/installation/#install-submodules) for full details. + +## Launching xorl-sglang + +### Single GPU (Qwen3-8B FP8) -# Launch (single GPU) +```bash python -m sglang.launch_server \ - --model-path Qwen/Qwen3-8B \ + --model-path Qwen/Qwen3-8B-FP8 \ + --port 30000 \ + --rl-on-policy-target xorl \ + --enable-fp32-router \ + --enable-fp32-lm-head \ + --mem-fraction-static 0.88 +``` + +### Tensor Parallel (Qwen3-30B FP8, 2 GPUs) + +```bash +CUDA_VISIBLE_DEVICES=4,5 python -m sglang.launch_server \ + --model-path Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 \ --port 30000 \ - --router-fp32 \ - --lm-head-fp32 + --tp-size 2 \ + --rl-on-policy-target xorl \ + --enable-fp32-router \ + --enable-fp32-lm-head \ + --mem-fraction-static 0.88 +``` -# Launch (tensor parallel, 4 GPUs) +### Tensor Parallel (Qwen3-235B FP8, 4 GPUs, remote node) + +```bash python -m sglang.launch_server \ - --model-path Qwen/Qwen3-8B \ + --model-path Qwen/Qwen3-235B-A22B-Instruct-2507-FP8 \ --port 30000 \ + --host 0.0.0.0 \ --tp-size 4 \ - --router-fp32 \ - --lm-head-fp32 + --rl-on-policy-target xorl \ + --enable-fp32-router \ + --enable-fp32-lm-head \ + --mem-fraction-static 0.88 ``` -Wait for the server to be ready: +### Key launch flags + +| Flag | Description | +|---|---| +| `--rl-on-policy-target xorl` | Enable xorl weight sync endpoints and routing capture. Use `xorl-batch-invariant` for batch-invariant mode. | +| `--enable-fp32-router` | MoE router gate in FP32 (must match training config `router_fp32: true`) | +| `--enable-fp32-lm-head` | LM head logits in FP32 (must match training config `lm_head_fp32: true`) | +| `--enable-return-routed-experts` | Return expert routing indices in response metadata (for R3) | +| `--enable-return-expert-logits` | Return expert routing weights alongside indices | +| `--tp-size N` | Tensor parallelism across N GPUs | +| `--mem-fraction-static 0.88` | Fraction of GPU memory for KV cache (leave headroom for weight sync buffers) | + +### Wait for ready + ```python import requests, time while True: @@ -145,9 +217,10 @@ Before weight sync can happen, register the xorl-sglang instance with the traini ```python import requests -requests.post("http://training-server:5555/add_inference_endpoint", json={ +requests.post("http://training-server:6000/add_inference_endpoint", json={ "host": "inference-node-01", "port": 30000, + "worker_port": 30000, "world_size": 4, # match --tp-size }) ``` @@ -156,8 +229,11 @@ Multiple replicas can be registered — weight sync broadcasts to all of them in ```python for host, port, tp_size in inference_replicas: - requests.post("http://training-server:5555/add_inference_endpoint", json={ - "host": host, "port": port, "world_size": tp_size, + requests.post("http://training-server:6000/add_inference_endpoint", json={ + "host": host, + "port": port, + "worker_port": port, + "world_size": tp_size, }) ``` @@ -177,7 +253,7 @@ for _ in range(n_train_steps): training.optim_step(...) # Sync new weights and resume inference -service.sync_inference_weights(master_address=TRAIN_HOST, master_port=29600) +training.sync_inference_weights(master_address=TRAIN_HOST, master_port=29600).result() requests.post("http://inference-node:30000/wake_up") ``` @@ -185,13 +261,17 @@ requests.post("http://inference-node:30000/wake_up") ## Upstream Compatibility -xorl-sglang tracks upstream SGLang closely and aims to rebase regularly. The modifications are minimal and isolated: -- Weight update endpoints (new files) -- Routing data export (additions to existing generation path) -- Numerical alignment CLI flags (additions to launch args) -- Per-token logprob format compatibility with xorl's `LossFnOutput` +xorl-sglang tracks upstream SGLang closely and aims to rebase regularly. The modifications touch 32 files but are concentrated in a few areas: + +| Area | Files modified | Nature of change | +|---|---|---| +| Weight sync protocol | `model_runner.py`, `scheduler_update_weights_mixin.py`, `tokenizer_communicator_mixin.py`, `http_server.py` | New endpoints and NCCL recv logic (~1,700 lines added) | +| Routing data export | `routed_experts_capturer.py`, `io_struct.py`, `detokenizer_manager.py`, `scheduler_output_processor_mixin.py` | Additions to existing generation path | +| Numerical alignment | `qwen3_moe.py`, `qwen3.py`, `qwen2_moe.py`, `qwen2.py`, `layernorm.py` | Conditional paths gated on `rl_on_policy_target` | +| CLI args | `server_args.py` | New flags: `--enable-fp32-router`, `--enable-return-expert-logits`, `--enable-rdma-weight-updates`, `--rl-on-policy-target xorl` | +| Bug fixes | `schedule_batch.py`, `scheduler.py` | Prefill-only slot leak fix | -If you need a feature from a newer upstream SGLang release, file an issue at the xorl-sglang repo. +The `main` branch is periodically rebased onto upstream SGLang. The `xorl-rl-target` branch contains the xorl-specific commits on top. If you need a feature from a newer upstream SGLang release, file an issue at the xorl-sglang repo. ## Source diff --git a/docs/src/content/docs/server-training/api-reference.md b/docs/src/content/docs/server-training/training-server/api-reference.md similarity index 86% rename from docs/src/content/docs/server-training/api-reference.md rename to docs/src/content/docs/server-training/training-server/api-reference.md index 2a0fecd1..d683a9df 100644 --- a/docs/src/content/docs/server-training/api-reference.md +++ b/docs/src/content/docs/server-training/training-server/api-reference.md @@ -9,7 +9,7 @@ Poll `POST /api/v1/retrieve_future` with that ID to get the actual result. The `xorl-client` SDK handles polling automatically. ::: -All endpoints are served at `http://:/`. Training operations use a two-phase async protocol — see [Server Architecture](/server-training/architecture/#api-server) for details. +All endpoints are served at `http://:/`. Training operations use a two-phase async protocol — see [Launching & Configuration](/server-training/training-server/launching/#api-server) for details. ## Training Operations @@ -28,8 +28,8 @@ All endpoints are served at `http://:/`. Training operations use a t | `POST` | `/api/v1/unload_model` | Unload a session, freeing associated adapter state. | | `POST` | `/api/v1/kill_session` | Kill an active session; optionally reload weights from checkpoint. | | `GET` | `/api/v1/session_info` | List active sessions and their state. | -| `POST` | `/api/v1/create_session` | Tinker SDK compatibility alias for `create_model`. | -| `POST` | `/api/v1/session_heartbeat` | Tinker SDK keepalive. | +| `POST` | `/api/v1/create_session` | Tinker compatibility stub. Returns a `session_id` but does not register a model session. | +| `POST` | `/api/v1/session_heartbeat` | Tinker compatibility stub; currently a no-op. | ## Checkpointing @@ -40,7 +40,7 @@ All endpoints are served at `http://:/`. Training operations use a t | `POST` | `/api/v1/list_checkpoints` | List available checkpoints under `output_dir`. | | `POST` | `/api/v1/delete_checkpoint` | Delete a checkpoint by ID. | | `POST` | `/api/v1/weights_info` | Return checkpoint metadata for a model (used by xorl-client to load weights). | -| `POST` | `/api/v1/save_weights_for_sampler` | Save sampler-format weights for inference. | +| `POST` | `/api/v1/save_weights_for_sampler` | Save inference weights under `sampler_weights/` (LoRA adapter or full HF checkpoint, depending on training mode). | | `GET` | `/api/v1/training_runs` | List training runs. | ## Inference Integration diff --git a/docs/src/content/docs/server-training/architecture.mdx b/docs/src/content/docs/server-training/training-server/launching.mdx similarity index 99% rename from docs/src/content/docs/server-training/architecture.mdx rename to docs/src/content/docs/server-training/training-server/launching.mdx index 1b4bdefe..e024679a 100644 --- a/docs/src/content/docs/server-training/architecture.mdx +++ b/docs/src/content/docs/server-training/training-server/launching.mdx @@ -1,5 +1,5 @@ --- -title: "Server Architecture" +title: "Launching & Configuration" --- The xorl training server is built as a layered system of independent processes that communicate over ZMQ. This page covers how those layers fit together, what each component does, and the full API surface. diff --git a/docs/src/content/docs/server-training/weight-sync/overview.mdx b/docs/src/content/docs/server-training/weight-sync/overview.mdx index ae1eed14..2c7282a3 100644 --- a/docs/src/content/docs/server-training/weight-sync/overview.mdx +++ b/docs/src/content/docs/server-training/weight-sync/overview.mdx @@ -85,21 +85,25 @@ xorl's training server broadcasts weights to registered inference endpoints usin Before syncing, register the inference server's address with xorl: ```http -POST /api/v1/add_inference_endpoint +POST /add_inference_endpoint ``` ```json -{"host": "inference-node-01", "port": 30000, "world_size": 8} +{"host": "inference-node-01", "port": 30000, "worker_port": 30000, "world_size": 8} ``` Multiple endpoints can be registered (e.g. for multi-replica inference): ```python for host, port in inference_servers: - requests.post(f"{base_url}/api/v1/add_inference_endpoint", json={ - "host": host, "port": port, "world_size": 8 + requests.post(f"{base_url}/add_inference_endpoint", json={ + "host": host, + "port": port, + "worker_port": port, + "world_size": 8, }) ``` Endpoints are deduplicated by `(host, port)` — registering the same endpoint twice is safe. +For the common same-port xorl-sglang setup, set `worker_port` equal to `port`. ## Triggering a Sync @@ -131,11 +135,11 @@ The sync is synchronous from the caller's perspective — the endpoint returns o ```python # Every N RL steps, sync weights to inference if rl_step % sync_every == 0: - training_client.sync_weights( + training_client.sync_inference_weights( master_address=TRAINING_HEAD_IP, master_port=29600, group_name="policy_sync", - ) + ).result() # Now SGLang is serving the latest policy weights ``` @@ -172,7 +176,7 @@ For QLoRA, the sync dequantizes and merges: ## Removing Inference Endpoints ```http -POST /api/v1/remove_inference_endpoint +POST /remove_inference_endpoint {"host": "inference-node-01", "port": 30000} ``` @@ -190,7 +194,7 @@ for _ in range(n_steps): training_client.optim_step(...) # Sync and resume inference -training_client.sync_weights(...) +training_client.sync_inference_weights(...).result() requests.post(f"{inference_url}/wake_up") ``` diff --git a/examples/server/no_robot_sft/README.md b/examples/server/no_robot_sft/README.md index 1eb3aa19..32cdf51f 100644 --- a/examples/server/no_robot_sft/README.md +++ b/examples/server/no_robot_sft/README.md @@ -17,7 +17,7 @@ export CUDA_VISIBLE_DEVICES=0,1,2,3 python -m xorl.server.launcher \ --mode auto \ - --config ../configs/qwen3_4b_instruct_2507.yaml \ + --config ../configs/lora/qwen3_8b_lora.yaml \ --api-port 6000 \ --log-level DEBUG ``` @@ -32,10 +32,13 @@ The launcher will: ### SFT (Supervised Fine-Tuning) -[`example_sft.py`](example_sft.py) -- Minimal SFT training loop on the [No Robots](https://huggingface.co/datasets/HuggingFaceH4/no_robots) dataset. +[`run_sft.py`](run_sft.py) -- Minimal LoRA SFT training loop on the [No Robots](https://huggingface.co/datasets/HuggingFaceH4/no_robots) dataset. ```bash -python example_sft.py --config.base_url http://localhost:6000 +python run_sft.py \ + --config.base_url http://localhost:6000 \ + --config.model_name Qwen/Qwen3-8B \ + --config.lora_rank 32 ``` Key features: @@ -45,6 +48,9 @@ Key features: - Periodic checkpoint saving and resume support - Per-token NLL metrics via `compute_mean_nll` +The checked-in server config above loads `Qwen/Qwen3-8B` with LoRA rank 32. +`run_sft.py` still defaults to an older 4B example, so the command overrides those defaults explicitly. + **Config options** (passed via `--config.`): | Field | Default | Description | @@ -59,43 +65,45 @@ Key features: ## Server Configs -### [`configs/qwen3_4b_instruct_2507.yaml`](../configs/qwen3_4b_instruct_2507.yaml) +### [`configs/lora/qwen3_8b_lora.yaml`](../configs/lora/qwen3_8b_lora.yaml) -Qwen3-4B dense model on 4 GPUs with Ulysses sequence parallelism (SP=4). +Qwen3-8B dense model on 4 GPUs with LoRA enabled. -### [`configs/qwen3_coder_30b_a3b.yaml`](../configs/qwen3_coder_30b_a3b.yaml) +### [`configs/lora/qwen3_coder_30b_a3b_lora.yaml`](../configs/lora/qwen3_coder_30b_a3b_lora.yaml) -Qwen3-Coder-30B-A3B MoE model on 4 GPUs with Ulysses SP=4 and Flash Attention 3. +Qwen3-Coder-30B-A3B MoE model on 4 GPUs with LoRA enabled. ### Config Reference ```yaml # Model -model_path: Qwen/Qwen3-4B-Instruct-2507 -attn_implementation: flash_attention_3 # eager, sdpa, flash_attention_3, flash_attention_4 +model_path: Qwen/Qwen3-8B +tokenizer_path: Qwen/Qwen3-8B +attn_implementation: flash_attention_3 # Parallelism (world_size = dp_rep * dp_shard * ulysses * cp) data_parallel_mode: fsdp2 # ddp, fsdp, fsdp2 -ulysses_parallel_size: 4 # Ulysses sequence parallelism +ulysses_parallel_size: 1 # Ulysses sequence parallelism +data_parallel_replicate_size: 1 +data_parallel_shard_size: 4 # Memory & Performance ce_mode: compiled # eager, compiled enable_mixed_precision: true enable_gradient_checkpointing: true enable_full_shard: true +enable_activation_offload: false +init_device: meta # Data Processing -sample_packing_sequence_len: 128000 # Max packed sequence length +sample_packing_sequence_len: 128000 # Max packed sequence length enable_packing: true # Pack multiple samples per micro-batch # LoRA enable_lora: true lora_rank: 32 lora_alpha: 32 -lora_target_modules: ["qkv_proj", "o_proj", "gate_up_proj", "down_proj"] - -# Workers -worker_bind_address: tcp://127.0.0.1:5556 +lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] ``` ## Training Loop Pattern @@ -108,8 +116,8 @@ import xorl_client # Connect to server service = xorl_client.ServiceClient(base_url="http://localhost:6000") client = service.create_lora_training_client( - base_model="Qwen/Qwen3-4B-Instruct-2507", - rank=64, + base_model="Qwen/Qwen3-8B", + rank=32, model_id="my-run", ) diff --git a/pyproject.sglang.toml b/pyproject.sglang.toml new file mode 100644 index 00000000..f9cd32d3 --- /dev/null +++ b/pyproject.sglang.toml @@ -0,0 +1,123 @@ +# Alternative pyproject.toml that pins PyTorch 2.9.1 so that xorl, xorl-client, +# and xorl-sglang can all be installed in the same environment. +# +# Usage: +# cp pyproject.sglang.toml pyproject.toml +# uv sync # (uv) +# pip install -e . # (conda/pip) + +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project.urls] +"Homepage" = "https://github.com/xorl-org/xorl" +"Source Code" = "https://github.com/xorl-org/xorl" + +[project] +name = "xorl" +dynamic = ["version"] +description = "Xorl: Simple and high-performance distributed training framework for large language models" +requires-python = "==3.12.*" + +dependencies = [ + "boto3>=1.35.0", + "datasets>=4.4.0", + "packaging>=23.0,<26.0", + "tenacity>=8.0.0", + "torchdata>=0.8.0,<1.0", + "transformers[torch]>=5.0", + "psutil", + "wandb", + "safetensors", + "einops", + "numpy<2.4", + "numba", + "pydantic", + "uvicorn", + "pyzmq", + "msgpack", + "fastapi", + "xorl-client @ git+https://github.com/togethercomputer/xorl-client.git", + # PyTorch 2.9.1 with CUDA 12.9 (compatible with xorl-sglang) + "torch @ https://download.pytorch.org/whl/cu129/torch-2.9.1%2Bcu129-cp312-cp312-manylinux_2_28_x86_64.whl", + "torchvision @ https://download.pytorch.org/whl/cu129/torchvision-0.24.1%2Bcu129-cp312-cp312-manylinux_2_28_x86_64.whl", + "triton @ https://download.pytorch.org/whl/triton-3.5.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", + # Flash Attention + "flash-attn-3 @ https://github.com/windreamer/flash-attention3-wheels/releases/download/2026.02.17-06dc5e7/flash_attn_3-3.0.0%2B20260216.cu129torch291cxx11abitrue.fec3a6-cp39-abi3-linux_x86_64.whl", + "flash-attn-cute @ git+https://github.com/Dao-AILab/flash-attention.git@5678dd909aca97f925957dab716022046fb1e44f#subdirectory=flash_attn/cute", + # Submodule + "sglang[all] @ file:submodules/xorl-sglang/python", +] + + +[dependency-groups] +# Follow the best practice in https://docs.astral.sh/uv/concepts/projects/dependencies/#development-dependencies +# to manage dev dependencies (i.e., dependencies that are only used in development) like +# test, lint and doc tools. +dev = [ + {include-group = "lint"}, + {include-group = "test"}, +] +lint = [ + "pre-commit", + "ruff", +] +test = [ + "pytest", + "expecttest" +] + + +[tool.uv] +# This locks the uv version so that we have a consistent uv behavior across the board. +# Inconsistent uv versions might generate different uv lock files which creates chaos. +# +# NOTE 1: Update this at least once per month as uv releases new version every week. +# NOTE 2: When updating this line, make sure to update Dockerfile under docker/ to the same +# version and release new docker images. +required-version = ">=0.8.14" + + +[tool.setuptools.dynamic] +version = {attr = "xorl.__version__"} + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.ruff] +target-version = "py312" +line-length = 120 + +[tool.ruff.lint] +ignore = ["C901", "E501", "E741", "W605", "C408"] +select = ["C", "E", "F", "I", "W"] + +[tool.ruff.lint.per-file-ignores] +"__init__.py" = ["E402", "F401", "F403", "F811"] + +[tool.ruff.lint.isort] +lines-after-imports = 2 +known-first-party = ["xorl"] +known-third-party = ["torch", "transformers", "wandb"] + + +[tool.pytest.ini_options] +addopts = "-v" +markers = [ + "cpu: CPU-only tests (no GPU required)", + "gpu: Tests that require GPU", + "e2e: End-to-end tests (full system, require GPU and torchrun)", + "server: Tests for the training server (API, orchestrator, runner)", + "distributed: Tests that require distributed setup (torchrun)", + "slow: Tests that take a long time to run", + "benchmark: Performance benchmark tests", + "collator: Tests for data collators", + "dataloader: Tests for data loaders", +] +testpaths = ["tests"] +norecursedirs = ["submodules"] +filterwarnings = [ + "ignore::DeprecationWarning:multiprocessing.popen_fork", + "ignore:The argument 'device' of Tensor:DeprecationWarning", +] diff --git a/pyproject.toml b/pyproject.toml index 0cd2f9c7..9ba9a30d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,7 @@ dependencies = [ "pyzmq", "msgpack", "fastapi", + "xorl-client @ git+https://github.com/togethercomputer/xorl-client.git", # PyTorch 2.10.0 with CUDA 12.9 "torch @ https://download.pytorch.org/whl/cu129/torch-2.10.0%2Bcu129-cp312-cp312-manylinux_2_28_x86_64.whl", "torchvision @ https://download.pytorch.org/whl/cu129/torchvision-0.25.0%2Bcu129-cp312-cp312-manylinux_2_28_x86_64.whl", diff --git a/submodules/xorl-client b/submodules/xorl-client new file mode 160000 index 00000000..2a3a60a7 --- /dev/null +++ b/submodules/xorl-client @@ -0,0 +1 @@ +Subproject commit 2a3a60a783c98e2a8ff722bad06dab18caee350c diff --git a/submodules/xorl-sglang b/submodules/xorl-sglang new file mode 160000 index 00000000..2b61bf30 --- /dev/null +++ b/submodules/xorl-sglang @@ -0,0 +1 @@ +Subproject commit 2b61bf30bcb9c51bfafb22a9b178d45d3876b0c6 From 6ebc5a53a96163f7504d8ec3c28a93539cd5229c Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Fri, 10 Apr 2026 12:41:01 -0700 Subject: [PATCH 24/41] Fix NCCL weight sync rendezvous port reuse (#70) * Fix nccl weight sync rendezvous reuse * Default weight sync rendezvous to ephemeral ports (cherry picked from commit 101334268440cc57b874c41c7a2c3b2f18ac28cf) --- src/xorl/server/api_server/api_types.py | 4 +- src/xorl/server/api_server/endpoints.py | 2 +- src/xorl/server/backend/base.py | 2 +- src/xorl/server/backend/dummy.py | 2 +- src/xorl/server/backend/remote.py | 2 +- src/xorl/server/protocol/operations.py | 2 +- src/xorl/server/weight_sync/README.md | 4 +- src/xorl/server/weight_sync/backends/base.py | 2 +- .../weight_sync/backends/nccl_broadcast.py | 97 +++++++---- .../test_nccl_broadcast_lifecycle.py | 153 ++++++++++++++++++ 10 files changed, 232 insertions(+), 38 deletions(-) create mode 100644 tests/server/weight_sync/test_nccl_broadcast_lifecycle.py diff --git a/src/xorl/server/api_server/api_types.py b/src/xorl/server/api_server/api_types.py index 7e7c85ff..98a8b75c 100644 --- a/src/xorl/server/api_server/api_types.py +++ b/src/xorl/server/api_server/api_types.py @@ -608,7 +608,7 @@ class AddInferenceEndpointRequest(BaseModel): master_address: Optional[str] = Field( default=None, description="Master address for NCCL rendezvous (auto-detected if not provided)" ) - master_port: int = Field(default=29600, description="Master port for NCCL rendezvous") + master_port: int = Field(default=0, description="Master port for NCCL rendezvous (0 selects an ephemeral port)") group_name: str = Field(default="weight_sync_group", description="Name of the NCCL process group") buffer_size_mb: int = Field(default=1024, description="Size of each transfer bucket in MB (to avoid OOM)") @@ -655,7 +655,7 @@ class SyncInferenceWeightsRequest(BaseModel): master_address: str = Field( default="", description="Master address for NCCL rendezvous (training server address). Auto-detected if empty." ) - master_port: int = Field(default=29600, description="Master port for NCCL rendezvous") + master_port: int = Field(default=0, description="Master port for NCCL rendezvous (0 selects an ephemeral port)") group_name: str = Field(default="weight_sync_group", description="Name of the NCCL process group") buffer_size_mb: int = Field(default=1024, description="Size of each transfer bucket in MB (to avoid OOM)") flush_cache: bool = Field( diff --git a/src/xorl/server/api_server/endpoints.py b/src/xorl/server/api_server/endpoints.py index b92ecd0c..6c88290e 100644 --- a/src/xorl/server/api_server/endpoints.py +++ b/src/xorl/server/api_server/endpoints.py @@ -645,7 +645,7 @@ async def add_inference_endpoint_endpoint(request: AddInferenceEndpointRequest, Parameters for auto-sync: - sync_weights: Whether to sync weights after adding (default: False) - master_address: Training server hostname for NCCL rendezvous (auto-detected if not provided) - - master_port: Port for NCCL rendezvous (default: 29600) + - master_port: Port for NCCL rendezvous (default: 0 for ephemeral) - group_name: NCCL process group name (default: weight_sync_group) - buffer_size_mb: Transfer bucket size in MB (default: 1024) """ diff --git a/src/xorl/server/backend/base.py b/src/xorl/server/backend/base.py index 45d94670..7f195f44 100644 --- a/src/xorl/server/backend/base.py +++ b/src/xorl/server/backend/base.py @@ -126,7 +126,7 @@ async def sync_inference_weights( self, endpoints: List[Dict[str, Any]], master_address: str = "localhost", - master_port: int = 29600, + master_port: int = 0, request_id: Optional[str] = None, **kwargs, ) -> Dict[str, Any]: diff --git a/src/xorl/server/backend/dummy.py b/src/xorl/server/backend/dummy.py index 6d33cc6c..ecbba254 100644 --- a/src/xorl/server/backend/dummy.py +++ b/src/xorl/server/backend/dummy.py @@ -96,7 +96,7 @@ async def wake_up(self, request_id=None): return {"status": "awake", "load_time": 0.0} async def sync_inference_weights( - self, endpoints, master_address="localhost", master_port=29600, request_id=None, **kwargs + self, endpoints, master_address="localhost", master_port=0, request_id=None, **kwargs ): return { "success": True, diff --git a/src/xorl/server/backend/remote.py b/src/xorl/server/backend/remote.py index 60c18083..7f4870b2 100644 --- a/src/xorl/server/backend/remote.py +++ b/src/xorl/server/backend/remote.py @@ -330,7 +330,7 @@ async def sync_inference_weights( self, endpoints, master_address="localhost", - master_port=29600, + master_port=0, group_name="weight_sync_group", buffer_size_mb=1024, sync_method="nccl_broadcast", diff --git a/src/xorl/server/protocol/operations.py b/src/xorl/server/protocol/operations.py index cc370a4a..42542e1f 100644 --- a/src/xorl/server/protocol/operations.py +++ b/src/xorl/server/protocol/operations.py @@ -99,7 +99,7 @@ class SyncWeightsData: endpoints: List[Dict[str, Any]] = field(default_factory=list) master_address: str = "localhost" - master_port: int = 29600 + master_port: int = 0 group_name: str = "weight_sync_group" buffer_size_mb: int = 1024 sync_method: str = "nccl_broadcast" diff --git a/src/xorl/server/weight_sync/README.md b/src/xorl/server/weight_sync/README.md index afdc20f2..2eb3f8ac 100644 --- a/src/xorl/server/weight_sync/README.md +++ b/src/xorl/server/weight_sync/README.md @@ -46,7 +46,7 @@ the command internally via Gloo so all ranks enter the sync together. ```python resp = requests.post("http://localhost:6000/api/v1/sync_inference_weights", json={ "master_address": "localhost", # training server address for NCCL rendezvous - "master_port": 29600, # default + "master_port": 0, # default; asks TCPStore to bind an ephemeral port "buffer_size_mb": 1024, # bucket size; reduce if OOM during sync "flush_cache": True, # flush KV cache after sync (default) "pause_mode": "retract", # "retract" | "abort" | "in_place" @@ -179,7 +179,7 @@ Populated by the handler from the sync request payload: class TransportConfig: endpoints: List[EndpointConfig] # host, port, world_size (TP size) master_address: str - master_port: int + master_port: int # 0 selects an ephemeral port on the training rank group_name: str buffer_size_mb: int device: str diff --git a/src/xorl/server/weight_sync/backends/base.py b/src/xorl/server/weight_sync/backends/base.py index e6f65a53..8d534f87 100644 --- a/src/xorl/server/weight_sync/backends/base.py +++ b/src/xorl/server/weight_sync/backends/base.py @@ -64,7 +64,7 @@ class TransportConfig: endpoints: List[EndpointConfig] = field(default_factory=list) master_address: str = "localhost" - master_port: int = 29600 + master_port: int = 0 group_name: str = "weight_sync_group" buffer_size_mb: int = 1024 device: str = "cuda:0" diff --git a/src/xorl/server/weight_sync/backends/nccl_broadcast.py b/src/xorl/server/weight_sync/backends/nccl_broadcast.py index 31f4c697..ec4f3f65 100644 --- a/src/xorl/server/weight_sync/backends/nccl_broadcast.py +++ b/src/xorl/server/weight_sync/backends/nccl_broadcast.py @@ -16,6 +16,7 @@ import os import time from concurrent.futures import ThreadPoolExecutor, as_completed +from contextlib import contextmanager from dataclasses import dataclass from threading import Thread from typing import Any, Dict, FrozenSet, List, Optional, Tuple @@ -94,7 +95,7 @@ def __init__( self, endpoints: List[EndpointInfo], master_address: str = "localhost", - master_port: int = 29600, + master_port: int = 0, group_name: str = "weight_sync_group", buffer_size_mb: int = 1024, device: str = "cuda:0", @@ -105,7 +106,7 @@ def __init__( Args: endpoints: List of inference endpoints to sync weights to master_address: Address for NCCL rendezvous (training server) - master_port: Port for NCCL rendezvous + master_port: Port for NCCL rendezvous (0 selects an ephemeral port) group_name: Name of the NCCL process group buffer_size_mb: Size of each transfer bucket in MB device: Device to use for NCCL operations @@ -122,6 +123,9 @@ def __init__( # Process group (initialized during sync) self.process_group: Optional[dist.ProcessGroup] = None + self._training_raw_store = None + self._training_prefix_store = None + self._active_master_port = master_port logger.info( f"NCCLWeightSynchronizer initialized: " @@ -133,6 +137,54 @@ def __init__( # NCCL process group management # ======================================================================== + @contextmanager + def _without_torchelastic_agent_store(self): + """Temporarily disable the elastic agent store override for custom groups.""" + old_agent_store = os.environ.pop("TORCHELASTIC_USE_AGENT_STORE", None) + try: + yield + finally: + if old_agent_store is not None: + os.environ["TORCHELASTIC_USE_AGENT_STORE"] = old_agent_store + + def _create_training_store(self) -> None: + """Create and retain the rendezvous store used by the training process group.""" + self._cleanup_training_store() + + requested_port = self.master_port + logger.info(f"[Training] Creating TCPStore (requested_port={requested_port}, is_master=True)...") + + with self._without_torchelastic_agent_store(): + raw_store = TCPStore( + host_name=self.master_address, + port=requested_port, + world_size=self.world_size, + is_master=True, + timeout=default_pg_timeout, + ) + + self._training_raw_store = raw_store + self._active_master_port = raw_store.port + self._training_prefix_store = PrefixStore(self.group_name, raw_store) + logger.info( + f"[Training] Rendezvous store ready: tcp://{self.master_address}:{self._active_master_port}, " + f"group_name={self.group_name}" + ) + + def _cleanup_training_store(self) -> None: + """Drop references to the rendezvous store so the TCP listener can be reclaimed.""" + raw_store = self._training_raw_store + self._training_prefix_store = None + self._training_raw_store = None + self._active_master_port = self.master_port + + close = getattr(raw_store, "close", None) + if callable(close): + try: + close() + except Exception as exc: + logger.warning(f"[Training] Failed to close rendezvous store cleanly: {exc}") + def _init_training_process_group(self) -> dist.ProcessGroup: """ Initialize NCCL process group on the training side (rank 0). @@ -147,15 +199,10 @@ def _init_training_process_group(self) -> dist.ProcessGroup: # Important: Only set NCCL_CUMEM_ENABLE=0 os.environ["NCCL_CUMEM_ENABLE"] = "0" - # torchrun sets TORCHELASTIC_USE_AGENT_STORE=True which forces all ranks - # to be TCPStore clients. We need to unset this for our separate weight - # sync process group where rank 0 must be the store master. - old_agent_store = os.environ.pop("TORCHELASTIC_USE_AGENT_STORE", None) - rank = 0 # Training is always rank 0 logger.info( - f"[Training] Initializing process group: tcp://{self.master_address}:{self.master_port}, " + f"[Training] Initializing process group: tcp://{self.master_address}:{self._active_master_port}, " f"rank={rank}, world_size={self.world_size}, group_name={self.group_name}" ) @@ -168,19 +215,9 @@ def _init_training_process_group(self) -> dist.ProcessGroup: backend_obj = Backend("nccl") timeout = default_pg_timeout - # Create TCPStore directly - rank 0 is the master - is_master = rank == 0 - logger.info(f"[Training] Creating TCPStore (is_master={is_master})...") - store = TCPStore( - host_name=self.master_address, - port=self.master_port, - world_size=self.world_size, - is_master=is_master, - timeout=timeout, - ) + if self._training_prefix_store is None: + self._create_training_store() - # Use PrefixStore with group_name to namespace keys - store = PrefixStore(self.group_name, store) logger.info("[Training] TCPStore created, creating process group...") # Handle different PyTorch versions by inspecting the actual signature @@ -195,22 +232,18 @@ def _init_training_process_group(self) -> dist.ProcessGroup: torch.cuda.set_device(self.device) device_id = torch.device(self.device) - try: + with self._without_torchelastic_agent_store(): pg, _ = _new_process_group_helper( self.world_size, rank, [], backend_obj, - store, + self._training_prefix_store, group_name=self.group_name, **{pg_options_param_name: None}, timeout=timeout, device_id=device_id, ) - finally: - # Restore the environment variable - if old_agent_store is not None: - os.environ["TORCHELASTIC_USE_AGENT_STORE"] = old_agent_store _world.pg_group_ranks[pg] = {i: i for i in range(self.world_size)} @@ -233,7 +266,6 @@ def init_nccl_group(self) -> bool: logger.info("=" * 70) logger.info("Initializing NCCL process groups...") logger.info("=" * 70) - logger.info(f"[Training] NCCL rendezvous: tcp://{self.master_address}:{self.master_port}") logger.info(f"[Training] World size: {self.world_size} (1 training + {self.world_size - 1} inference)") training_error = None @@ -253,6 +285,13 @@ def init_inference(): # 3. Run training (rank 0) in main thread - this completes NCCL rendezvous # 4. Join inference thread after training completes + try: + self._create_training_store() + except Exception as exc: + logger.error(f"[Training] Failed to create rendezvous store: {exc}") + return False + + logger.info(f"[Training] NCCL rendezvous: tcp://{self.master_address}:{self._active_master_port}") logger.info("[Training] Starting inference endpoint initialization in background...") inference_thread = Thread(target=init_inference) inference_thread.start() @@ -324,6 +363,8 @@ def destroy_nccl_group(self) -> None: logger.error(f"Failed to destroy training process group: {e}") self.process_group = None + self._cleanup_training_store() + # ======================================================================== # Inference endpoint management # ======================================================================== @@ -343,7 +384,7 @@ def init_single(rank_offset: int, endpoint: EndpointInfo) -> Dict[str, Any]: url = f"http://{endpoint.host}:{endpoint.port}/init_weights_update_group" payload = { "master_address": self.master_address, - "master_port": self.master_port, + "master_port": self._active_master_port, "rank_offset": rank_offset, "world_size": self.world_size, "group_name": self.group_name, diff --git a/tests/server/weight_sync/test_nccl_broadcast_lifecycle.py b/tests/server/weight_sync/test_nccl_broadcast_lifecycle.py new file mode 100644 index 00000000..93c8d727 --- /dev/null +++ b/tests/server/weight_sync/test_nccl_broadcast_lifecycle.py @@ -0,0 +1,153 @@ +"""Lifecycle regression tests for NCCL weight sync rendezvous.""" + +from types import SimpleNamespace + +import torch +import torch.distributed as dist + +from xorl.server.weight_sync.backends import nccl_broadcast as nccl_broadcast_module +from xorl.server.weight_sync.backends.nccl_broadcast import EndpointInfo, NCCLWeightSynchronizer + + +class _JoinDrivenThread: + """Test double that defers target execution until ``join()``.""" + + instances = [] + + def __init__(self, target): + self._target = target + self.join_called = False + _JoinDrivenThread.instances.append(self) + + def start(self): + return None + + def join(self): + self.join_called = True + self._target() + + +class _StickyTCPStore: + """Fake rendezvous store that keeps ports reserved after destroy.""" + + open_ports = set() + next_ephemeral_port = 31000 + + def __init__(self, host_name, port, world_size, is_master, timeout): + del host_name, world_size, is_master, timeout + if port == 0: + port = self.next_ephemeral_port + _StickyTCPStore.next_ephemeral_port += 1 + if port in self.open_ports: + raise RuntimeError(f"EADDRINUSE: {port}") + self.port = port + self.open_ports.add(port) + + def close(self): + return None + + +class _FakePrefixStore: + def __init__(self, prefix, store): + self.prefix = prefix + self.store = store + + +class _FakeProcessGroup: + def __init__(self, store): + self.store = store + self.destroyed = False + + +def _make_sync(master_port: int = 0) -> NCCLWeightSynchronizer: + return NCCLWeightSynchronizer( + endpoints=[EndpointInfo(host="127.0.0.1", port=12345, world_size=1)], + master_address="127.0.0.1", + master_port=master_port, + group_name="weight_sync_group", + device="cuda:0", + ) + + +def test_init_nccl_group_fails_before_starting_inference_when_store_bind_fails(monkeypatch): + sync = _make_sync() + inference_started = False + + _JoinDrivenThread.instances.clear() + monkeypatch.setattr(nccl_broadcast_module, "Thread", _JoinDrivenThread) + monkeypatch.setattr(sync, "_create_training_store", lambda: (_ for _ in ()).throw(RuntimeError("EADDRINUSE"))) + + def fake_inference_init(): + nonlocal inference_started + inference_started = True + return [] + + monkeypatch.setattr(sync, "_init_inference_endpoints", fake_inference_init) + + assert sync.init_nccl_group() is False + assert inference_started is False + assert _JoinDrivenThread.instances == [] + + +def test_destroy_nccl_group_reinit_uses_fresh_ephemeral_rendezvous_port_when_old_one_is_sticky(monkeypatch): + sync = _make_sync() + destroyed_groups = [] + + _StickyTCPStore.open_ports.clear() + _StickyTCPStore.next_ephemeral_port = 31000 + + def fake_new_process_group_helper(world_size, rank, group_ranks, backend, store, **kwargs): + del world_size, rank, group_ranks, backend, kwargs + return _FakeProcessGroup(store.store), None + + def fake_destroy_process_group(process_group): + process_group.destroyed = True + destroyed_groups.append(process_group) + + monkeypatch.setattr(nccl_broadcast_module, "TCPStore", _StickyTCPStore) + monkeypatch.setattr(nccl_broadcast_module, "PrefixStore", _FakePrefixStore) + monkeypatch.setattr(dist, "is_initialized", lambda: False) + monkeypatch.setattr(dist, "destroy_process_group", fake_destroy_process_group) + monkeypatch.setattr(torch.cuda, "set_device", lambda *_args, **_kwargs: None) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(nccl_broadcast_module, "Backend", lambda name: name) + monkeypatch.setattr(nccl_broadcast_module, "_new_process_group_helper", fake_new_process_group_helper) + monkeypatch.setattr(nccl_broadcast_module, "_world", SimpleNamespace(pg_group_ranks={})) + monkeypatch.setattr(nccl_broadcast_module, "default_pg_timeout", object()) + monkeypatch.setattr(sync, "_destroy_inference_endpoints", lambda: None) + + first_pg = sync._init_training_process_group() + sync.process_group = first_pg + first_port = sync._active_master_port + sync.destroy_nccl_group() + + assert destroyed_groups == [first_pg] + + sync.process_group = sync._init_training_process_group() + assert first_port == 31000 + assert sync._active_master_port == 31001 + + +def test_explicit_master_port_is_honored(monkeypatch): + sync = _make_sync(master_port=29600) + + _StickyTCPStore.open_ports.clear() + + def fake_new_process_group_helper(world_size, rank, group_ranks, backend, store, **kwargs): + del world_size, rank, group_ranks, backend, kwargs + return _FakeProcessGroup(store.store), None + + monkeypatch.setattr(nccl_broadcast_module, "TCPStore", _StickyTCPStore) + monkeypatch.setattr(nccl_broadcast_module, "PrefixStore", _FakePrefixStore) + monkeypatch.setattr(dist, "is_initialized", lambda: False) + monkeypatch.setattr(torch.cuda, "set_device", lambda *_args, **_kwargs: None) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(nccl_broadcast_module, "Backend", lambda name: name) + monkeypatch.setattr(nccl_broadcast_module, "_new_process_group_helper", fake_new_process_group_helper) + monkeypatch.setattr(nccl_broadcast_module, "_world", SimpleNamespace(pg_group_ranks={})) + monkeypatch.setattr(nccl_broadcast_module, "default_pg_timeout", object()) + + process_group = sync._init_training_process_group() + + assert process_group.store.port == 29600 + assert sync._active_master_port == 29600 From e10eb5613a7dab616c823f0b2a646b69fe28c033 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Fri, 10 Apr 2026 12:44:24 -0700 Subject: [PATCH 25/41] Make create_session return usable sessions (#111) (cherry picked from commit c49840b4e90ca895de3b224eeed0654e9a6c80a6) --- .../content/docs/server-training/overview.md | 19 ++ .../training-server/api-reference.md | 4 +- src/xorl/server/api_server/api_types.py | 166 +++++++++++++++++- src/xorl/server/api_server/endpoints.py | 87 +++++++-- src/xorl/server/api_server/utils.py | 3 +- tests/server/api_server/test_api_server.py | 79 ++++++++- tests/server/api_server/test_api_types.py | 29 ++- 7 files changed, 356 insertions(+), 31 deletions(-) diff --git a/docs/src/content/docs/server-training/overview.md b/docs/src/content/docs/server-training/overview.md index 1c7c37d6..190200fe 100644 --- a/docs/src/content/docs/server-training/overview.md +++ b/docs/src/content/docs/server-training/overview.md @@ -110,6 +110,25 @@ The [xorl-client](/server-training/client-sdk/overview/) Python SDK drives the t --- +## Tinker API Compatibility + +Tinker clients can create a usable session with `POST /api/v1/create_session`. The returned +`session_id` is registered as xorl's backing `model_id`, so follow-up requests may send either +`session_id` (Tinker-style) or `model_id` (xorl-native), and the server normalizes both forms. + +| Endpoint | Description | +|---|---| +| `POST /api/v1/create_session` | Create/register a Tinker-compatible session ID | +| `POST /api/v1/session_heartbeat` | Refresh a session's idle timeout | +| `POST /api/v1/create_model` | Create/register a model session with explicit metadata | +| `POST /api/v1/unload_model` | Unload and release a session | +| `POST /api/v1/forward_backward` | Forward + backward pass | +| `POST /api/v1/optim_step` | Optimizer step | +| `POST /api/v1/weights_info` | Checkpoint metadata for model loading | +| `GET /api/v1/training_runs` | List training runs | + +--- + ## Recommended Reading XoRL's server training mode is designed for online RL with LLMs. The following papers provide background on the algorithms and system designs that XoRL supports: diff --git a/docs/src/content/docs/server-training/training-server/api-reference.md b/docs/src/content/docs/server-training/training-server/api-reference.md index d683a9df..9dddf172 100644 --- a/docs/src/content/docs/server-training/training-server/api-reference.md +++ b/docs/src/content/docs/server-training/training-server/api-reference.md @@ -28,8 +28,8 @@ All endpoints are served at `http://:/`. Training operations use a t | `POST` | `/api/v1/unload_model` | Unload a session, freeing associated adapter state. | | `POST` | `/api/v1/kill_session` | Kill an active session; optionally reload weights from checkpoint. | | `GET` | `/api/v1/session_info` | List active sessions and their state. | -| `POST` | `/api/v1/create_session` | Tinker compatibility stub. Returns a `session_id` but does not register a model session. | -| `POST` | `/api/v1/session_heartbeat` | Tinker compatibility stub; currently a no-op. | +| `POST` | `/api/v1/create_session` | Create and register a Tinker-compatible session ID for follow-up calls. | +| `POST` | `/api/v1/session_heartbeat` | Refresh a session's last-activity timestamp for idle cleanup. | ## Checkpointing diff --git a/src/xorl/server/api_server/api_types.py b/src/xorl/server/api_server/api_types.py index 98a8b75c..5f0df7c0 100644 --- a/src/xorl/server/api_server/api_types.py +++ b/src/xorl/server/api_server/api_types.py @@ -9,6 +9,20 @@ from pydantic import BaseModel, ConfigDict, Field, field_serializer, model_validator +def _map_session_id_to_model_id(data: Any) -> Any: + """Accept Tinker's session_id field anywhere xorl expects model_id.""" + if isinstance(data, dict) and "session_id" in data and "model_id" not in data: + data["model_id"] = data["session_id"] + return data + + +def _map_model_id_to_session_id(data: Any) -> Any: + """Accept xorl's model_id field anywhere the Tinker endpoint expects session_id.""" + if isinstance(data, dict) and "model_id" in data and "session_id" not in data: + data["session_id"] = data["model_id"] + return data + + # ============================================================================ # Two-Phase Async Response Types # ============================================================================ @@ -180,11 +194,18 @@ class ForwardRequest(BaseModel): model_config = ConfigDict(extra="ignore") model_id: str = Field( - default="default", description="Model identifier (must be created via /api/v1/create_model first)" + default="default", + description="Model identifier (must be created via /api/v1/create_model or /api/v1/create_session first)", ) seq_id: Optional[int] = Field(default=None, description="Sequence ID for request ordering") forward_input: DatumInput = Field(..., description="Forward input data") + @model_validator(mode="before") + @classmethod + def _map_tinker_fields(cls, data): + """Map tinker's session_id to model_id.""" + return _map_session_id_to_model_id(data) + class ForwardResponse(BaseModel): """API response for forward operation.""" @@ -206,7 +227,8 @@ class ForwardBackwardRequest(BaseModel): model_config = ConfigDict(extra="ignore") model_id: str = Field( - default="default", description="Model identifier (must be created via /api/v1/create_model first)" + default="default", + description="Model identifier (must be created via /api/v1/create_model or /api/v1/create_session first)", ) seq_id: Optional[int] = Field( default=None, @@ -214,6 +236,12 @@ class ForwardBackwardRequest(BaseModel): ) forward_backward_input: DatumInput = Field(..., description="Forward-backward input data") + @model_validator(mode="before") + @classmethod + def _map_tinker_fields(cls, data): + """Map tinker's session_id to model_id.""" + return _map_session_id_to_model_id(data) + class ForwardBackwardResponse(BaseModel): """API response for forward-backward operation.""" @@ -230,7 +258,8 @@ class OptimStepRequest(BaseModel): model_config = ConfigDict(extra="ignore") model_id: str = Field( - default="default", description="Model identifier (must be created via /api/v1/create_model first)" + default="default", + description="Model identifier (must be created via /api/v1/create_model or /api/v1/create_session first)", ) seq_id: Optional[int] = Field( default=None, @@ -239,6 +268,12 @@ class OptimStepRequest(BaseModel): adam_params: AdamParams = Field(default_factory=AdamParams, description="AdamW optimizer parameters") gradient_clip: Optional[float] = Field(default=None, description="Gradient clipping value") + @model_validator(mode="before") + @classmethod + def _map_tinker_fields(cls, data): + """Map tinker's session_id to model_id.""" + return _map_session_id_to_model_id(data) + class OptimStepResponse(BaseModel): """API response for optimizer step.""" @@ -263,6 +298,8 @@ class SaveWeightsRequest(BaseModel): Endpoint: POST /api/v1/save_weights """ + model_config = ConfigDict(extra="ignore") + model_id: str = Field(default="default", description="Model identifier") path: Optional[str] = Field( default=None, description="Checkpoint name (e.g., 'checkpoint-001'). Auto-generated if not specified." @@ -270,6 +307,12 @@ class SaveWeightsRequest(BaseModel): seq_id: Optional[int] = Field(default=None, description="Sequence ID for request ordering") # Note: save_optimizer is always True - we always save full state for checkpointing + @model_validator(mode="before") + @classmethod + def _map_tinker_fields(cls, data): + """Map tinker's session_id to model_id.""" + return _map_session_id_to_model_id(data) + class SaveWeightsResponse(BaseModel): """API response for saving weights. @@ -287,11 +330,19 @@ class LoadWeightsRequest(BaseModel): Endpoint: POST /api/v1/load_weights """ + model_config = ConfigDict(extra="ignore") + model_id: str = Field(default="default", description="Model identifier") path: str = Field(..., description="Xorl URI to load from (e.g., xorl://default/weights/checkpoint-001)") optimizer: bool = Field(default=True, description="Whether to load optimizer state") seq_id: Optional[int] = Field(default=None, description="Sequence ID for request ordering") + @model_validator(mode="before") + @classmethod + def _map_tinker_fields(cls, data): + """Map tinker's session_id to model_id.""" + return _map_session_id_to_model_id(data) + class LoadWeightsResponse(BaseModel): """API response for loading weights.""" @@ -335,9 +386,8 @@ class CreateModelRequest(BaseModel): @classmethod def _map_tinker_fields(cls, data): """Map tinker's session_id to model_id if model_id not provided.""" + data = _map_session_id_to_model_id(data) if isinstance(data, dict): - if "session_id" in data and "model_id" not in data: - data["model_id"] = data["session_id"] # Tinker sends lora_config with "rank" key; normalize to also include lora_rank lora_cfg = data.get("lora_config") if isinstance(lora_cfg, dict) and "rank" in lora_cfg and "lora_rank" not in lora_cfg: @@ -352,12 +402,54 @@ class CreateModelResponse(BaseModel): type: Literal["create_model"] = Field(default="create_model", description="Response type identifier") +class CreateSessionRequest(BaseModel): + """API request for creating a Tinker-compatible session. + + The request body may be empty. If session_id is omitted, the server generates one. + """ + + model_config = ConfigDict(extra="ignore") + + session_id: Optional[str] = Field(default=None, description="Optional session identifier to register") + base_model: Optional[str] = Field(default=None, description="Optional base model metadata for this session") + lora_config: Dict[str, Any] = Field(default_factory=dict, description="Optional LoRA metadata for this session") + + @model_validator(mode="before") + @classmethod + def _normalize_lora_config(cls, data): + """Normalize optional LoRA metadata for parity with create_model.""" + data = _map_model_id_to_session_id(data) + if isinstance(data, dict): + lora_cfg = data.get("lora_config") + if isinstance(lora_cfg, dict) and "rank" in lora_cfg and "lora_rank" not in lora_cfg: + lora_cfg["lora_rank"] = lora_cfg["rank"] + return data + + +class CreateSessionResponse(BaseModel): + """API response for creating a Tinker-compatible session.""" + + type: Literal["create_session"] = Field(default="create_session", description="Response type identifier") + session_id: str = Field(..., description="Registered session identifier") + info_message: Optional[str] = Field(default=None, description="Informational message") + warning_message: Optional[str] = Field(default=None, description="Warning message") + error_message: Optional[str] = Field(default=None, description="Error message") + + class UnloadModelRequest(BaseModel): """API request for unloading a model (Tinker-compatible).""" + model_config = ConfigDict(extra="ignore") + model_id: str = Field(..., description="Model identifier to unload") type: Literal["unload_model"] = Field(default="unload_model", description="Request type identifier") + @model_validator(mode="before") + @classmethod + def _map_tinker_fields(cls, data): + """Map tinker's session_id to model_id.""" + return _map_session_id_to_model_id(data) + class UnloadModelResponse(BaseModel): """API response for unloading a model (Tinker-compatible).""" @@ -366,6 +458,27 @@ class UnloadModelResponse(BaseModel): type: Optional[Literal["unload_model"]] = Field(default=None, description="Response type identifier") +class SessionHeartbeatRequest(BaseModel): + """API request for session heartbeat / idle timeout refresh.""" + + model_config = ConfigDict(extra="ignore") + + session_id: str = Field(..., description="Session identifier to refresh") + + @model_validator(mode="before") + @classmethod + def _map_xorl_fields(cls, data): + """Accept model_id as an alias for session_id on the heartbeat endpoint.""" + return _map_model_id_to_session_id(data) + + +class SessionHeartbeatResponse(BaseModel): + """API response for session heartbeat.""" + + type: Literal["session_heartbeat"] = Field(default="session_heartbeat", description="Response type identifier") + session_id: str = Field(..., description="Session identifier that was refreshed") + + class SessionInfoResponse(BaseModel): """API response with session information for monitoring.""" @@ -398,6 +511,8 @@ class SessionInfoResponse(BaseModel): class SaveAdapterStateRequest(BaseModel): """API request for saving adapter state (weights + optimizer).""" + model_config = ConfigDict(extra="ignore") + model_id: str = Field(..., description="Adapter/session identifier to save") path: Optional[str] = Field( default=None, description="Directory to save adapter state to. Auto-generated if not specified." @@ -405,6 +520,12 @@ class SaveAdapterStateRequest(BaseModel): save_optimizer: bool = Field(default=True, description="Whether to save optimizer state for resuming training") seq_id: Optional[int] = Field(default=None, description="Sequence ID for request ordering") + @model_validator(mode="before") + @classmethod + def _map_tinker_fields(cls, data): + """Map tinker's session_id to model_id.""" + return _map_session_id_to_model_id(data) + class SaveAdapterStateResponse(BaseModel): """API response for saving adapter state.""" @@ -432,11 +553,20 @@ class RegisterWorkersResponse(BaseModel): class SaveWeightsForSamplerRequest(BaseModel): """API request for saving weights in inference-compatible format.""" + model_config = ConfigDict(extra="ignore") + model_id: str = Field( - default="default", description="Model identifier (must be created via /api/v1/create_model first)" + default="default", + description="Model identifier (must be created via /api/v1/create_model or /api/v1/create_session first)", ) name: str = Field(..., description="Checkpoint name (e.g., 'step-100')") + @model_validator(mode="before") + @classmethod + def _map_tinker_fields(cls, data): + """Map tinker's session_id to model_id.""" + return _map_session_id_to_model_id(data) + class SaveWeightsForSamplerResponse(BaseModel): """API response for saving weights for sampler.""" @@ -466,8 +596,16 @@ class CheckpointInfo(BaseModel): class ListCheckpointsRequest(BaseModel): """API request for listing checkpoints.""" + model_config = ConfigDict(extra="ignore") + model_id: str = Field(default="default", description="Model identifier") + @model_validator(mode="before") + @classmethod + def _map_tinker_fields(cls, data): + """Map tinker's session_id to model_id.""" + return _map_session_id_to_model_id(data) + class ListCheckpointsResponse(BaseModel): """API response for listing checkpoints. @@ -486,11 +624,19 @@ class DeleteCheckpointRequest(BaseModel): - 'sampler_weights/{name}' for sampler checkpoints (flat, no model_id) """ + model_config = ConfigDict(extra="ignore") + model_id: str = Field(default="default", description="Model identifier (used for training checkpoints)") checkpoint_id: str = Field( ..., description="Checkpoint ID to delete (e.g., 'weights/model_id/ckpt-001' or 'sampler_weights/adapter-001')" ) + @model_validator(mode="before") + @classmethod + def _map_tinker_fields(cls, data): + """Map tinker's session_id to model_id.""" + return _map_session_id_to_model_id(data) + class DeleteCheckpointResponse(BaseModel): """API response for deleting a checkpoint.""" @@ -513,10 +659,18 @@ class KillSessionRequest(BaseModel): start a new one. """ + model_config = ConfigDict(extra="ignore") + model_id: str = Field(..., description="Session to kill (must match active session)") save_checkpoint: bool = Field(default=True, description="Save checkpoint before killing") reset_weights: bool = Field(default=False, description="Reload base model weights after killing session") + @model_validator(mode="before") + @classmethod + def _map_tinker_fields(cls, data): + """Map tinker's session_id to model_id.""" + return _map_session_id_to_model_id(data) + class KillSessionResponse(BaseModel): """API response for killing a full-weights training session.""" diff --git a/src/xorl/server/api_server/endpoints.py b/src/xorl/server/api_server/endpoints.py index 6c88290e..f300bc7b 100644 --- a/src/xorl/server/api_server/endpoints.py +++ b/src/xorl/server/api_server/endpoints.py @@ -17,6 +17,8 @@ CreateModelResponse, CreateSamplingSessionRequest, CreateSamplingSessionResponse, + CreateSessionRequest, + CreateSessionResponse, DeleteCheckpointRequest, DeleteCheckpointResponse, ErrorResponse, @@ -35,6 +37,8 @@ RemoveInferenceEndpointResponse, SaveWeightsForSamplerRequest, SaveWeightsRequest, + SessionHeartbeatRequest, + SessionHeartbeatResponse, SessionInfoResponse, SetSyncQuantizationRequest, SetSyncQuantizationResponse, @@ -47,6 +51,7 @@ WeightsInfoRequest, WeightsInfoResponse, ) +from xorl.server.api_server.utils import validate_model_id from xorl.server.protocol.api_orchestrator import OrchestratorRequest from xorl.server.protocol.operations import KillSessionData @@ -784,22 +789,76 @@ async def create_sampling_session_endpoint(request: CreateSamplingSessionRequest # ============================================================================ -@router.post("/api/v1/create_session", tags=["Tinker Compatibility"]) -async def create_session_endpoint(): - """Stub for tinker SDK session creation.""" - return { - "type": "create_session", - "session_id": str(uuid.uuid4()), - "info_message": None, - "warning_message": None, - "error_message": None, - } +@router.post( + "/api/v1/create_session", + response_model=CreateSessionResponse, + responses={ + 400: {"model": ErrorResponse}, + 503: {"model": ErrorResponse}, + }, + tags=["Tinker Compatibility"], +) +async def create_session_endpoint( + request: CreateSessionRequest | None = None, + server=Depends(require_api_server), +): + """Create and register a usable Tinker-compatible session.""" + if request is None: + request = CreateSessionRequest() + + session_id = validate_model_id(request.session_id) if request.session_id else str(uuid.uuid4()) + already_registered = session_id in server.registered_model_ids + + server.registered_model_ids.add(session_id) + + model_config = server.model_configs.setdefault( + session_id, + { + "base_model": request.base_model or server.base_model or "unknown", + "lora_config": request.lora_config, + }, + ) + + if request.base_model: + model_config["base_model"] = request.base_model + if request.lora_config: + model_config["lora_config"] = request.lora_config + + server._update_session_activity(session_id) + + warning_message = None + info_message = f"Session '{session_id}' registered successfully." + if already_registered: + info_message = f"Session '{session_id}' refreshed successfully." + warning_message = f"Session '{session_id}' already existed; refreshed activity timestamp." + + return CreateSessionResponse( + session_id=session_id, + info_message=info_message, + warning_message=warning_message, + error_message=None, + ) + + +@router.post( + "/api/v1/session_heartbeat", + response_model=SessionHeartbeatResponse, + responses={ + 400: {"model": ErrorResponse}, + 404: {"model": ErrorResponse}, + 503: {"model": ErrorResponse}, + }, + tags=["Tinker Compatibility"], +) +async def session_heartbeat_endpoint(request: SessionHeartbeatRequest, server=Depends(require_api_server)): + """Refresh a session's last-activity timestamp for idle cleanup.""" + session_id = validate_model_id(request.session_id) + if session_id not in server.registered_model_ids: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Session {session_id} not found") -@router.post("/api/v1/session_heartbeat", tags=["Tinker Compatibility"]) -async def session_heartbeat_endpoint(): - """Stub for tinker SDK session heartbeat.""" - return {"type": "session_heartbeat"} + server._update_session_activity(session_id) + return SessionHeartbeatResponse(session_id=session_id) @router.get("/api/v1/healthz", tags=["Tinker Compatibility"]) diff --git a/src/xorl/server/api_server/utils.py b/src/xorl/server/api_server/utils.py index 338c6eb3..84d46555 100644 --- a/src/xorl/server/api_server/utils.py +++ b/src/xorl/server/api_server/utils.py @@ -10,7 +10,8 @@ # Error message for model_id not registered MODEL_ID_NOT_REGISTERED_ERROR = ( "model_id '{model_id}' has not been registered. " - "You must call /api/v1/create_model first to register your model_id before calling other /api/v1/ endpoints. " + "You must call /api/v1/create_model or /api/v1/create_session first to register your model_id before " + "calling other /api/v1/ endpoints. " "Try initializing the TrainingClient by calling create_lora_training_client on the ServiceClient." ) diff --git a/tests/server/api_server/test_api_server.py b/tests/server/api_server/test_api_server.py index 112117fd..b6171b91 100644 --- a/tests/server/api_server/test_api_server.py +++ b/tests/server/api_server/test_api_server.py @@ -5,21 +5,28 @@ Integration tests with Orchestrator are excluded as they require complex infrastructure. """ -import pytest - +import asyncio +import time -pytestmark = [pytest.mark.cpu, pytest.mark.server] +import pytest from xorl.server.api_server.api_types import ( AdamParams, + CreateSessionRequest, Datum, DatumInput, ForwardBackwardRequest, LoadWeightsRequest, OptimStepRequest, SaveWeightsRequest, + SessionHeartbeatRequest, + UntypedAPIFuture, ) -from xorl.server.api_server.server import APIServer +from xorl.server.api_server.endpoints import create_session_endpoint, save_weights_endpoint, session_heartbeat_endpoint +from xorl.server.api_server.server import APIServer, app + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] class TestAPIServerConfiguration: @@ -104,6 +111,70 @@ def test_request_creation_and_serialization(self): request = LoadWeightsRequest(path="/tmp/checkpoint", optimizer=True) assert request.optimizer is True +class TestTinkerSessionCompatibility: + """Test Tinker-compatible session creation and heartbeats.""" + + def test_tinker_session_endpoints_are_public_in_openapi(self): + """The repaired Tinker session endpoints should stay in the public schema.""" + app.openapi_schema = None + schema_paths = app.openapi()["paths"] + + assert "/api/v1/create_session" in schema_paths + assert "/api/v1/session_heartbeat" in schema_paths + + def test_create_session_registers_usable_session_id(self, tmp_path): + """Returned session IDs should work in follow-up calls that send session_id.""" + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17010", + engine_output_addr="tcp://127.0.0.1:17011", + output_dir=str(tmp_path), + ) + seen_model_ids = [] + + async def fake_submit_save_weights_async(request): + seen_model_ids.append(request.model_id) + return UntypedAPIFuture(request_id="req-1", model_id=request.model_id) + + server.submit_save_weights_async = fake_submit_save_weights_async + + create_response = asyncio.run(create_session_endpoint(CreateSessionRequest(), server=server)) + session_id = create_response.session_id + + assert session_id in server.registered_model_ids + assert session_id in server.session_last_activity + + save_response = asyncio.run( + save_weights_endpoint( + SaveWeightsRequest(session_id=session_id, path="checkpoint-001"), + server=server, + ) + ) + + assert seen_model_ids == [session_id] + assert save_response.request_id == "req-1" + assert save_response.model_id == session_id + + def test_session_heartbeat_refreshes_activity(self): + """Heartbeats should update the activity timestamp for registered sessions.""" + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17012", + engine_output_addr="tcp://127.0.0.1:17013", + ) + + session_id = "heartbeat-session" + asyncio.run(create_session_endpoint(CreateSessionRequest(session_id=session_id), server=server)) + initial_activity = server.session_last_activity[session_id] + + time.sleep(0.01) + + heartbeat_response = asyncio.run( + session_heartbeat_endpoint( + SessionHeartbeatRequest(session_id=session_id), + server=server, + ) + ) + assert heartbeat_response.session_id == session_id + assert server.session_last_activity[session_id] > initial_activity if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/server/api_server/test_api_types.py b/tests/server/api_server/test_api_types.py index b2f1276a..098b4ab2 100644 --- a/tests/server/api_server/test_api_types.py +++ b/tests/server/api_server/test_api_types.py @@ -5,14 +5,11 @@ """ import pytest - - -pytestmark = [pytest.mark.cpu, pytest.mark.server] - from pydantic import ValidationError from xorl.server.api_server.api_types import ( AdamParams, + CreateModelRequest, Datum, DatumInput, ErrorResponse, @@ -31,6 +28,9 @@ ) +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + class TestDatumAndForwardBackward: """Test Datum, DatumInput, ForwardBackwardRequest and ForwardBackwardResponse.""" @@ -90,6 +90,14 @@ def test_datum_forward_backward_request_and_response(self): assert request.model_id == "default" assert request.forward_backward_input.loss_fn == "causallm_loss" + request = ForwardBackwardRequest( + session_id="session-123", + forward_backward_input=DatumInput( + data=[Datum(model_input={"input_ids": [1, 2, 3]}, loss_fn_inputs={"labels": [2, 3, 4]})], + ), + ) + assert request.model_id == "session-123" + request = ForwardBackwardRequest( forward_backward_input=DatumInput( data=[ @@ -167,6 +175,8 @@ def test_weights_health_error_and_serialization(self): request = SaveWeightsRequest(model_id="test-model", path="/tmp/checkpoint") assert request.model_id == "test-model" assert request.path == "/tmp/checkpoint" + request = SaveWeightsRequest(session_id="session-123", path="/tmp/checkpoint") + assert request.model_id == "session-123" request = SaveWeightsRequest() assert request.model_id == "default" assert request.path is None @@ -180,12 +190,23 @@ def test_weights_health_error_and_serialization(self): # LoadWeightsRequest request = LoadWeightsRequest(model_id="test-model", path="/tmp/checkpoint", optimizer=True) assert request.optimizer is True + request = LoadWeightsRequest(session_id="session-123", path="/tmp/checkpoint") + assert request.model_id == "session-123" request = LoadWeightsRequest(path="/tmp/checkpoint") assert request.model_id == "default" assert request.optimizer is True with pytest.raises(ValidationError): LoadWeightsRequest() + # CreateModelRequest + request = CreateModelRequest( + session_id="session-123", + base_model="Qwen/Qwen3-8B", + lora_config={"rank": 64}, + ) + assert request.model_id == "session-123" + assert request.lora_config["lora_rank"] == 64 + # LoadWeightsResponse response = LoadWeightsResponse(path="xorl://default/weights/checkpoint-001") assert response.path == "xorl://default/weights/checkpoint-001" From 38732cc3c88b440fe4933dfa0a6687413091c481 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Fri, 10 Apr 2026 13:36:29 -0700 Subject: [PATCH 26/41] Fix causal attention FLOP accounting (#104) * Fix causal attention FLOP accounting * Remove unsupported FLOPs model aliases (cherry picked from commit cd140158137790ac61f9466c8dab3ebc5564aaad) --- src/xorl/utils/count_flops.py | 47 ++++++++++++--------- tests/e2e/qwen3_8b/test_tflops_threshold.py | 24 +++++------ 2 files changed, 38 insertions(+), 33 deletions(-) diff --git a/src/xorl/utils/count_flops.py b/src/xorl/utils/count_flops.py index 71601390..cd946afa 100644 --- a/src/xorl/utils/count_flops.py +++ b/src/xorl/utils/count_flops.py @@ -9,6 +9,17 @@ logger = logging.get_logger(__name__) +def _attention_score_elements(batch_seqlens: List[int], causal: bool) -> int: + """Return attended query-key elements across a batch of packed sequences. + + Decoder self-attention is lower-triangular, so the quadratic term is + ``s * (s + 1) / 2`` per sequence instead of ``s * s``. + """ + if causal: + return sum(seqlen * (seqlen + 1) // 2 for seqlen in batch_seqlens) + return sum(seqlen * seqlen for seqlen in batch_seqlens) + + def _gc_multipliers( gc_enabled: bool, recompute_modules: Optional[List[str]], @@ -26,7 +37,7 @@ def _gc_multipliers( Returns a dict with keys: attn_linear - Q/K/V/O projection FLOPs multiplier - attn_qkv - attention softmax/QK^T/SV FLOPs multiplier + attn_qkv - attention QK^T/SV training FLOPs multiplier per attended element router - MoE gate router multiplier gate - MoE gate_proj multiplier up - MoE up_proj multiplier @@ -93,14 +104,12 @@ def __init__( # all_reduce(sum, dp) / world_size gives correct per-GPU TFlops. self._cp_size = cp_size self.estimate_func = { - "qwen2_vl": self._estimate_qwen2_vl_flops, # the only difference between Qwen2 and Qwen2.5 for counting flops is the window attention # used in the ViT for Qwen2.5VL which is considered in the _estimate_qwen2_vl_flops function. "qwen2_5_vl": self._estimate_qwen2_vl_flops, "deepseek_v3": self._estimate_deepseek_v3_flops, "qwen3_moe": self._estimate_qwen3_moe_flops, "llama": self._estimate_llama_flops, - "qwen2": self._estimate_qwen2_flops, # qwen3 reused _estimate_qwen2_flops func because the only model structure diff between qwen2 dense and qwen3 dense is that # qwen3 has additional RMSNorm layers for q and k. # RMSNorm layers have minimal impact at the MFU and can be ignored. @@ -162,8 +171,8 @@ def _estimate_deepseek_v3_flops(self, tokens_sum, batch_seqlens, delta_time): + 6 * embed_lm_N * tokens_sum ) - seqlen_square_sum = sum(s * s * num_hidden_layers for s in batch_seqlens) - attn_qkv_flops = m["attn_qkv"] * seqlen_square_sum * q_head_dim * num_query_heads + attn_score_elements = _attention_score_elements(batch_seqlens, causal=True) + attn_qkv_flops = m["attn_qkv"] * attn_score_elements * q_head_dim * num_query_heads * num_hidden_layers return (dense_N_flops + attn_qkv_flops) / delta_time / 1e12 @@ -201,8 +210,8 @@ def _estimate_qwen3_moe_flops(self, tokens_sum, batch_seqlens, delta_time): dense_N_flops = per_layer_flops * num_hidden_layers + 6 * embed_lm_N * tokens_sum # Attention QKV FLOPs (quadratic in sequence length) - seqlen_square_sum = sum(s * s for s in batch_seqlens) - attn_qkv_flops = m["attn_qkv"] * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers + attn_score_elements = _attention_score_elements(batch_seqlens, causal=True) + attn_qkv_flops = m["attn_qkv"] * attn_score_elements * head_dim * num_attention_heads * num_hidden_layers flops_all_token = dense_N_flops + attn_qkv_flops return flops_all_token / delta_time / 1e12 @@ -264,8 +273,8 @@ def _estimate_qwen3_5_flops(self, tokens_sum, batch_seqlens, delta_time): + 6 * embed_lm_N * tokens_sum ) - seqlen_square_sum = sum(s * s for s in batch_seqlens) - attn_qkv_flops = m["attn_qkv"] * seqlen_square_sum * head_dim * num_attention_heads * full_attn_layers + attn_score_elements = _attention_score_elements(batch_seqlens, causal=True) + attn_qkv_flops = m["attn_qkv"] * attn_score_elements * head_dim * num_attention_heads * full_attn_layers return (dense_N_flops + attn_qkv_flops) / delta_time / 1e12 @@ -335,8 +344,8 @@ def _estimate_qwen3_5_moe_flops(self, tokens_sum, batch_seqlens, delta_time): + 6 * embed_lm_N * tokens_sum ) - seqlen_square_sum = sum(s * s for s in batch_seqlens) - attn_qkv_flops = m["attn_qkv"] * seqlen_square_sum * head_dim * num_attention_heads * full_attn_layers + attn_score_elements = _attention_score_elements(batch_seqlens, causal=True) + attn_qkv_flops = m["attn_qkv"] * attn_score_elements * head_dim * num_attention_heads * full_attn_layers return (dense_N_flops + attn_qkv_flops) / delta_time / 1e12 @@ -361,8 +370,8 @@ def _estimate_qwen2_flops(self, tokens_sum, batch_seqlens, delta_time): per_layer_flops = m["dense_mlp"] * mlp_N * tokens_sum + m["attn_linear"] * attn_linear_N * tokens_sum dense_N_flops = per_layer_flops * num_hidden_layers + 6 * embed_lm_N * tokens_sum - seqlen_square_sum = sum(s * s for s in batch_seqlens) - attn_qkv_flops = m["attn_qkv"] * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers + attn_score_elements = _attention_score_elements(batch_seqlens, causal=True) + attn_qkv_flops = m["attn_qkv"] * attn_score_elements * head_dim * num_attention_heads * num_hidden_layers flops_all_token = dense_N_flops + attn_qkv_flops return flops_all_token / delta_time / 1e12 @@ -394,10 +403,8 @@ def _estimate_qwen2_vl_flops(self, tokens_sum, batch_seqlens, delta_time, **karg dense_N_flops = 6 * dense_N * tokens_sum # attn all_layer & all_token fwd & bwd flops - seqlen_square_sum = 0 - for seqlen in batch_seqlens: - seqlen_square_sum += seqlen * seqlen - attn_qkv_flops = 12 * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers + attn_score_elements = _attention_score_elements(batch_seqlens, causal=True) + attn_qkv_flops = 12 * attn_score_elements * head_dim * num_attention_heads * num_hidden_layers # vit flops image_seqlens = kargs.get("image_seqlens", None) @@ -457,10 +464,8 @@ def _estimate_qwen_vit_flop(self, image_seqlens, config): window_attn_layer_num = config.depth - full_attn_layer_num # full attn layer & all_token fwd & bwd flops - seqlen_square_sum = 0 - for seqlen in image_seqlens: - seqlen_square_sum += seqlen * seqlen - attn_qkv_flops = 12 * seqlen_square_sum * head_dim * num_heads * full_attn_layer_num + attn_score_elements = _attention_score_elements(image_seqlens, causal=False) + attn_qkv_flops = 12 * attn_score_elements * head_dim * num_heads * full_attn_layer_num # If window attention is used, add the window attention flops if window_attn_layer_num > 0: diff --git a/tests/e2e/qwen3_8b/test_tflops_threshold.py b/tests/e2e/qwen3_8b/test_tflops_threshold.py index 10856469..4a03286f 100644 --- a/tests/e2e/qwen3_8b/test_tflops_threshold.py +++ b/tests/e2e/qwen3_8b/test_tflops_threshold.py @@ -2,12 +2,12 @@ Runs LoRA SFT benchmarks with the real Qwen3-8B model and asserts that achieved TFLOPS meets minimum thresholds. Thresholds are set at -80% of measured baselines on H100 80GB HBM3 with flash_attention_3. +80% of corrected baselines on H100 80GB HBM3 with flash_attention_3. -Baselines (Qwen3-8B LoRA rank=8, seq_len=4096, 4 samples, 10 steps): - 1 GPU: ~447 TFLOPS -> threshold 358 - 2 GPUs: ~420 TFLOPS -> threshold 336 - 4 GPUs: ~411 TFLOPS -> threshold 329 +Corrected baselines (Qwen3-8B LoRA rank=8, seq_len=4096, 4 samples, 10 steps): + 1 GPU: ~416 TFLOPS -> threshold 333 + 2 GPUs: ~391 TFLOPS -> threshold 313 + 4 GPUs: ~383 TFLOPS -> threshold 306 These tests are slow (~2-3 min each) and require real Qwen3-8B weights. Run with: pytest tests/e2e/qwen3_8b/test_tflops_threshold.py -v -s @@ -47,11 +47,11 @@ QWEN3_8B_DIR = "/data/shared/huggingface/hub/models--Qwen--Qwen3-8B/snapshots/b968826d9c46dd6066d109eabc6255188de91218" VOCAB_SIZE = 151936 -# Minimum TFLOPS per GPU (80% of H100 baseline, flash_attention_3, seq_len=4096) +# Minimum TFLOPS per GPU (80% of corrected H100 baseline, flash_attention_3, seq_len=4096) MIN_TFLOPS = { - 1: 358, - 2: 336, - 4: 329, + 1: 333, + 2: 313, + 4: 306, } # Benchmark parameters @@ -164,7 +164,7 @@ class TestQwen3_8B_TFLOPS: @skip_if_gpu_count_less_than(1) @_skip_if_no_qwen3_8b() def test_1gpu_tflops(self, tmp_path): - """1-GPU Qwen3-8B LoRA must achieve >= 358 TFLOPS on H100.""" + """1-GPU Qwen3-8B LoRA must achieve >= 333 TFLOPS on H100.""" result = _run_tflops_benchmark(num_gpus=1, dp_shard_size=1, tmp_path=str(tmp_path)) assert result["tflops"] >= MIN_TFLOPS[1], f"1-GPU TFLOPS {result['tflops']:.1f} below threshold {MIN_TFLOPS[1]}" assert result["losses"][-1] < result["losses"][0], "Loss should decrease" @@ -172,7 +172,7 @@ def test_1gpu_tflops(self, tmp_path): @skip_if_gpu_count_less_than(2) @_skip_if_no_qwen3_8b() def test_2gpu_tflops(self, tmp_path): - """2-GPU Qwen3-8B LoRA FSDP2 must achieve >= 336 TFLOPS on H100.""" + """2-GPU Qwen3-8B LoRA FSDP2 must achieve >= 313 TFLOPS on H100.""" result = _run_tflops_benchmark(num_gpus=2, dp_shard_size=2, tmp_path=str(tmp_path)) assert result["tflops"] >= MIN_TFLOPS[2], f"2-GPU TFLOPS {result['tflops']:.1f} below threshold {MIN_TFLOPS[2]}" assert result["losses"][-1] < result["losses"][0], "Loss should decrease" @@ -180,7 +180,7 @@ def test_2gpu_tflops(self, tmp_path): @skip_if_gpu_count_less_than(4) @_skip_if_no_qwen3_8b() def test_4gpu_tflops(self, tmp_path): - """4-GPU Qwen3-8B LoRA FSDP2 must achieve >= 329 TFLOPS on H100.""" + """4-GPU Qwen3-8B LoRA FSDP2 must achieve >= 306 TFLOPS on H100.""" result = _run_tflops_benchmark(num_gpus=4, dp_shard_size=4, tmp_path=str(tmp_path)) assert result["tflops"] >= MIN_TFLOPS[4], f"4-GPU TFLOPS {result['tflops']:.1f} below threshold {MIN_TFLOPS[4]}" assert result["losses"][-1] < result["losses"][0], "Loss should decrease" From 6f9e78200f42ec3a568ff9e115f3f547ff954e59 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Fri, 10 Apr 2026 19:16:51 -0700 Subject: [PATCH 27/41] Optimizer refactoring + Muon efficiency improvements (#90) * Add Gram Newton-Schulz backend for Muon (#97) * Add Gram Newton-Schulz backend for Muon * Fix QLoRA import cycles for benchmark entrypoints * Prefer installed quack for Muon Gram-NS * Use upstream quack for Muon Gram-NS * Batch Gram Newton-Schulz updates in Muon * Add Muon dtype controls and split optimizer modules * Add Muon transient update dtype control * Add Muon force momentum path flag * Fix Muon dtype handling for server and Gram-NS * Fallback to torch matmuls for unsupported Quack dtypes * Add Gram Newton-Schulz backend for Muon (#97) * Add Gram Newton-Schulz backend for Muon * Fix QLoRA import cycles for benchmark entrypoints * Prefer installed quack for Muon Gram-NS * Use upstream quack for Muon Gram-NS * Batch Gram Newton-Schulz updates in Muon * Reconcile Muon Gram-NS batching changes * Format Muon batched NS list comprehension * Batch grouped Gram-NS by matrix shape (cherry picked from commit 37205d38c4f51053b33a42970264a0e04ec900f4) --- src/xorl/arguments.py | 70 +++- src/xorl/lora/mapping.py | 4 +- src/xorl/optim/__init__.py | 14 +- src/xorl/optim/anyprecision_adamw.py | 96 ++++++ src/xorl/optim/gram_newton_schulz.py | 332 ++++++++++++++++++ src/xorl/optim/multi_optimizer.py | 96 ++++++ src/xorl/optim/muon.py | 282 +++++++++++++-- src/xorl/optim/optimizer.py | 273 ++------------- src/xorl/optim/signsgd.py | 50 +++ src/xorl/qlora/__init__.py | 102 ++---- src/xorl/qlora/modules/__init__.py | 27 +- src/xorl/server/runner/model_runner.py | 18 +- src/xorl/server/server_arguments.py | 60 ++++ tests/optim/test_muon.py | 456 +++++++++++++++++++++++++ tests/qlora/test_imports.py | 13 + tests/server/test_server_arguments.py | 41 +++ tests/test_arguments.py | 50 +++ 17 files changed, 1616 insertions(+), 368 deletions(-) create mode 100644 src/xorl/optim/anyprecision_adamw.py create mode 100644 src/xorl/optim/gram_newton_schulz.py create mode 100644 src/xorl/optim/multi_optimizer.py create mode 100644 src/xorl/optim/signsgd.py create mode 100644 tests/optim/test_muon.py create mode 100644 tests/qlora/test_imports.py diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 4afb0306..7505f45a 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -580,6 +580,55 @@ class TrainingArguments: "None defaults to 'original'." }, ) + muon_ns_algorithm: Literal["standard_newton_schulz", "gram_newton_schulz"] = field( + default="standard_newton_schulz", + metadata={ + "help": "Newton-Schulz backend for Muon. 'standard_newton_schulz' keeps the PyTorch Muon path; " + "'gram_newton_schulz' uses Dao-AILab's Gram Newton-Schulz formulation." + }, + ) + muon_ns_use_quack_kernels: bool = field( + default=True, + metadata={ + "help": "Allow Muon Gram Newton-Schulz to use Quack symmetric GEMM kernels on supported Hopper/Blackwell GPUs. " + "Falls back to torch matmuls when unavailable." + }, + ) + muon_gram_ns_num_restarts: int = field( + default=1, + metadata={ + "help": "Number of restart locations to autotune for Muon's Gram Newton-Schulz backend when explicit " + "muon_gram_ns_restart_iterations are not provided." + }, + ) + muon_gram_ns_restart_iterations: Optional[List[int]] = field( + default=None, + metadata={ + "help": "Explicit restart iteration indices for Muon's Gram Newton-Schulz backend. " + "A value of 2 means restart after the second iteration." + }, + ) + muon_grad_dtype: Optional[Literal["fp32", "bf16"]] = field( + default=None, + metadata={ + "help": "Optional dtype cast for the gradient tensor used inside Muon. " + "Use this to force the Muon optimizer path to fp32 or bf16 independently of momentum state dtype." + }, + ) + muon_update_dtype: Optional[Literal["fp32", "bf16"]] = field( + default=None, + metadata={ + "help": "Optional dtype cast for the transient Muon update tensor passed into Newton-Schulz. " + "Use this to decouple compute dtype from gradient and momentum-buffer storage dtype." + }, + ) + muon_force_momentum_path: bool = field( + default=False, + metadata={ + "help": "Force Muon to build the update through the momentum-buffer path even when muon_momentum=0. " + "Intended for debugging and ablations." + }, + ) @property def optimizer_kwargs(self) -> Dict[str, Any]: @@ -591,9 +640,24 @@ def optimizer_kwargs(self) -> Dict[str, Any]: kwargs["muon_nesterov"] = self.muon_nesterov kwargs["muon_ns_steps"] = self.muon_ns_steps kwargs["muon_adjust_lr_fn"] = self.muon_adjust_lr_fn + kwargs["muon_ns_algorithm"] = self.muon_ns_algorithm + kwargs["muon_ns_use_quack_kernels"] = self.muon_ns_use_quack_kernels + kwargs["muon_gram_ns_num_restarts"] = self.muon_gram_ns_num_restarts + if self.muon_gram_ns_restart_iterations is not None: + kwargs["muon_gram_ns_restart_iterations"] = self.muon_gram_ns_restart_iterations # Wire optimizer_dtype -> muon_momentum_dtype so "bf16" sets bf16 Muon momentum if self.optimizer_dtype == "bf16": kwargs["muon_momentum_dtype"] = torch.bfloat16 + if self.muon_grad_dtype == "bf16": + kwargs["muon_grad_dtype"] = torch.bfloat16 + elif self.muon_grad_dtype == "fp32": + kwargs["muon_grad_dtype"] = torch.float32 + if self.muon_update_dtype == "bf16": + kwargs["muon_update_dtype"] = torch.bfloat16 + elif self.muon_update_dtype == "fp32": + kwargs["muon_update_dtype"] = torch.float32 + if self.muon_force_momentum_path: + kwargs["muon_force_momentum_path"] = True return kwargs max_grad_norm: float = field( @@ -1037,12 +1101,6 @@ class LoRAArguments: default=False, metadata={"help": "Only save LoRA weights (not full model) in HF checkpoints"}, ) - moe_hybrid_shared_lora: bool = field( - default=False, - metadata={ - "help": "Enable hybrid shared LoRA for MoE: share lora_A for gate/up_proj, lora_B for down_proj across experts" - }, - ) # QLoRA: quantize base weights for memory savings enable_qlora: bool = field( default=False, diff --git a/src/xorl/lora/mapping.py b/src/xorl/lora/mapping.py index 3c207016..25222143 100644 --- a/src/xorl/lora/mapping.py +++ b/src/xorl/lora/mapping.py @@ -8,8 +8,6 @@ import torch.nn as nn -from xorl.models.layers.moe import MoEExperts, MoEExpertsLoRA - from .modules.base import LoraModule from .modules.linear import LoraLinear @@ -29,6 +27,8 @@ def _ensure_moe_mapping(): global _moe_registered if not _moe_registered: + from xorl.models.layers.moe import MoEExperts, MoEExpertsLoRA # noqa: PLC0415 + _moe_registered = True LORA_MAPPING[MoEExperts] = MoEExpertsLoRA diff --git a/src/xorl/optim/__init__.py b/src/xorl/optim/__init__.py index 1b7f3111..cef34cee 100644 --- a/src/xorl/optim/__init__.py +++ b/src/xorl/optim/__init__.py @@ -1,6 +1,16 @@ +from .anyprecision_adamw import AnyPrecisionAdamW from .lr_scheduler import build_lr_scheduler +from .multi_optimizer import MultiOptimizer from .muon import Muon -from .optimizer import SignSGD, build_optimizer +from .optimizer import build_optimizer +from .signsgd import SignSGD -__all__ = ["build_lr_scheduler", "build_optimizer", "Muon", "SignSGD"] +__all__ = [ + "AnyPrecisionAdamW", + "build_lr_scheduler", + "build_optimizer", + "MultiOptimizer", + "Muon", + "SignSGD", +] diff --git a/src/xorl/optim/anyprecision_adamw.py b/src/xorl/optim/anyprecision_adamw.py new file mode 100644 index 00000000..3785fbc0 --- /dev/null +++ b/src/xorl/optim/anyprecision_adamw.py @@ -0,0 +1,96 @@ +import torch +from torch.optim.optimizer import Optimizer + + +class AnyPrecisionAdamW(Optimizer): + def __init__( + self, + params, + lr=1e-3, + betas=(0.9, 0.95), + eps=1e-8, + weight_decay=0.0, + use_kahan_summation=False, + momentum_dtype=torch.bfloat16, + variance_dtype=torch.bfloat16, + compensation_buffer_dtype=torch.bfloat16, + ): + defaults = { + "lr": lr, + "betas": betas, + "eps": eps, + "weight_decay": weight_decay, + "use_kahan_summation": use_kahan_summation, + "momentum_dtype": momentum_dtype, + "variance_dtype": variance_dtype, + "compensation_buffer_dtype": compensation_buffer_dtype, + } + super().__init__(params, defaults) + + @torch.no_grad() + def step(self, closure=None): + """ + Performs a single optimization step. + + Args: + closure (callable, optional): A closure that reevaluates the model and returns the loss. + """ + + if closure is not None: + with torch.enable_grad(): + closure() + + for group in self.param_groups: + beta1, beta2 = group["betas"] + lr = group["lr"] + weight_decay = group["weight_decay"] + eps = group["eps"] + use_kahan_summation = group["use_kahan_summation"] + + momentum_dtype = group["momentum_dtype"] + variance_dtype = group["variance_dtype"] + compensation_buffer_dtype = group["compensation_buffer_dtype"] + for p in group["params"]: + if p.grad is None: + continue + + if p.grad.is_sparse: + raise RuntimeError("AnyPrecisionAdamW does not support sparse gradients.") + + state = self.state[p] + if len(state) == 0: + state["step"] = torch.tensor(0.0) + state["exp_avg"] = torch.zeros_like(p, dtype=momentum_dtype) + state["exp_avg_sq"] = torch.zeros_like(p, dtype=variance_dtype) + + if use_kahan_summation: + state["compensation"] = torch.zeros_like(p, dtype=compensation_buffer_dtype) + + state["step"] += 1 + step = state["step"] + + exp_avg = state["exp_avg"] + exp_avg_sq = state["exp_avg_sq"] + grad = p.grad + + if weight_decay: + p.data.mul_(1 - lr * weight_decay) + + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + + bias_correction1 = 1 - beta1**step + step_size = lr / bias_correction1 + + denom_correction = (1 - beta2**step) ** 0.5 + centered_variance = (exp_avg_sq.sqrt() / denom_correction).add_(eps, alpha=1) + + if use_kahan_summation: + compensation = state["compensation"] + compensation.addcdiv_(exp_avg, centered_variance, value=-step_size) + + temp_buffer = p.detach().clone() + p.data.add_(compensation) + compensation.add_(temp_buffer.sub_(p.data)) + else: + p.data.addcdiv_(exp_avg, centered_variance, value=-step_size) diff --git a/src/xorl/optim/gram_newton_schulz.py b/src/xorl/optim/gram_newton_schulz.py new file mode 100644 index 00000000..be319044 --- /dev/null +++ b/src/xorl/optim/gram_newton_schulz.py @@ -0,0 +1,332 @@ +""" +Gram Newton-Schulz orthogonalization for Muon. + +This is an in-tree adaptation of Dao-AILab/gram-newton-schulz that keeps the +algorithm local to Xorl so it can be used inside the existing FSDP/DTensor +Muon optimizer without depending on the external package's single-GPU Muon +wrapper. +""" + +import importlib +import os +from functools import lru_cache +from itertools import combinations +from types import SimpleNamespace +from typing import Iterable, List, Optional, Sequence, Tuple + +import numpy as np +import torch +from torch import Tensor + +from ..utils import logging + + +logger = logging.get_logger(__name__) + +SYMMETRIC_KERNEL_TILE_SIZE = 256 +MOST_NEGATIVE_GRAM_EIGENVALUE = -4e-4 + + +def _torch_sym_baddbmm(A: Tensor, B: Tensor, C: Tensor, alpha: float = 1.0, beta: float = 1.0) -> Tensor: + if A.ndim == 2: + return torch.addmm(C, A, B, beta=beta, alpha=alpha) + return torch.baddbmm(C, A, B, beta=beta, alpha=alpha) + + +def _torch_mm_add(A: Tensor, B: Tensor, C: Tensor, beta: float = 1.0) -> Tensor: + if A.ndim == 2: + return torch.addmm(C, A, B, beta=beta) + return torch.baddbmm(C, A, B, beta=beta) + + +_TORCH_BACKEND = SimpleNamespace( + sym_mm=lambda A, B: A @ B, + sym_baddbmm=_torch_sym_baddbmm, + mm=lambda A, B: A @ B, + mm_add=_torch_mm_add, +) + + +def _quack_supports_symmetric_gemm_dtype(dtype: torch.dtype, capability: tuple[int, int]) -> bool: + if dtype == torch.float32: + return capability[0] >= 10 + return dtype in (torch.float16, torch.bfloat16) + + +def _ensure_cutlass_arch_for_current_device() -> None: + if os.environ.get("CUTE_DSL_ARCH") or not torch.cuda.is_available(): + return + + major, minor = torch.cuda.get_device_capability() + if major == 11 and minor == 0: + major, minor = 10, 1 + + suffix = "a" if major >= 9 else "" + os.environ["CUTE_DSL_ARCH"] = f"sm_{major}{minor}{suffix}" + + +def _import_quack_gemm_interface(): + try: + return importlib.import_module("quack.gemm_interface") + except Exception as exc: + raise ImportError("Muon Gram Newton-Schulz requires the upstream `quack-kernels` package") from exc + + +@lru_cache(maxsize=1) +def _make_quack_backend(): + _ensure_cutlass_arch_for_current_device() + gemm_interface = _import_quack_gemm_interface() + gemm = gemm_interface.gemm + gemm_add = gemm_interface.gemm_add + gemm_symmetric = gemm_interface.gemm_symmetric + + return SimpleNamespace( + sym_mm=lambda A, B: gemm_symmetric(A, B), + sym_baddbmm=lambda A, B, C, alpha=1.0, beta=1.0: gemm_symmetric(A, B, C=C, alpha=alpha, beta=beta), + mm=lambda A, B: gemm(A, B), + mm_add=lambda A, B, C, beta=1.0: gemm_add(A, B, C=C, beta=beta), + ) + + +def expand_ns_coefficients( + ns_coefficients: Sequence[float] | Sequence[Sequence[float]], + ns_steps: int, +) -> List[Tuple[float, float, float]]: + """ + Expand Muon's coefficient config into one coefficient triple per iteration. + + PyTorch Muon uses a single `(a, b, c)` tuple repeated `ns_steps` times. + Gram Newton-Schulz needs the per-iteration coefficients explicitly. + """ + if ns_steps <= 0: + raise ValueError(f"ns_steps must be positive, got {ns_steps}") + + if len(ns_coefficients) == 3 and isinstance(ns_coefficients[0], (float, int)): + triple = tuple(float(v) for v in ns_coefficients) + return [triple] * ns_steps + + expanded = [tuple(float(v) for v in coeff) for coeff in ns_coefficients] + if len(expanded) != ns_steps: + raise ValueError( + f"Expected {ns_steps} coefficient triples, got {len(expanded)}. " + "Pass a single (a, b, c) tuple or one triple per step." + ) + for coeff in expanded: + if len(coeff) != 3: + raise ValueError(f"Each Newton-Schulz coefficient set must have 3 values, got {len(coeff)}") + return expanded + + +def simulate_perturbed_gram_newton_schulz( + x_eigenvalues: np.ndarray, + coefficients: Sequence[Tuple[float, float, float]], + perturbation: float, + reset_iterations: Optional[Iterable[int]] = None, +) -> dict[str, np.ndarray]: + """Mirror the restart autotune heuristic used in Dao-AILab/gram-newton-schulz.""" + if perturbation >= 0: + raise ValueError(f"perturbation must be negative, got {perturbation}") + + q_values: dict[str, np.ndarray] = {} + x_eigenvalues = np.asarray(x_eigenvalues, dtype=np.float64).copy() + reset_iterations = set(reset_iterations or ()) + q = np.ones_like(x_eigenvalues) + + with np.errstate(over="ignore", invalid="ignore"): + for iteration, (a, b, c) in enumerate(coefficients): + if iteration == 0 or iteration in reset_iterations: + if iteration != 0: + x_eigenvalues *= q + r = x_eigenvalues**2 + perturbation + q = np.ones_like(x_eigenvalues) + + z = a + r * (b + r * c) + q *= z + r *= z**2 + q_values[f"Q_{iteration}"] = q.astype(np.float64) + + return q_values + + +def stability_metric(q_values: dict[str, np.ndarray]) -> float: + def condition(values: np.ndarray) -> float: + abs_values = np.abs(values) + return float(abs_values.max() / abs_values.min()) + + return max(condition(values) for values in q_values.values()) + + +def find_best_restarts( + coefficients: Sequence[Tuple[float, float, float]], + *, + num_restarts: int = 1, + most_negative_gram_eigenvalue: float = MOST_NEGATIVE_GRAM_EIGENVALUE, +) -> List[int]: + """ + Find restart positions with the same heuristic as Dao-AILab's autotuner. + + Restart positions are zero-based iteration indices at which the algorithm + restarts before executing that iteration. For example, `[2]` means restart + after finishing the second iteration. + """ + if num_restarts < 0: + raise ValueError(f"num_restarts must be non-negative, got {num_restarts}") + + possible_positions = list(range(1, len(coefficients))) + x_eigenvalues = np.logspace(0, -10, 10000) + + if num_restarts == 0: + return [] + if num_restarts > len(possible_positions): + raise ValueError(f"Cannot have {num_restarts} restarts with only {len(coefficients)} Newton-Schulz iterations") + + best_restarts: Optional[List[int]] = None + best_max_q = float("inf") + + for restart_combo in combinations(possible_positions, num_restarts): + restart_positions = list(restart_combo) + q_values = simulate_perturbed_gram_newton_schulz( + x_eigenvalues, + coefficients, + most_negative_gram_eigenvalue, + reset_iterations=restart_positions, + ) + max_q = stability_metric(q_values) + if max_q < best_max_q: + best_max_q = max_q + best_restarts = restart_positions + + if best_restarts is None or not np.isfinite(best_max_q) or best_max_q >= 1e8: + raise ValueError( + "Failed to find numerically stable Gram Newton-Schulz restarts. " + "Increase the restart count or provide explicit restart positions." + ) + + return best_restarts + + +class GramNewtonSchulzOrthogonalizer: + """Orthogonalize a 2D matrix with Gram Newton-Schulz and optional Quack kernels.""" + + def __init__( + self, + *, + ns_coefficients: Sequence[Tuple[float, float, float]], + ns_epsilon: float = 1e-7, + ns_use_quack_kernels: bool = True, + gram_newton_schulz_restart_iterations: Optional[Sequence[int]] = None, + ) -> None: + if ns_epsilon <= 0: + raise ValueError(f"ns_epsilon must be positive, got {ns_epsilon}") + + self.ns_coefficients = tuple(tuple(float(v) for v in coeff) for coeff in ns_coefficients) + if not self.ns_coefficients: + raise ValueError("ns_coefficients must contain at least one iteration") + for coeff in self.ns_coefficients: + if len(coeff) != 3: + raise ValueError(f"Each Newton-Schulz coefficient set must have 3 values, got {len(coeff)}") + + self.ns_epsilon = ns_epsilon + self.ns_use_quack_kernels = ns_use_quack_kernels + self.reset_iterations = tuple(sorted(set(gram_newton_schulz_restart_iterations or ()))) + self._logged_quack_fallback = False + self._logged_quack_dtype_fallback = False + + def orthogonalize(self, X: Tensor) -> Tensor: + """Orthogonalize a 2D matrix and return a tensor with the original shape and dtype.""" + original_shape = X.shape + if X.ndim == 2: + X = X.unsqueeze(0) + elif X.ndim > 3: + X = X.reshape(-1, *X.shape[-2:]) + + original_dtype = X.dtype + X = X.to(torch.float32) + + should_transpose = X.size(-2) > X.size(-1) + if should_transpose: + X = X.mT + + X /= X.norm(dim=(-2, -1), keepdim=True).clamp_min(self.ns_epsilon) + + if X.is_cuda: + # Respect the caller-selected Muon compute dtype on CUDA so dtype + # sweeps can actually compare bf16 vs fp32 Newton-Schulz behavior. + compute_dtype = ( + original_dtype if original_dtype in (torch.float16, torch.bfloat16, torch.float32) else torch.float32 + ) + else: + compute_dtype = torch.float32 + X = X.to(compute_dtype) + + if X.size(-2) == X.size(-1): + X = self._standard_newton_schulz(X) + else: + X = self._gram_newton_schulz(X) + + if should_transpose: + X = X.mT + + return X.to(original_dtype).reshape(original_shape) + + def _select_backend(self, X: Tensor): + if not self.ns_use_quack_kernels or not X.is_cuda or min(X.size(-2), X.size(-1)) <= SYMMETRIC_KERNEL_TILE_SIZE: + return _TORCH_BACKEND + + capability = torch.cuda.get_device_capability(X.device) + if capability[0] < 9: + return _TORCH_BACKEND + if not _quack_supports_symmetric_gemm_dtype(X.dtype, capability): + if not self._logged_quack_dtype_fallback: + logger.warning_rank0( + "Gram Newton-Schulz could not use Quack symmetric GEMM for " + f"dtype={X.dtype} on sm_{capability[0]}{capability[1]}; falling back to torch matmuls." + ) + self._logged_quack_dtype_fallback = True + return _TORCH_BACKEND + + try: + return _make_quack_backend() + except Exception as exc: # pragma: no cover - exercised only in incomplete CUDA environments + if not self._logged_quack_fallback: + logger.warning_rank0( + f"Gram Newton-Schulz could not load Quack kernels ({exc!r}); falling back to torch matmuls." + ) + self._logged_quack_fallback = True + return _TORCH_BACKEND + + def _gram_newton_schulz(self, X: Tensor) -> Tensor: + ops = self._select_backend(X) + R = ops.sym_mm(X, X.mT) + + batch_size = R.size(0) + identity = ( + torch.eye(R.size(-1), device=X.device, dtype=X.dtype).unsqueeze(0).expand(batch_size, -1, -1).contiguous() + ) + Q = None + + for iteration, (a, b, c) in enumerate(self.ns_coefficients): + if iteration in self.reset_iterations and iteration != 0: + X = ops.mm(Q, X).contiguous() + R = ops.sym_mm(X, X.mT) + Q = None + + Z = ops.sym_baddbmm(R, R, C=R, alpha=c, beta=b) + if iteration == 0 or iteration in self.reset_iterations: + Q = Z + a * identity + else: + Q = ops.sym_baddbmm(Q, Z, C=Q, beta=a) + + if iteration < len(self.ns_coefficients) - 1 and iteration + 1 not in self.reset_iterations: + RZ = ops.sym_baddbmm(R, Z, C=R, beta=a) + R = ops.sym_baddbmm(Z, RZ, C=RZ, beta=a) + + return ops.mm(Q, X) + + def _standard_newton_schulz(self, X: Tensor) -> Tensor: + ops = self._select_backend(X) + for a, b, c in self.ns_coefficients: + gram_matrix = ops.sym_mm(X, X.mT) + gram_update = ops.sym_baddbmm(gram_matrix, gram_matrix, C=gram_matrix, alpha=c, beta=b) + X = ops.mm_add(gram_update, X, C=X, beta=a) + return X diff --git a/src/xorl/optim/multi_optimizer.py b/src/xorl/optim/multi_optimizer.py new file mode 100644 index 00000000..4d6ce2e2 --- /dev/null +++ b/src/xorl/optim/multi_optimizer.py @@ -0,0 +1,96 @@ +from typing import Any, Dict, List + +import torch +import torch.nn as nn +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_optimizer_state_dict, + set_optimizer_state_dict, +) +from torch.distributed.checkpoint.stateful import Stateful +from torch.optim.optimizer import Optimizer + +from ..utils import logging + + +logger = logging.get_logger(__name__) + + +class MultiOptimizer(Optimizer, Stateful): + """ + A container that handles multiple optimizers (for ep and non-ep parameters when ep+fsdp2 is enabled) + + Mapping of name -> torch.optim.Optimizer with convenience methods. + Compatible with torch.distributed.checkpoint optimizer APIs that accept a Mapping. + + This class is needed for EP+FSDP2 case because EP and non-EP param have different FSDP sharding dimension (dim-0 vs. dim-1). + """ + + def __init__( + self, + root_model: nn.Module, + optimizers: dict, + key_names: list[str], + ): + self.model = root_model + self.optimizers_dict = optimizers + self._is_multi_optimizer: bool = True + self.key_names = key_names + + @property + def param_groups(self) -> List[Dict[str, Any]]: + """Return all param_groups from all internal optimizers.""" + all_groups = [] + for opt in self.optimizers_dict.values(): + all_groups.extend(opt.param_groups) + return all_groups + + @property + def state(self) -> Dict[torch.nn.Parameter, Any]: + """Return merged state dict from all internal optimizers.""" + merged_state: Dict[torch.nn.Parameter, Any] = {} + for opt in self.optimizers_dict.values(): + merged_state.update(opt.state) + return merged_state + + def step(self) -> None: + for opt in self.optimizers_dict.values(): + opt.step() + + def zero_grad(self) -> None: + for opt in self.optimizers_dict.values(): + opt.zero_grad() + + def state_dict(self) -> Dict[str, Any]: + merged: Dict[str, Any] = {} + for name in self.key_names: + opt = self.optimizers_dict.get(name) + sd = get_optimizer_state_dict(self.model, opt, options=StateDictOptions(flatten_optimizer_state_dict=True)) + overlap = set(merged.keys()) & set(sd.keys()) + if overlap: + raise KeyError( + f"Key clash detected while merging state dict for optimizer '{name}': {', '.join(sorted(overlap))}" + ) + else: + logger.info_rank0( + f"MultiOptimizer merged '{name}' state dict ({len(sd)} keys, total {len(merged) + len(sd)})" + ) + merged.update(sd) + + return merged + + def load_state_dict(self, state_dict: Dict[str, Any]) -> None: + for name in self.key_names: + opt = self.optimizers_dict.get(name) + set_optimizer_state_dict( + self.model, + opt, + optim_state_dict=state_dict, + options=StateDictOptions(flatten_optimizer_state_dict=True), + ) + + def register_step_pre_hook(self, hook): + return [opt.register_step_pre_hook(hook) for opt in self.optimizers_dict.values()] + + def __len__(self) -> int: + return len(self.optimizers_dict) diff --git a/src/xorl/optim/muon.py b/src/xorl/optim/muon.py index 90ba0b01..1f84e1f2 100644 --- a/src/xorl/optim/muon.py +++ b/src/xorl/optim/muon.py @@ -4,7 +4,7 @@ Extends ``torch.optim.Muon`` with: - Mixed param groups: ``use_muon=True`` (Newton-Schulz) / ``False`` (AdamW fallback) - FSDP2/EP DTensor support (shard-local Newton-Schulz) - - 3D+ MoE expert tensor support (reshape to 2D for NS) + - 3D+ MoE expert tensor support (preserve leading dims as matrix batches) The core Muon algorithm is aligned with PyTorch's implementation: B_t = (1-μ) * g_t + μ * B_{t-1} (EMA momentum; skipped if momentum=0) @@ -22,6 +22,8 @@ - PyTorch torch.optim.Muon """ +from collections import defaultdict +from dataclasses import dataclass from typing import Iterable, Optional, Tuple import torch @@ -31,10 +33,30 @@ from torch.optim.optimizer import Optimizer from ..utils import logging +from .gram_newton_schulz import GramNewtonSchulzOrthogonalizer, expand_ns_coefficients, find_best_restarts logger = logging.get_logger(__name__) +GROUPED_GRAM_NS_FP32_BYTE_LIMIT = 2 * 1024**3 + + +@dataclass +class _MuonUpdatePlan: + param: torch.Tensor + adjusted_lr: float + orig_shape: Optional[torch.Size] + pieces: list[Optional[torch.Tensor]] + + +@dataclass +class _GroupedOrthogonalizationEntry: + plan: _MuonUpdatePlan + piece_index: int + tensor: torch.Tensor + batched_tensor: torch.Tensor + transposed_for_batching: bool + class Muon(TorchMuon): """ @@ -57,10 +79,31 @@ class Muon(TorchMuon): ``0.2 * sqrt(max(A, B))`` so Muon can reuse AdamW hyperparams. adamw_betas: Beta coefficients for AdamW groups. adamw_eps: Epsilon for AdamW groups. + grad_dtype: If set, force the gradient tensor used inside Muon to this + dtype before momentum/update construction. momentum_dtype: If set, force Muon momentum buffers to this dtype (e.g. ``torch.bfloat16``). Saves memory when params/grads are fp32 but bf16 momentum is sufficient for Newton-Schulz. Default ``None`` inherits dtype from the gradient. + update_dtype: If set, force the transient update tensor passed into + Newton-Schulz to this dtype. This decouples Muon compute dtype from + the stored gradient or momentum-buffer dtype. + force_momentum_path: If true, still route update construction through + the momentum-buffer path when ``momentum=0``. Intended for + debugging/ablations that separate path effects from optimizer + coefficient effects. + ns_algorithm: Newton-Schulz backend. ``"standard_newton_schulz"`` + preserves the current PyTorch Muon path. ``"gram_newton_schulz"`` + enables the Dao-AILab Gram Newton-Schulz update path. + ns_use_quack_kernels: Whether Gram Newton-Schulz may use Quack GEMM + kernels on SM90+/SM100 devices. Falls back to torch matmuls when + unavailable or unsupported. + gram_newton_schulz_num_restarts: If using Gram Newton-Schulz and + explicit restart iterations are not provided, autotune this many + restart locations from the chosen coefficients. + gram_newton_schulz_restart_iterations: Explicit Gram Newton-Schulz + restart iteration indices. A value of ``2`` means restart after + finishing the second iteration. adamw_state_dtype: If set, force AdamW fallback optimizer states (``exp_avg``, ``exp_avg_sq``) to this dtype (e.g. ``torch.bfloat16``). Default ``None`` inherits dtype from @@ -79,11 +122,32 @@ def __init__( adamw_betas: Tuple[float, float] = (0.9, 0.95), adamw_eps: float = 1e-8, momentum_dtype: Optional[torch.dtype] = None, + grad_dtype: Optional[torch.dtype] = None, + update_dtype: Optional[torch.dtype] = None, + force_momentum_path: bool = False, + ns_algorithm: str = "standard_newton_schulz", + ns_use_quack_kernels: bool = True, + gram_newton_schulz_num_restarts: int = 1, + gram_newton_schulz_restart_iterations: Optional[Iterable[int]] = None, adamw_state_dtype: Optional[torch.dtype] = None, ): + if ns_algorithm not in {"standard_newton_schulz", "gram_newton_schulz"}: + raise ValueError( + f"Unsupported Muon ns_algorithm: {ns_algorithm!r}. " + "Expected 'standard_newton_schulz' or 'gram_newton_schulz'." + ) + if gram_newton_schulz_num_restarts < 0: + raise ValueError( + f"gram_newton_schulz_num_restarts must be non-negative, got {gram_newton_schulz_num_restarts}" + ) + self._momentum_dtype = momentum_dtype + self._grad_dtype = grad_dtype + self._update_dtype = update_dtype + self._force_momentum_path = force_momentum_path self._adamw_state_dtype = adamw_state_dtype self._logged_dtypes = False + self._gram_ns_orthogonalizers = {} # Skip TorchMuon.__init__ (which enforces 2D-only) and call # Optimizer.__init__ directly so we can accept 3D MoE params and # non-Muon (AdamW) param groups. @@ -96,6 +160,14 @@ def __init__( eps=1e-7, ns_steps=ns_steps, adjust_lr_fn=adjust_lr_fn, + ns_algorithm=ns_algorithm, + ns_use_quack_kernels=ns_use_quack_kernels, + gram_newton_schulz_num_restarts=gram_newton_schulz_num_restarts, + gram_newton_schulz_restart_iterations=( + tuple(gram_newton_schulz_restart_iterations) + if gram_newton_schulz_restart_iterations is not None + else None + ), use_muon=True, adamw_betas=adamw_betas, adamw_eps=adamw_eps, @@ -140,6 +212,10 @@ def _muon_step(self, group: dict) -> None: eps = group["eps"] weight_decay = group["weight_decay"] adjust_lr_fn = group["adjust_lr_fn"] + uses_grouped_gram_ns = group["ns_algorithm"] == "gram_newton_schulz" + grouped_updates: dict[tuple[tuple[int, int], torch.dtype, torch.device], list[_GroupedOrthogonalizationEntry]] + grouped_updates = defaultdict(list) + update_plans: list[_MuonUpdatePlan] = [] for p in group["params"]: if p.grad is None: @@ -154,9 +230,12 @@ def _muon_step(self, group: dict) -> None: else: grad_local = grad p_local = p.data + raw_grad_dtype = grad_local.dtype + if self._grad_dtype is not None and grad_local.dtype != self._grad_dtype: + grad_local = grad_local.to(self._grad_dtype) - # Handle 3D+ MoE expert tensors: [num_experts, hidden, intermediate] - # For fused gate_up_proj [E, H, 2I], split into two halves for NS + # Handle 3D+ tensors as batches of matrices: [..., hidden, intermediate]. + # For fused gate_up_proj [E, H, 2I], split into two [..., H, I] halves. orig_shape = None fused_split = None if grad_local.ndim >= 3: @@ -164,15 +243,18 @@ def _muon_step(self, group: dict) -> None: fused_gate_up_ids = group.get("_fused_gate_up_ids", set()) if id(p) in fused_gate_up_ids: fused_split = grad_local.shape[-1] // 2 - grad_local = grad_local.reshape(-1, grad_local.shape[-1]) + grad_local = grad_local.reshape(-1, *grad_local.shape[-2:]) # --- PyTorch Muon algorithm (aligned with torch.optim._muon) --- - if momentum == 0: + if momentum == 0 and not self._force_momentum_path: # No momentum: apply NS directly to the raw gradient, no buffer needed. update = grad_local + if self._update_dtype is not None and update.dtype != self._update_dtype: + update = update.to(self._update_dtype) if not self._logged_dtypes: logger.info_rank0( - f"Muon dtypes (no momentum): param={p_local.dtype}, grad={grad_local.dtype} " + f"Muon dtypes (no momentum): param={p_local.dtype}, raw_grad={raw_grad_dtype}, " + f"grad={grad_local.dtype}, update={update.dtype} " f"(shape={list(grad_local.shape)})" ) self._logged_dtypes = True @@ -184,12 +266,6 @@ def _muon_step(self, group: dict) -> None: grad_local, dtype=buf_dtype, ) - if not self._logged_dtypes: - logger.info_rank0( - f"Muon dtypes: param={p_local.dtype}, grad={grad_local.dtype}, " - f"momentum={buf_dtype} (shape={list(grad_local.shape)})" - ) - self._logged_dtypes = True buf = state["momentum_buffer"] # EMA momentum: B = (1-μ)*g + μ*B @@ -197,33 +273,183 @@ def _muon_step(self, group: dict) -> None: buf.lerp_(grad_local.to(buf.dtype), 1 - momentum) # Nesterov: ~B = g + μ*B (or just B if nesterov=False) - # Work in buf dtype (may be bf16 even if grad is fp32) update = grad_local.to(buf.dtype).lerp(buf, momentum) if nesterov else buf + if self._update_dtype is not None and update.dtype != self._update_dtype: + update = update.to(self._update_dtype) + if not self._logged_dtypes: + logger.info_rank0( + f"Muon dtypes: param={p_local.dtype}, raw_grad={raw_grad_dtype}, grad={grad_local.dtype}, " + f"momentum={buf.dtype}, update={update.dtype}, " + f"force_momentum_path={self._force_momentum_path} " + f"(shape={list(grad_local.shape)})" + ) + self._logged_dtypes = True - # Newton-Schulz orthogonalization - if fused_split is not None: - u1 = _zeropower_via_newtonschulz(update[..., :fused_split], ns_coefficients, ns_steps, eps) - u2 = _zeropower_via_newtonschulz(update[..., fused_split:], ns_coefficients, ns_steps, eps) - update = torch.cat([u1, u2], dim=-1) - del u1, u2 + adjusted_lr = _adjust_lr( + lr, + adjust_lr_fn, + grad_local.shape[-2:] if grad_local.ndim > 2 else grad_local.shape, + ) + pieces = [update[..., :fused_split], update[..., fused_split:]] if fused_split is not None else [update] + plan = _MuonUpdatePlan( + param=p_local, + adjusted_lr=adjusted_lr, + orig_shape=orig_shape, + pieces=[None] * len(pieces), + ) + update_plans.append(plan) + + for piece_index, piece in enumerate(pieces): + if uses_grouped_gram_ns: + batched_piece = piece.reshape(-1, *piece.shape[-2:]) if piece.ndim > 2 else piece.unsqueeze(0) + transposed_for_batching = batched_piece.shape[-2] > batched_piece.shape[-1] + if transposed_for_batching: + batched_piece = batched_piece.mT + batch_key = (tuple(batched_piece.shape[-2:]), piece.dtype, piece.device) + grouped_updates[batch_key].append( + _GroupedOrthogonalizationEntry( + plan=plan, + piece_index=piece_index, + tensor=piece, + batched_tensor=batched_piece, + transposed_for_batching=transposed_for_batching, + ) + ) + else: + plan.pieces[piece_index] = self._orthogonalize_update(piece, group, ns_coefficients, ns_steps, eps) + + if uses_grouped_gram_ns and grouped_updates: + orthogonalizer = self._get_gram_ns_orthogonalizer(group) + self._orthogonalize_grouped_gram_ns_updates(grouped_updates, orthogonalizer) + + for plan in update_plans: + update_pieces = plan.pieces + if any(piece is None for piece in update_pieces): + raise RuntimeError("Grouped Muon orthogonalization left an update piece uninitialized") + if len(update_pieces) == 1: + update = update_pieces[0] else: - update = _zeropower_via_newtonschulz(update, ns_coefficients, ns_steps, eps) - - # LR adjustment based on the 2D shape fed to NS - adjusted_lr = _adjust_lr(lr, adjust_lr_fn, grad_local.shape) + update = torch.cat(update_pieces, dim=-1) # Restore original shape - if orig_shape is not None: - update = update.reshape(orig_shape) + if plan.orig_shape is not None: + update = update.reshape(plan.orig_shape) # Cast back to param dtype - update = update.to(p_local.dtype) + update = update.to(plan.param.dtype) # Decoupled weight decay - p_local.mul_(1 - lr * weight_decay) + plan.param.mul_(1 - lr * weight_decay) # Parameter update - p_local.add_(update, alpha=-adjusted_lr) + plan.param.add_(update, alpha=-plan.adjusted_lr) + + def _orthogonalize_update( + self, + update: torch.Tensor, + group: dict, + ns_coefficients: Tuple[float, float, float], + ns_steps: int, + eps: float, + ) -> torch.Tensor: + if group["ns_algorithm"] == "standard_newton_schulz": + if update.ndim <= 2: + return _zeropower_via_newtonschulz(update, ns_coefficients, ns_steps, eps) + + original_shape = update.shape + flat_update = update.reshape(-1, *update.shape[-2:]) + orthogonalized = [ + _zeropower_via_newtonschulz(matrix, ns_coefficients, ns_steps, eps) for matrix in flat_update.unbind(0) + ] + return torch.stack(orthogonalized, dim=0).reshape(original_shape) + if group["ns_algorithm"] == "gram_newton_schulz": + return self._get_gram_ns_orthogonalizer(group).orthogonalize(update) + raise ValueError( + f"Unsupported Muon ns_algorithm: {group['ns_algorithm']!r}. " + "Expected 'standard_newton_schulz' or 'gram_newton_schulz'." + ) + + def _orthogonalize_grouped_gram_ns_updates( + self, + grouped_updates: dict[tuple[tuple[int, int], torch.dtype, torch.device], list[_GroupedOrthogonalizationEntry]], + orthogonalizer: GramNewtonSchulzOrthogonalizer, + ) -> None: + for batch_entries in grouped_updates.values(): + for chunk in self._iter_grouped_gram_ns_chunks(batch_entries): + if len(chunk) == 1 and chunk[0].batched_tensor.shape[0] == 1 and chunk[0].tensor.ndim == 2: + entry = chunk[0] + entry.plan.pieces[entry.piece_index] = orthogonalizer.orthogonalize(entry.tensor) + continue + + stacked = torch.cat([entry.batched_tensor for entry in chunk], dim=0) + outputs = orthogonalizer.orthogonalize(stacked) + output_offset = 0 + for entry in chunk: + batch_size = entry.batched_tensor.shape[0] + output = outputs[output_offset : output_offset + batch_size] + output_offset += batch_size + if entry.transposed_for_batching: + output = output.mT + if entry.tensor.ndim == 2: + entry.plan.pieces[entry.piece_index] = output.squeeze(0) + else: + entry.plan.pieces[entry.piece_index] = output.reshape(entry.tensor.shape) + + def _iter_grouped_gram_ns_chunks( + self, + batch_entries: list[_GroupedOrthogonalizationEntry], + ): + per_matrix_numel = batch_entries[0].batched_tensor.shape[-2] * batch_entries[0].batched_tensor.shape[-1] + max_matrix_batch = max(1, GROUPED_GRAM_NS_FP32_BYTE_LIMIT // (per_matrix_numel * 4)) + chunk: list[_GroupedOrthogonalizationEntry] = [] + chunk_matrix_batch = 0 + + for entry in batch_entries: + entry_matrix_batch = entry.batched_tensor.shape[0] + if chunk and chunk_matrix_batch + entry_matrix_batch > max_matrix_batch: + yield chunk + chunk = [] + chunk_matrix_batch = 0 + + chunk.append(entry) + chunk_matrix_batch += entry_matrix_batch + + if chunk: + yield chunk + + def _get_gram_ns_orthogonalizer(self, group: dict) -> GramNewtonSchulzOrthogonalizer: + expanded_coefficients = tuple(expand_ns_coefficients(group["ns_coefficients"], group["ns_steps"])) + restart_iterations = group["gram_newton_schulz_restart_iterations"] + if restart_iterations is not None: + restart_iterations = tuple(int(i) for i in restart_iterations) + else: + restart_iterations = tuple( + find_best_restarts( + expanded_coefficients, + num_restarts=group["gram_newton_schulz_num_restarts"], + ) + ) + + cache_key = ( + expanded_coefficients, + float(group["eps"]), + bool(group["ns_use_quack_kernels"]), + restart_iterations, + ) + orthogonalizer = self._gram_ns_orthogonalizers.get(cache_key) + if orthogonalizer is None: + orthogonalizer = GramNewtonSchulzOrthogonalizer( + ns_coefficients=expanded_coefficients, + ns_epsilon=group["eps"], + ns_use_quack_kernels=group["ns_use_quack_kernels"], + gram_newton_schulz_restart_iterations=restart_iterations, + ) + self._gram_ns_orthogonalizers[cache_key] = orthogonalizer + logger.info_rank0( + "Initialized Muon Gram Newton-Schulz " + f"(restarts={list(restart_iterations)}, quack_kernels={group['ns_use_quack_kernels']})" + ) + return orthogonalizer def _adamw_step(self, group: dict) -> None: """Standard AdamW update for non-Muon parameters.""" diff --git a/src/xorl/optim/optimizer.py b/src/xorl/optim/optimizer.py index 72a4165e..84ddf988 100644 --- a/src/xorl/optim/optimizer.py +++ b/src/xorl/optim/optimizer.py @@ -6,18 +6,15 @@ import torch import torch.nn as nn from torch.distributed._tensor import DTensor -from torch.distributed.checkpoint.state_dict import ( - StateDictOptions, - get_optimizer_state_dict, - set_optimizer_state_dict, -) -from torch.distributed.checkpoint.stateful import Stateful from torch.optim import AdamW from torch.optim.optimizer import Optimizer from ..distributed.parallel_state import get_parallel_state from ..utils import logging +from .anyprecision_adamw import AnyPrecisionAdamW +from .multi_optimizer import MultiOptimizer from .muon import Muon +from .signsgd import SignSGD logger = logging.get_logger(__name__) @@ -28,244 +25,6 @@ _MUON_EXCLUDE_PARAM_PATTERNS = {"embed_tokens", "lm_head", "norm", "gate.weight"} -# https://github.com/meta-llama/llama-recipes/blob/v0.0.4/src/llama_recipes/policies/anyprecision_optimizer.py -class SignSGD(Optimizer): - """Sign-based SGD optimizer with no optimizer-state tensors.""" - - def __init__( - self, - params, - lr: float = 1e-3, - weight_decay: float = 0.0, - ): - if lr < 0.0: - raise ValueError(f"Invalid learning rate: {lr}") - if weight_decay < 0.0: - raise ValueError(f"Invalid weight_decay: {weight_decay}") - - defaults = { - "lr": lr, - "weight_decay": weight_decay, - } - super().__init__(params, defaults) - - @torch.no_grad() - def step(self, closure=None): - """Perform a single optimization step.""" - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - for group in self.param_groups: - lr = group["lr"] - weight_decay = group["weight_decay"] - - for p in group["params"]: - grad = p.grad - if grad is None: - continue - if grad.is_sparse: - raise RuntimeError("SignSGD does not support sparse gradients.") - - if weight_decay: - p.add_(p, alpha=-lr * weight_decay) - - p.add_(torch.sign(grad), alpha=-lr) - - return loss - - -class AnyPrecisionAdamW(Optimizer): - def __init__( - self, - params, - lr=1e-3, - betas=(0.9, 0.95), - eps=1e-8, - weight_decay=0.0, - use_kahan_summation=False, - momentum_dtype=torch.bfloat16, - variance_dtype=torch.bfloat16, - compensation_buffer_dtype=torch.bfloat16, - ): - defaults = { - "lr": lr, - "betas": betas, - "eps": eps, - "weight_decay": weight_decay, - "use_kahan_summation": use_kahan_summation, - "momentum_dtype": momentum_dtype, - "variance_dtype": variance_dtype, - "compensation_buffer_dtype": compensation_buffer_dtype, - } - super().__init__(params, defaults) - - @torch.no_grad() - def step(self, closure=None): - """ - Performs a single optimization step. - - Args: - closure (callable, optional): A closure that reevaluates the model and returns the loss. - """ - - if closure is not None: - with torch.enable_grad(): - closure() - - for group in self.param_groups: - beta1, beta2 = group["betas"] - lr = group["lr"] - weight_decay = group["weight_decay"] - eps = group["eps"] - use_kahan_summation = group["use_kahan_summation"] - - momentum_dtype = group["momentum_dtype"] - variance_dtype = group["variance_dtype"] - compensation_buffer_dtype = group["compensation_buffer_dtype"] - for p in group["params"]: - if p.grad is None: - continue - - if p.grad.is_sparse: - raise RuntimeError("AnyPrecisionAdamW does not support sparse gradients.") - - state = self.state[p] - # State initialization - if len(state) == 0: - state["step"] = torch.tensor(0.0) - - # momentum - EMA of gradient values - state["exp_avg"] = torch.zeros_like(p, dtype=momentum_dtype) - - # variance uncentered - EMA of squared gradient values - state["exp_avg_sq"] = torch.zeros_like(p, dtype=variance_dtype) - - # optional Kahan summation - accumulated error tracker - if use_kahan_summation: - state["compensation"] = torch.zeros_like(p, dtype=compensation_buffer_dtype) - - # Main processing - # update the steps for each param group update - state["step"] += 1 - step = state["step"] - - exp_avg = state["exp_avg"] - exp_avg_sq = state["exp_avg_sq"] - grad = p.grad - - if weight_decay: # weight decay, AdamW style - p.data.mul_(1 - lr * weight_decay) - - exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) # update momentum - exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) # update uncentered variance - - bias_correction1 = 1 - beta1**step # adjust using bias1 - step_size = lr / bias_correction1 - - denom_correction = (1 - beta2**step) ** 0.5 # adjust using bias2 and avoids math import - centered_variance = (exp_avg_sq.sqrt() / denom_correction).add_(eps, alpha=1) - - if use_kahan_summation: # lr update to compensation - compensation = state["compensation"] - compensation.addcdiv_(exp_avg, centered_variance, value=-step_size) - - # update weights with compensation (Kahan summation) - # save error back to compensation for next iteration - temp_buffer = p.detach().clone() - p.data.add_(compensation) - compensation.add_(temp_buffer.sub_(p.data)) - else: # usual AdamW updates - p.data.addcdiv_(exp_avg, centered_variance, value=-step_size) - - -class MultiOptimizer(Optimizer, Stateful): - """ - A container that handles multiple optimizers (for ep and non-ep parameters when ep+fsdp2 is enabled) - - Mapping of name -> torch.optim.Optimizer with convenience methods. - Compatible with torch.distributed.checkpoint optimizer APIs that accept a Mapping. - - This class is needed for EP+FSDP2 case because EP and non-EP param have different FSDP sharding dimension (dim-0 vs. dim-1). - """ - - def __init__( - self, - root_model: nn.Module, - optimizers: dict, # {"ep": opt1, "non_ep": opt2} - key_names: list[str], - ): - self.model = root_model - self.optimizers_dict = optimizers - self._is_multi_optimizer: bool = True - self.key_names = key_names - - @property - def param_groups(self) -> List[Dict[str, Any]]: - """Return all param_groups from all internal optimizers.""" - all_groups = [] - for opt in self.optimizers_dict.values(): - all_groups.extend(opt.param_groups) - return all_groups - - @property - def state(self) -> Dict[torch.nn.Parameter, Any]: - """Return merged state dict from all internal optimizers.""" - merged_state: Dict[torch.nn.Parameter, Any] = {} - for opt in self.optimizers_dict.values(): - merged_state.update(opt.state) - return merged_state - - def step(self) -> None: - for opt in self.optimizers_dict.values(): - opt.step() - - def zero_grad(self) -> None: - for opt in self.optimizers_dict.values(): - opt.zero_grad() - - def state_dict( - self, - ) -> Dict[str, Any]: - # get the flatten state dict for multi-optimizer - merged: Dict[str, Any] = {} - for name in self.key_names: - opt = self.optimizers_dict.get(name) - sd = get_optimizer_state_dict(self.model, opt, options=StateDictOptions(flatten_optimizer_state_dict=True)) - # check for key clashes before merging - overlap = set(merged.keys()) & set(sd.keys()) - if overlap: - raise KeyError( - f"Key clash detected while merging state dict for optimizer '{name}': {', '.join(sorted(overlap))}" - ) - else: - logger.info_rank0( - f"MultiOptimizer merged '{name}' state dict ({len(sd)} keys, total {len(merged) + len(sd)})" - ) - merged.update(sd) - - return merged - - def load_state_dict(self, state_dict: Dict[str, Any]) -> None: - # Feed the same merged flattened dict to each sub-optimizer; PyTorch will - # pick out only the entries for parameters that belong to that optimizer. - for name in self.key_names: - opt = self.optimizers_dict.get(name) - set_optimizer_state_dict( - self.model, - opt, - optim_state_dict=state_dict, - options=StateDictOptions(flatten_optimizer_state_dict=True), - ) - - def register_step_pre_hook(self, hook): - return [opt.register_step_pre_hook(hook) for opt in self.optimizers_dict.values()] - - def __len__(self) -> int: - return len(self.optimizers_dict) - - def _should_build_ep_aware(model: "nn.Module", param_groups: Optional[Sequence[Dict[str, Any]]]) -> bool: # Only auto-split when using FSDP2 with EP and no explicit param_groups if param_groups is not None: @@ -337,6 +96,14 @@ def get_parameter_names(model, forbidden_layer_types, forbidden_param_names): } +def _normalize_optional_dtype(dtype: Any, *, field_name: str) -> Optional[torch.dtype]: + if dtype is None or isinstance(dtype, torch.dtype): + return dtype + if dtype in _ANYPRECISION_STATE_DTYPES: + return _ANYPRECISION_STATE_DTYPES[dtype] + raise ValueError(f"Unsupported {field_name}: {dtype!r}. Expected 'fp32', 'bf16', a torch.dtype, or None.") + + def _get_optimizer_cls_and_kwargs( optimizer_type: str, *, @@ -386,6 +153,9 @@ def _get_optimizer_cls_and_kwargs( return SignSGD, ctor_kwargs elif optimizer_type == "muon": adamw_state_dtype = _ANYPRECISION_STATE_DTYPES.get(optimizer_dtype) + momentum_dtype = kwargs.get("muon_momentum_dtype") + if momentum_dtype is None: + momentum_dtype = _ANYPRECISION_STATE_DTYPES.get(optimizer_dtype) ctor_kwargs = dict( lr=kwargs.get("muon_lr", 0.02), momentum=kwargs.get("muon_momentum", 0.95), @@ -393,9 +163,16 @@ def _get_optimizer_cls_and_kwargs( ns_steps=kwargs.get("muon_ns_steps", 5), weight_decay=weight_decay, adjust_lr_fn=kwargs.get("muon_adjust_lr_fn"), + ns_algorithm=kwargs.get("muon_ns_algorithm", "standard_newton_schulz"), + ns_use_quack_kernels=kwargs.get("muon_ns_use_quack_kernels", True), + gram_newton_schulz_num_restarts=kwargs.get("muon_gram_ns_num_restarts", 1), + gram_newton_schulz_restart_iterations=kwargs.get("muon_gram_ns_restart_iterations"), adamw_betas=betas, adamw_eps=eps, - momentum_dtype=kwargs.get("muon_momentum_dtype"), + momentum_dtype=_normalize_optional_dtype(momentum_dtype, field_name="muon_momentum_dtype"), + grad_dtype=_normalize_optional_dtype(kwargs.get("muon_grad_dtype"), field_name="muon_grad_dtype"), + update_dtype=_normalize_optional_dtype(kwargs.get("muon_update_dtype"), field_name="muon_update_dtype"), + force_momentum_path=kwargs.get("muon_force_momentum_path", False), adamw_state_dtype=adamw_state_dtype, ) return Muon, ctor_kwargs @@ -575,7 +352,11 @@ def build_optimizer( - sgd: {"momentum": 0.9, "nesterov": True} - signsgd: no optimizer-specific kwargs - muon: {"muon_lr": 0.02, "muon_momentum": 0.95, "muon_nesterov": True, - "muon_ns_steps": 5, "muon_adjust_lr_fn": None, "muon_momentum_dtype": None} + "muon_ns_steps": 5, "muon_adjust_lr_fn": None, + "muon_ns_algorithm": "standard_newton_schulz", + "muon_ns_use_quack_kernels": True, "muon_gram_ns_num_restarts": 1, + "muon_gram_ns_restart_iterations": None, "muon_momentum_dtype": None, + "muon_grad_dtype": None, "muon_update_dtype": None, "muon_force_momentum_path": False} - adamw/anyprecision_adamw: any extra kwargs forwarded to constructor """ # EP-aware routing: for FSDP2+EP, split params into EP and non-EP groups and build two optimizers. diff --git a/src/xorl/optim/signsgd.py b/src/xorl/optim/signsgd.py new file mode 100644 index 00000000..9de7eab7 --- /dev/null +++ b/src/xorl/optim/signsgd.py @@ -0,0 +1,50 @@ +import torch +from torch.optim.optimizer import Optimizer + + +# https://github.com/meta-llama/llama-recipes/blob/v0.0.4/src/llama_recipes/policies/anyprecision_optimizer.py +class SignSGD(Optimizer): + """Sign-based SGD optimizer with no optimizer-state tensors.""" + + def __init__( + self, + params, + lr: float = 1e-3, + weight_decay: float = 0.0, + ): + if lr < 0.0: + raise ValueError(f"Invalid learning rate: {lr}") + if weight_decay < 0.0: + raise ValueError(f"Invalid weight_decay: {weight_decay}") + + defaults = { + "lr": lr, + "weight_decay": weight_decay, + } + super().__init__(params, defaults) + + @torch.no_grad() + def step(self, closure=None): + """Perform a single optimization step.""" + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + lr = group["lr"] + weight_decay = group["weight_decay"] + + for p in group["params"]: + grad = p.grad + if grad is None: + continue + if grad.is_sparse: + raise RuntimeError("SignSGD does not support sparse gradients.") + + if weight_decay: + p.add_(p, alpha=-lr * weight_decay) + + p.add_(torch.sign(grad), alpha=-lr) + + return loss diff --git a/src/xorl/qlora/__init__.py b/src/xorl/qlora/__init__.py index 0d273463..2e9c889f 100644 --- a/src/xorl/qlora/__init__.py +++ b/src/xorl/qlora/__init__.py @@ -1,25 +1,4 @@ -"""Lazy exports for QLoRA to avoid package-level circular imports.""" - -_MODULE_ATTRS = { - "BlockFP8QLoRALinear", - "NF4QLoRALinear", - "NvFP4QLoRALinear", - "QLoRALinear", - "QLoRAMoeExperts", - "prefetch_aqn_noise", -} - -_UTIL_ATTRS = { - "detect_prequantized_block_fp8", - "detect_prequantized_nvfp4", - "get_prequantized_exclude_modules", - "inject_qlora_into_model", - "maybe_load_and_quantize_moe_qlora", - "maybe_load_prequantized_qlora", - "maybe_quantize_qlora", - "maybe_requant_qlora", - "save_qlora_checkpoint", -} +import importlib __all__ = [ @@ -37,62 +16,33 @@ "detect_prequantized_nvfp4", "detect_prequantized_block_fp8", "maybe_load_prequantized_qlora", - "get_prequantized_exclude_modules", ] -def __getattr__(name): - if name in _MODULE_ATTRS: - from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear # noqa: PLC0415 - from xorl.qlora.modules.linear import QLoRALinear, prefetch_aqn_noise # noqa: PLC0415 - from xorl.qlora.modules.moe_experts import QLoRAMoeExperts # noqa: PLC0415 - from xorl.qlora.modules.nf4_linear import NF4QLoRALinear # noqa: PLC0415 - from xorl.qlora.modules.nvfp4_linear import NvFP4QLoRALinear # noqa: PLC0415 - - globals().update( - { - "QLoRALinear": QLoRALinear, - "NvFP4QLoRALinear": NvFP4QLoRALinear, - "BlockFP8QLoRALinear": BlockFP8QLoRALinear, - "NF4QLoRALinear": NF4QLoRALinear, - "QLoRAMoeExperts": QLoRAMoeExperts, - "prefetch_aqn_noise": prefetch_aqn_noise, - } - ) - return globals()[name] - - if name in _UTIL_ATTRS: - from xorl.qlora.utils import ( # noqa: PLC0415 - detect_prequantized_block_fp8, - detect_prequantized_nvfp4, - inject_qlora_into_model, - maybe_load_and_quantize_moe_qlora, - maybe_load_prequantized_qlora, - maybe_quantize_qlora, - maybe_requant_qlora, - save_qlora_checkpoint, - ) - - globals().update( - { - "detect_prequantized_block_fp8": detect_prequantized_block_fp8, - "detect_prequantized_nvfp4": detect_prequantized_nvfp4, - "inject_qlora_into_model": inject_qlora_into_model, - "maybe_load_and_quantize_moe_qlora": maybe_load_and_quantize_moe_qlora, - "maybe_load_prequantized_qlora": maybe_load_prequantized_qlora, - "maybe_quantize_qlora": maybe_quantize_qlora, - "maybe_requant_qlora": maybe_requant_qlora, - "save_qlora_checkpoint": save_qlora_checkpoint, - } - ) - - if "get_prequantized_exclude_modules" not in globals(): - from xorl.models.checkpoint_handlers.buffers import ( # noqa: PLC0415 - get_prequantized_exclude_modules, - ) - - globals()["get_prequantized_exclude_modules"] = get_prequantized_exclude_modules +_QLORA_ATTRS = { + "BlockFP8QLoRALinear": ("xorl.qlora.modules.block_fp8_linear", "BlockFP8QLoRALinear"), + "NF4QLoRALinear": ("xorl.qlora.modules.nf4_linear", "NF4QLoRALinear"), + "NvFP4QLoRALinear": ("xorl.qlora.modules.nvfp4_linear", "NvFP4QLoRALinear"), + "QLoRALinear": ("xorl.qlora.modules.linear", "QLoRALinear"), + "QLoRAMoeExperts": ("xorl.qlora.modules.moe_experts", "QLoRAMoeExperts"), + "prefetch_aqn_noise": ("xorl.qlora.modules.linear", "prefetch_aqn_noise"), + "detect_prequantized_block_fp8": ("xorl.qlora.utils", "detect_prequantized_block_fp8"), + "detect_prequantized_nvfp4": ("xorl.qlora.utils", "detect_prequantized_nvfp4"), + "inject_qlora_into_model": ("xorl.qlora.utils", "inject_qlora_into_model"), + "maybe_load_and_quantize_moe_qlora": ("xorl.qlora.utils", "maybe_load_and_quantize_moe_qlora"), + "maybe_load_prequantized_qlora": ("xorl.qlora.utils", "maybe_load_prequantized_qlora"), + "maybe_quantize_qlora": ("xorl.qlora.utils", "maybe_quantize_qlora"), + "maybe_requant_qlora": ("xorl.qlora.utils", "maybe_requant_qlora"), + "save_qlora_checkpoint": ("xorl.qlora.utils", "save_qlora_checkpoint"), +} - return globals()[name] - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") +def __getattr__(name): + if name not in _QLORA_ATTRS: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + module_name, attr_name = _QLORA_ATTRS[name] + module = importlib.import_module(module_name) + value = getattr(module, attr_name) + globals()[name] = value + return value diff --git a/src/xorl/qlora/modules/__init__.py b/src/xorl/qlora/modules/__init__.py index c5c88910..78336c9a 100644 --- a/src/xorl/qlora/modules/__init__.py +++ b/src/xorl/qlora/modules/__init__.py @@ -1,8 +1,4 @@ -from .block_fp8_linear import BlockFP8QLoRALinear -from .linear import QLoRALinear, prefetch_aqn_noise -from .moe_experts import QLoRAMoeExperts -from .nf4_linear import NF4QLoRALinear -from .nvfp4_linear import NvFP4QLoRALinear +import importlib __all__ = [ @@ -13,3 +9,24 @@ "QLoRAMoeExperts", "prefetch_aqn_noise", ] + + +_MODULE_ATTRS = { + "BlockFP8QLoRALinear": ("xorl.qlora.modules.block_fp8_linear", "BlockFP8QLoRALinear"), + "NF4QLoRALinear": ("xorl.qlora.modules.nf4_linear", "NF4QLoRALinear"), + "NvFP4QLoRALinear": ("xorl.qlora.modules.nvfp4_linear", "NvFP4QLoRALinear"), + "QLoRALinear": ("xorl.qlora.modules.linear", "QLoRALinear"), + "QLoRAMoeExperts": ("xorl.qlora.modules.moe_experts", "QLoRAMoeExperts"), + "prefetch_aqn_noise": ("xorl.qlora.modules.linear", "prefetch_aqn_noise"), +} + + +def __getattr__(name): + if name not in _MODULE_ATTRS: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + module_name, attr_name = _MODULE_ATTRS[name] + module = importlib.import_module(module_name) + value = getattr(module, attr_name) + globals()[name] = value + return value diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index 477d1a0b..dec1ef9c 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -13,6 +13,7 @@ See xorl.server.runner.runner_dispatcher for the entry point. """ +import gc import logging import os import time @@ -563,7 +564,20 @@ def _initialize_optimizer(self): if optimizer_type == "muon": optimizer_kwargs = { k: self.train_config[k] - for k in ("muon_lr", "muon_momentum", "muon_nesterov", "muon_ns_steps", "muon_adjust_lr_fn") + for k in ( + "muon_lr", + "muon_momentum", + "muon_nesterov", + "muon_ns_steps", + "muon_adjust_lr_fn", + "muon_ns_algorithm", + "muon_ns_use_quack_kernels", + "muon_gram_ns_num_restarts", + "muon_gram_ns_restart_iterations", + "muon_grad_dtype", + "muon_update_dtype", + "muon_force_momentum_path", + ) if k in self.train_config } self.optimizer = build_optimizer( @@ -1261,8 +1275,6 @@ def forward_backward( # After weight-sync + optim_step from the previous step the CUDA # allocator can have many small free blocks; without this, CUBLAS # handle creation or Triton autotuner workspace allocs can fail. - import gc - gc.collect() torch.cuda.empty_cache() diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index 173ca147..c2f452ab 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -278,6 +278,59 @@ class ServerArguments: }, ) + muon_ns_algorithm: Literal["standard_newton_schulz", "gram_newton_schulz"] = field( + default="standard_newton_schulz", + metadata={ + "help": "Newton-Schulz backend for Muon. 'standard_newton_schulz' keeps the PyTorch Muon path; " + "'gram_newton_schulz' uses Dao-AILab's Gram Newton-Schulz formulation." + }, + ) + + muon_ns_use_quack_kernels: bool = field( + default=True, + metadata={ + "help": "Allow Muon Gram Newton-Schulz to use Quack symmetric GEMM kernels on supported Hopper/Blackwell GPUs. " + "Falls back to torch matmuls when unavailable." + }, + ) + + muon_gram_ns_num_restarts: int = field( + default=1, + metadata={ + "help": "Number of restart locations to autotune for Muon's Gram Newton-Schulz backend when explicit " + "muon_gram_ns_restart_iterations are not provided." + }, + ) + + muon_gram_ns_restart_iterations: Optional[List[int]] = field( + default=None, + metadata={ + "help": "Explicit restart iteration indices for Muon's Gram Newton-Schulz backend. " + "A value of 2 means restart after the second iteration." + }, + ) + muon_grad_dtype: Optional[Literal["fp32", "bf16"]] = field( + default=None, + metadata={ + "help": "Optional dtype cast for the gradient tensor used inside Muon. " + "Use this to force the Muon optimizer path to fp32 or bf16 independently of momentum state dtype." + }, + ) + muon_update_dtype: Optional[Literal["fp32", "bf16"]] = field( + default=None, + metadata={ + "help": "Optional dtype cast for the transient Muon update tensor passed into Newton-Schulz. " + "Use this to decouple compute dtype from gradient and momentum-buffer storage dtype." + }, + ) + muon_force_momentum_path: bool = field( + default=False, + metadata={ + "help": "Force Muon to build the update through the momentum-buffer path even when muon_momentum=0. " + "Intended for debugging and ablations." + }, + ) + # ======================================================================== # Checkpointing & Output # ======================================================================== @@ -519,6 +572,13 @@ def to_config_dict(self) -> Dict[str, Any]: "muon_nesterov": self.muon_nesterov, "muon_ns_steps": self.muon_ns_steps, "muon_adjust_lr_fn": self.muon_adjust_lr_fn, + "muon_ns_algorithm": self.muon_ns_algorithm, + "muon_ns_use_quack_kernels": self.muon_ns_use_quack_kernels, + "muon_gram_ns_num_restarts": self.muon_gram_ns_num_restarts, + "muon_gram_ns_restart_iterations": self.muon_gram_ns_restart_iterations, + "muon_grad_dtype": self.muon_grad_dtype, + "muon_update_dtype": self.muon_update_dtype, + "muon_force_momentum_path": self.muon_force_momentum_path, "load_checkpoint_path": self.load_checkpoint_path, "ckpt_manager": self.ckpt_manager, "enable_self_test": self.enable_self_test, diff --git a/tests/optim/test_muon.py b/tests/optim/test_muon.py new file mode 100644 index 00000000..0b3ff118 --- /dev/null +++ b/tests/optim/test_muon.py @@ -0,0 +1,456 @@ +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +import xorl.optim.gram_newton_schulz as gram_newton_schulz +import xorl.optim.muon as muon_module +from xorl.optim import Muon, build_optimizer +from xorl.optim.gram_newton_schulz import GramNewtonSchulzOrthogonalizer, expand_ns_coefficients, find_best_restarts + + +pytestmark = [pytest.mark.cpu] + + +class TinyModule(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(3, 2) + + +class FakeCudaTensor: + def __init__(self, *, dtype: torch.dtype, shape: tuple[int, int]): + self.dtype = dtype + self.device = torch.device("cuda", 0) + self.is_cuda = True + self._shape = shape + + def size(self, dim=None): + if dim is None: + return self._shape + return self._shape[dim] + + +def test_find_best_restarts_matches_default_muon_coefficients(): + coefficients = expand_ns_coefficients((3.4445, -4.775, 2.0315), 5) + + assert find_best_restarts(coefficients, num_restarts=1) == [3] + + +def test_muon_gram_newton_schulz_updates_parameters_and_autotunes_restart(): + param = nn.Parameter(torch.tensor([[1.0, -2.0, 3.0], [4.0, -5.0, 6.0]], dtype=torch.float32)) + optimizer = Muon( + [param], + lr=0.1, + momentum=0.0, + nesterov=False, + ns_algorithm="gram_newton_schulz", + ns_use_quack_kernels=False, + gram_newton_schulz_num_restarts=1, + ) + + before = param.detach().clone() + param.grad = torch.tensor([[0.5, -0.25, 0.75], [-1.0, 0.5, -0.5]], dtype=torch.float32) + + optimizer.step() + + assert torch.isfinite(param).all() + assert not torch.allclose(param, before) + assert len(optimizer._gram_ns_orthogonalizers) == 1 + orthogonalizer = next(iter(optimizer._gram_ns_orthogonalizers.values())) + assert list(orthogonalizer.reset_iterations) == [3] + + +def test_build_optimizer_threads_gram_newton_schulz_kwargs(): + model = TinyModule() + + optimizer = build_optimizer( + model, + lr=0.1, + weight_decay=0.01, + optimizer_type="muon", + no_decay_params=["bias"], + optimizer_kwargs={ + "muon_lr": 0.02, + "muon_ns_algorithm": "gram_newton_schulz", + "muon_ns_use_quack_kernels": False, + "muon_gram_ns_num_restarts": 1, + "muon_grad_dtype": "fp32", + "muon_update_dtype": "bf16", + "muon_force_momentum_path": True, + }, + ) + + assert isinstance(optimizer, Muon) + muon_groups = [group for group in optimizer.param_groups if group["use_muon"]] + adamw_groups = [group for group in optimizer.param_groups if not group["use_muon"]] + + assert muon_groups + assert adamw_groups + assert all(group["ns_algorithm"] == "gram_newton_schulz" for group in muon_groups) + assert all(group["ns_use_quack_kernels"] is False for group in muon_groups) + assert all(group["gram_newton_schulz_num_restarts"] == 1 for group in muon_groups) + assert optimizer._momentum_dtype is torch.bfloat16 + assert optimizer._grad_dtype is torch.float32 + assert optimizer._update_dtype is torch.bfloat16 + assert optimizer._force_momentum_path is True + + +def test_make_quack_backend_prefers_installed_quack(monkeypatch): + fake_gemm_interface = SimpleNamespace( + gemm=lambda A, B: ("gemm", A, B), + gemm_add=lambda A, B, C=None, beta=1.0: ("gemm_add", A, B, C, beta), + gemm_symmetric=lambda A, B, C=None, alpha=1.0, beta=1.0: ("gemm_symmetric", A, B, C, alpha, beta), + ) + import_calls = [] + + def fake_import_module(module_name): + import_calls.append(module_name) + if module_name == "quack.gemm_interface": + return fake_gemm_interface + raise AssertionError(f"Unexpected import {module_name}") + + monkeypatch.setattr(gram_newton_schulz.importlib, "import_module", fake_import_module) + gram_newton_schulz._make_quack_backend.cache_clear() + + backend = gram_newton_schulz._make_quack_backend() + + assert import_calls == ["quack.gemm_interface"] + assert backend.sym_baddbmm("A", "B", "C", alpha=2.0, beta=3.0) == ("gemm_symmetric", "A", "B", "C", 2.0, 3.0) + + gram_newton_schulz._make_quack_backend.cache_clear() + + +def test_make_quack_backend_requires_installed_quack(monkeypatch): + def fake_import_module(module_name): + assert module_name == "quack.gemm_interface" + raise ModuleNotFoundError("No module named 'quack'") + + monkeypatch.setattr(gram_newton_schulz.importlib, "import_module", fake_import_module) + gram_newton_schulz._make_quack_backend.cache_clear() + + with pytest.raises(ImportError, match="upstream `quack-kernels` package"): + gram_newton_schulz._make_quack_backend() + + gram_newton_schulz._make_quack_backend.cache_clear() + + +def test_select_backend_falls_back_to_torch_for_fp32_on_sm90(monkeypatch): + orthogonalizer = GramNewtonSchulzOrthogonalizer( + ns_coefficients=((3.4445, -4.775, 2.0315),) * 5, + ns_use_quack_kernels=True, + ) + quack_backend = object() + + monkeypatch.setattr(gram_newton_schulz, "_make_quack_backend", lambda: quack_backend) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (9, 0)) + + backend = orthogonalizer._select_backend(FakeCudaTensor(dtype=torch.float32, shape=(512, 512))) + + assert backend is gram_newton_schulz._TORCH_BACKEND + + +def test_select_backend_keeps_quack_for_bf16_on_sm90(monkeypatch): + orthogonalizer = GramNewtonSchulzOrthogonalizer( + ns_coefficients=((3.4445, -4.775, 2.0315),) * 5, + ns_use_quack_kernels=True, + ) + quack_backend = object() + + monkeypatch.setattr(gram_newton_schulz, "_make_quack_backend", lambda: quack_backend) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (9, 0)) + + backend = orthogonalizer._select_backend(FakeCudaTensor(dtype=torch.bfloat16, shape=(512, 512))) + + assert backend is quack_backend + + +def test_muon_groups_gram_newton_schulz_updates_by_shape(monkeypatch): + class FakeOrthogonalizer: + def __init__(self): + self.seen_shapes = [] + + def orthogonalize(self, X): + self.seen_shapes.append(tuple(X.shape)) + if X.ndim == 3: + offsets = torch.arange(1, X.shape[0] + 1, device=X.device, dtype=X.dtype).view(-1, 1, 1) + return X + offsets + return X + 7 + + p1 = nn.Parameter(torch.zeros((2, 3), dtype=torch.float32)) + p2 = nn.Parameter(torch.zeros((2, 3), dtype=torch.float32)) + p3 = nn.Parameter(torch.zeros((2, 2), dtype=torch.float32)) + optimizer = Muon( + [p1, p2, p3], + lr=1.0, + momentum=0.0, + nesterov=False, + ns_algorithm="gram_newton_schulz", + ns_use_quack_kernels=False, + ) + orthogonalizer = FakeOrthogonalizer() + + monkeypatch.setattr(muon_module, "_adjust_lr", lambda lr, adjust_lr_fn, shape: lr) + monkeypatch.setattr(optimizer, "_get_gram_ns_orthogonalizer", lambda group: orthogonalizer) + + p1.grad = torch.ones((2, 3), dtype=torch.float32) + p2.grad = torch.full((2, 3), 2.0, dtype=torch.float32) + p3.grad = torch.full((2, 2), 3.0, dtype=torch.float32) + + optimizer.step() + + assert orthogonalizer.seen_shapes == [(2, 2, 3), (2, 2)] + assert torch.allclose(p1, torch.full((2, 3), -2.0)) + assert torch.allclose(p2, torch.full((2, 3), -4.0)) + assert torch.allclose(p3, torch.full((2, 2), -10.0)) + + +def test_muon_groups_gram_newton_schulz_updates_by_matrix_shape(monkeypatch): + class FakeOrthogonalizer: + def __init__(self): + self.seen_shapes = [] + + def orthogonalize(self, X): + self.seen_shapes.append(tuple(X.shape)) + if X.ndim == 3: + offsets = torch.arange(1, X.shape[0] + 1, device=X.device, dtype=X.dtype).view(-1, 1, 1) + return X + offsets + return X + 7 + + p1 = nn.Parameter(torch.zeros((2, 3, 4), dtype=torch.float32)) + p2 = nn.Parameter(torch.zeros((1, 3, 4), dtype=torch.float32)) + p3 = nn.Parameter(torch.zeros((3, 5), dtype=torch.float32)) + optimizer = Muon( + [p1, p2, p3], + lr=1.0, + momentum=0.0, + nesterov=False, + ns_algorithm="gram_newton_schulz", + ns_use_quack_kernels=False, + ) + orthogonalizer = FakeOrthogonalizer() + + monkeypatch.setattr(muon_module, "_adjust_lr", lambda lr, adjust_lr_fn, shape: lr) + monkeypatch.setattr(optimizer, "_get_gram_ns_orthogonalizer", lambda group: orthogonalizer) + + p1.grad = torch.ones((2, 3, 4), dtype=torch.float32) + p2.grad = torch.full((1, 3, 4), 2.0, dtype=torch.float32) + p3.grad = torch.full((3, 5), 3.0, dtype=torch.float32) + + optimizer.step() + + assert orthogonalizer.seen_shapes == [(3, 3, 4), (3, 5)] + expected_p1 = torch.tensor( + [ + [[-2.0, -2.0, -2.0, -2.0], [-2.0, -2.0, -2.0, -2.0], [-2.0, -2.0, -2.0, -2.0]], + [[-3.0, -3.0, -3.0, -3.0], [-3.0, -3.0, -3.0, -3.0], [-3.0, -3.0, -3.0, -3.0]], + ], + dtype=torch.float32, + ) + assert torch.allclose(p1, expected_p1) + assert torch.allclose(p2, torch.full((1, 3, 4), -5.0)) + assert torch.allclose(p3, torch.full((3, 5), -10.0)) + + +def test_muon_groups_gram_newton_schulz_transpose_equivalent_shapes(monkeypatch): + class FakeOrthogonalizer: + def __init__(self): + self.seen_shapes = [] + + def orthogonalize(self, X): + self.seen_shapes.append(tuple(X.shape)) + offsets = torch.arange(1, X.shape[0] + 1, device=X.device, dtype=X.dtype).view(-1, 1, 1) + return X + offsets + + p1 = nn.Parameter(torch.zeros((2, 3), dtype=torch.float32)) + p2 = nn.Parameter(torch.zeros((3, 2), dtype=torch.float32)) + optimizer = Muon( + [p1, p2], + lr=1.0, + momentum=0.0, + nesterov=False, + ns_algorithm="gram_newton_schulz", + ns_use_quack_kernels=False, + ) + orthogonalizer = FakeOrthogonalizer() + + monkeypatch.setattr(muon_module, "_adjust_lr", lambda lr, adjust_lr_fn, shape: lr) + monkeypatch.setattr(optimizer, "_get_gram_ns_orthogonalizer", lambda group: orthogonalizer) + + p1.grad = torch.ones((2, 3), dtype=torch.float32) + p2.grad = torch.full((3, 2), 2.0, dtype=torch.float32) + + optimizer.step() + + assert orthogonalizer.seen_shapes == [(2, 2, 3)] + assert torch.allclose(p1, torch.full((2, 3), -2.0)) + assert torch.allclose(p2, torch.full((3, 2), -4.0)) + + +def test_muon_groups_fused_gate_up_halves(monkeypatch): + class FakeOrthogonalizer: + def __init__(self): + self.seen_shapes = [] + + def orthogonalize(self, X): + self.seen_shapes.append(tuple(X.shape)) + offsets = torch.arange(1, X.shape[0] + 1, device=X.device, dtype=X.dtype).view(-1, *([1] * (X.ndim - 1))) + return X + offsets + + p1 = nn.Parameter(torch.zeros((2, 3, 4), dtype=torch.float32)) + p2 = nn.Parameter(torch.zeros((2, 3, 4), dtype=torch.float32)) + optimizer = Muon( + [ + { + "params": [p1, p2], + "_fused_gate_up_ids": {id(p1), id(p2)}, + } + ], + lr=1.0, + momentum=0.0, + nesterov=False, + ns_algorithm="gram_newton_schulz", + ns_use_quack_kernels=False, + ) + orthogonalizer = FakeOrthogonalizer() + + monkeypatch.setattr(muon_module, "_adjust_lr", lambda lr, adjust_lr_fn, shape: lr) + monkeypatch.setattr(optimizer, "_get_gram_ns_orthogonalizer", lambda group: orthogonalizer) + + p1.grad = torch.ones((2, 3, 4), dtype=torch.float32) + p2.grad = torch.ones((2, 3, 4), dtype=torch.float32) + + optimizer.step() + + assert orthogonalizer.seen_shapes == [(8, 2, 3)] + expected_p1 = torch.tensor( + [ + [[-2.0, -2.0, -4.0, -4.0], [-2.0, -2.0, -4.0, -4.0], [-2.0, -2.0, -4.0, -4.0]], + [[-3.0, -3.0, -5.0, -5.0], [-3.0, -3.0, -5.0, -5.0], [-3.0, -3.0, -5.0, -5.0]], + ], + dtype=torch.float32, + ) + expected_p2 = torch.tensor( + [ + [[-6.0, -6.0, -8.0, -8.0], [-6.0, -6.0, -8.0, -8.0], [-6.0, -6.0, -8.0, -8.0]], + [[-7.0, -7.0, -9.0, -9.0], [-7.0, -7.0, -9.0, -9.0], [-7.0, -7.0, -9.0, -9.0]], + ], + dtype=torch.float32, + ) + assert torch.allclose(p1, expected_p1) + assert torch.allclose(p2, expected_p2) + + +def test_muon_standard_newton_schulz_preserves_batched_leading_dims(monkeypatch): + seen_shapes = [] + call_count = 0 + + def fake_zeropower(update, ns_coefficients, ns_steps, eps): + nonlocal call_count + call_count += 1 + seen_shapes.append(tuple(update.shape)) + return update + call_count + + p = nn.Parameter(torch.zeros((2, 3, 4), dtype=torch.float32)) + optimizer = Muon( + [p], + lr=1.0, + momentum=0.0, + nesterov=False, + ns_algorithm="standard_newton_schulz", + ) + + monkeypatch.setattr(muon_module, "_adjust_lr", lambda lr, adjust_lr_fn, shape: lr) + monkeypatch.setattr(muon_module, "_zeropower_via_newtonschulz", fake_zeropower) + + p.grad = torch.ones((2, 3, 4), dtype=torch.float32) + + optimizer.step() + + assert seen_shapes == [(3, 4), (3, 4)] + expected = torch.tensor( + [ + [[-2.0, -2.0, -2.0, -2.0], [-2.0, -2.0, -2.0, -2.0], [-2.0, -2.0, -2.0, -2.0]], + [[-3.0, -3.0, -3.0, -3.0], [-3.0, -3.0, -3.0, -3.0], [-3.0, -3.0, -3.0, -3.0]], + ], + dtype=torch.float32, + ) + assert torch.allclose(p, expected) + + +def test_muon_chunks_grouped_gram_newton_schulz_batches(monkeypatch): + class FakeOrthogonalizer: + def __init__(self): + self.seen_shapes = [] + + def orthogonalize(self, X): + self.seen_shapes.append(tuple(X.shape)) + return X + 1 + + p1 = nn.Parameter(torch.zeros((2, 3), dtype=torch.float32)) + p2 = nn.Parameter(torch.zeros((2, 3), dtype=torch.float32)) + optimizer = Muon( + [p1, p2], + lr=1.0, + momentum=0.0, + nesterov=False, + ns_algorithm="gram_newton_schulz", + ns_use_quack_kernels=False, + ) + orthogonalizer = FakeOrthogonalizer() + + monkeypatch.setattr(muon_module, "_adjust_lr", lambda lr, adjust_lr_fn, shape: lr) + monkeypatch.setattr(muon_module, "GROUPED_GRAM_NS_FP32_BYTE_LIMIT", 23) + monkeypatch.setattr(optimizer, "_get_gram_ns_orthogonalizer", lambda group: orthogonalizer) + + p1.grad = torch.ones((2, 3), dtype=torch.float32) + p2.grad = torch.ones((2, 3), dtype=torch.float32) + + optimizer.step() + + assert orthogonalizer.seen_shapes == [(2, 3), (2, 3)] + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for dtype-preservation coverage") +def test_gram_newton_schulz_preserves_fp32_compute_dtype_on_cuda(): + seen_dtypes = [] + + def sym_mm(A, B): + seen_dtypes.append(A.dtype) + return A @ B + + def sym_baddbmm(A, B, C, alpha=1.0, beta=1.0): + seen_dtypes.extend((A.dtype, B.dtype, C.dtype)) + if A.ndim == 2: + return torch.addmm(C, A, B, beta=beta, alpha=alpha) + return torch.baddbmm(C, A, B, beta=beta, alpha=alpha) + + def mm(A, B): + seen_dtypes.append(A.dtype) + return A @ B + + def mm_add(A, B, C, beta=1.0): + seen_dtypes.extend((A.dtype, B.dtype, C.dtype)) + if A.ndim == 2: + return torch.addmm(C, A, B, beta=beta) + return torch.baddbmm(C, A, B, beta=beta) + + orthogonalizer = GramNewtonSchulzOrthogonalizer( + ns_coefficients=((3.4445, -4.775, 2.0315),) * 5, + ns_use_quack_kernels=False, + ) + orthogonalizer._select_backend = lambda X: SimpleNamespace( + sym_mm=sym_mm, + sym_baddbmm=sym_baddbmm, + mm=mm, + mm_add=mm_add, + ) + + input_matrix = torch.randn(8, 8, device="cuda", dtype=torch.float32) + output_matrix = orthogonalizer.orthogonalize(input_matrix) + + assert output_matrix.dtype is torch.float32 + assert seen_dtypes + assert all(dtype is torch.float32 for dtype in seen_dtypes) diff --git a/tests/qlora/test_imports.py b/tests/qlora/test_imports.py new file mode 100644 index 00000000..15c60c72 --- /dev/null +++ b/tests/qlora/test_imports.py @@ -0,0 +1,13 @@ +import importlib + +import pytest + + +pytestmark = [pytest.mark.cpu] + + +def test_import_xorl_qlora_module(): + qlora = importlib.import_module("xorl.qlora") + + assert hasattr(qlora, "QLoRALinear") + assert hasattr(qlora, "inject_qlora_into_model") diff --git a/tests/server/test_server_arguments.py b/tests/server/test_server_arguments.py index 9ccedd25..1f63a119 100644 --- a/tests/server/test_server_arguments.py +++ b/tests/server/test_server_arguments.py @@ -28,3 +28,44 @@ def test_load_server_arguments_threads_signsgd_through_nested_config(tmp_path): assert args.optimizer == "signsgd" assert args.to_config_dict()["train"]["optimizer"] == "signsgd" + + +def test_load_server_arguments_threads_muon_gram_newton_schulz_through_nested_config(tmp_path): + config_path = tmp_path / "server_config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model": { + "model_path": "Qwen/Qwen3-8B", + }, + "train": { + "optimizer": "muon", + "muon_ns_algorithm": "gram_newton_schulz", + "muon_ns_use_quack_kernels": False, + "muon_gram_ns_num_restarts": 2, + "muon_gram_ns_restart_iterations": [2], + "muon_grad_dtype": "fp32", + "muon_update_dtype": "bf16", + "muon_force_momentum_path": True, + "output_dir": str(tmp_path / "outputs"), + }, + } + ), + encoding="utf-8", + ) + + args = load_server_arguments(str(config_path)) + train_config = args.to_config_dict()["train"] + + assert args.optimizer == "muon" + assert args.muon_ns_algorithm == "gram_newton_schulz" + assert train_config["muon_ns_algorithm"] == "gram_newton_schulz" + assert train_config["muon_ns_use_quack_kernels"] is False + assert train_config["muon_gram_ns_num_restarts"] == 2 + assert train_config["muon_gram_ns_restart_iterations"] == [2] + assert args.muon_grad_dtype == "fp32" + assert args.muon_update_dtype == "bf16" + assert args.muon_force_momentum_path is True + assert train_config["muon_grad_dtype"] == "fp32" + assert train_config["muon_update_dtype"] == "bf16" + assert train_config["muon_force_momentum_path"] is True diff --git a/tests/test_arguments.py b/tests/test_arguments.py index 11728439..07d32bc6 100644 --- a/tests/test_arguments.py +++ b/tests/test_arguments.py @@ -1,6 +1,7 @@ import sys import pytest +import torch import yaml from xorl.arguments import Arguments, parse_args @@ -41,3 +42,52 @@ def test_parse_args_accepts_signsgd_from_yaml(tmp_path, monkeypatch): assert args.train.optimizer == "signsgd" assert args.train.optimizer_kwargs == {} + + +def test_parse_args_wires_muon_kwargs_from_yaml(tmp_path, monkeypatch): + config_path = tmp_path / "config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model": { + "model_path": "Qwen/Qwen3-8B", + }, + "data": { + "datasets": [{"path": "dummy", "type": "tokenized"}], + }, + "train": { + "init_device": "meta", + "output_dir": str(tmp_path / "outputs"), + "optimizer": "muon", + "optimizer_dtype": "bf16", + "muon_ns_algorithm": "gram_newton_schulz", + "muon_ns_use_quack_kernels": False, + "muon_gram_ns_num_restarts": 2, + "muon_gram_ns_restart_iterations": [2], + "muon_grad_dtype": "fp32", + "muon_update_dtype": "fp32", + "muon_force_momentum_path": True, + "use_wandb": False, + }, + } + ), + encoding="utf-8", + ) + + monkeypatch.setenv("WORLD_SIZE", "1") + monkeypatch.setenv("LOCAL_WORLD_SIZE", "1") + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("LOCAL_RANK", "0") + monkeypatch.setattr(sys, "argv", ["train.py", str(config_path)]) + + args = parse_args(Arguments) + + assert args.train.optimizer == "muon" + assert args.train.optimizer_kwargs["muon_ns_algorithm"] == "gram_newton_schulz" + assert args.train.optimizer_kwargs["muon_ns_use_quack_kernels"] is False + assert args.train.optimizer_kwargs["muon_gram_ns_num_restarts"] == 2 + assert args.train.optimizer_kwargs["muon_gram_ns_restart_iterations"] == [2] + assert args.train.optimizer_kwargs["muon_momentum_dtype"] is torch.bfloat16 + assert args.train.optimizer_kwargs["muon_grad_dtype"] is torch.float32 + assert args.train.optimizer_kwargs["muon_update_dtype"] is torch.float32 + assert args.train.optimizer_kwargs["muon_force_momentum_path"] is True From 99d3fa43e39c77c08058c141041b977aedad0e2c Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 10 Apr 2026 20:38:21 -0700 Subject: [PATCH 28/41] Fix Qwen3.5-35B MoE routing and fused_recurrent kwargs bug (#121) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Fix Qwen3.5-35B MoE routing and fused_recurrent kwargs bug Two fixes: 1. **norm_topk_prob default** (root cause of K3 ~40): `norm_topk_prob=None` in the HF checkpoint was interpreted as `False` by xorl, skipping routing weight normalization. HF always normalizes top-k weights. Fix: default to `True` when config value is `None`. Result: K3 drops from 38.76 → 0.003 2. **fused_recurrent kwargs bug**: `fused_recurrent_gated_delta_rule()` passed keyword arguments to `autograd.Function.apply()` which doesn't accept them, crashing any sequence ≤ 64 tokens. Closes #117 (cherry picked from commit e60e8817a0b9ca6bcd9762bdf8abf3c6505d4b5f) --- .../qwen3_5_moe/configuration_qwen3_5_moe.py | 2 +- .../ops/gated_delta_rule/fused_recurrent.py | 30 +++++++++++++++++-- 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/src/xorl/models/transformers/qwen3_5_moe/configuration_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/configuration_qwen3_5_moe.py index a028859d..2112f9d4 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/configuration_qwen3_5_moe.py +++ b/src/xorl/models/transformers/qwen3_5_moe/configuration_qwen3_5_moe.py @@ -251,7 +251,7 @@ def from_hf_config(cls, hf_config): moe_intermediate_size=getattr(text_config, "moe_intermediate_size", 768), num_experts_per_tok=getattr(text_config, "num_experts_per_tok", 8), num_experts=getattr(text_config, "num_experts", 128), - norm_topk_prob=getattr(text_config, "norm_topk_prob", False), + norm_topk_prob=getattr(text_config, "norm_topk_prob", None) is not False, output_router_logits=getattr(text_config, "output_router_logits", False), router_aux_loss_coef=getattr(text_config, "router_aux_loss_coef", 0.001), mlp_only_layers=getattr(text_config, "mlp_only_layers", []), diff --git a/src/xorl/ops/linear_attention/ops/gated_delta_rule/fused_recurrent.py b/src/xorl/ops/linear_attention/ops/gated_delta_rule/fused_recurrent.py index d084c54f..043209a2 100644 --- a/src/xorl/ops/linear_attention/ops/gated_delta_rule/fused_recurrent.py +++ b/src/xorl/ops/linear_attention/ops/gated_delta_rule/fused_recurrent.py @@ -219,5 +219,31 @@ def forward( ) -def fused_recurrent_gated_delta_rule(*args, **kwargs): - return FusedRecurrentFunction.apply(*args, **kwargs) +def fused_recurrent_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, +): + return FusedRecurrentFunction.apply( + q, + k, + v, + g, + gk, + gv, + beta, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + cu_seqlens, + ) From 6171928b64b4111acbc3477da18844fed756c442 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Tue, 14 Apr 2026 13:27:41 -0700 Subject: [PATCH 29/41] Unified gradient_checkpointing_method: +12-22% MoE throughput with memory trade-off (#85) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Unified gradient_checkpointing_method: +12-22% MoE throughput Replace fragmented gradient checkpointing config (recompute_modules + moe_checkpoint_method) with a single gradient_checkpointing_method field: - recompute_full_layer (default): recompute entire decoder layer - recompute_before_dispatch: checkpoint attn+router, keep dispatch+expert+combine (+12%) - no_recompute: no recomputation, max throughput (+22%) Key changes: - MoEGradientCheckpointingLayer base class: provides _pre_dispatch_forward and _moe_forward so new MoE models only implement _pre_mlp_forward - MoEBlock.route() + forward_experts_only(): extract routing logic into reusable methods, eliminating 40-line duplication in decoder layers - Fix routing replay: cache routing_weights alongside selected_experts for deterministic checkpoint recompute (fixes NCCL timeout with recompute_full_layer + EP) - Remove moe_act kernel variants (QuackEPGroupGemmMoeAct, TritonEPGroupGemmMoeAct, native moe_act): benchmarked no measurable impact with recompute_before_dispatch - Remove sonic dead code, quack kernel warmup - Works with both alltoall and DeepEP backends - Docs updated across 6 pages Benchmarks: Qwen3-Coder-30B-A3B, quack + alltoall, 1-node EP8, H100 80GB, 32k seq recompute_full_layer: 17,600 tok/s, 37.5 GB peak recompute_before_dispatch: 19,990 tok/s, 44.5 GB peak (+13.6%) no_recompute: 21,800 tok/s, 54.8 GB peak (+23.9%) Convergence verified: all methods identical step-1 loss, monotonic convergence over 50 steps (spread ~0.36 at step 50, normal fp non-determinism). * Fix EnvironMeter kwarg: gc_enabled -> gradient_checkpointing_enabled The EnvironMeter.__init__ renamed `gc_enabled` to `gradient_checkpointing_enabled` and added a new `gradient_checkpointing_method` parameter. Update the trainer call site to use the new kwarg names so XorlFlopsCounter receives the correct gradient checkpointing configuration. * Flatten hidden_states before route() in _pre_dispatch_forward In the selective-checkpoint path (recompute_before_dispatch), _pre_dispatch_forward calls self.mlp.route() with the 3-D (batch, seq, hidden) tensor from _pre_mlp_forward. However, route() expects a flattened (num_tokens, hidden_dim) tensor because it runs softmax routing. Flatten before calling route to match the expected shape; forward_experts_only already handles the 3-D input separately. * Accept expert_scores in native/quack EP compute adapters The EP forward path now passes 6 arguments to compute_fn including expert_scores. The native and quack adapter wrappers still only accepted 5 positional args, causing a TypeError at call time. Add expert_scores as an optional kwarg to both wrappers to match the fused call signature from experts.py. * Fix MoE decoder layer attention-weight indexing in model loops _moe_forward returns (hidden_states,) or (hidden_states, router_logits) and never includes attention weights. When output_attentions=True the model loop was reading layer_outputs[1] as self-attention weights, but that slot actually held router_logits (when output_router_logits was also True). Use a None placeholder for attention weights instead, since the MoE selective-checkpoint path does not expose per-layer attention matrices. Fixes both Qwen3Moe and Qwen3.5Moe model classes. * Fix: pass expert_scores through native/quack EP wrappers Previous fix accepted the parameter but silently dropped it. Both native_ep_compute and QuackEPGroupGemm.apply already support expert_scores — the fused wrappers just weren't forwarding it. --------- Co-authored-by: qywu (cherry picked from commit b859ea2877e52c63f33f97b76d5405fed16bbc69) --- .../content/docs/config-reference/local.md | 5 +- docs/src/content/docs/moe/deepep.mdx | 2 +- .../content/docs/moe/expert-parallelism.mdx | 4 +- docs/src/content/docs/moe/kernels.mdx | 48 ++- docs/src/content/docs/moe/overview.mdx | 4 +- .../content/docs/training/local_training.mdx | 28 +- src/xorl/arguments.py | 45 ++- src/xorl/cli/direct_train.py | 8 +- src/xorl/cli/train.py | 1 + src/xorl/distributed/torch_parallelize.py | 30 +- src/xorl/models/base.py | 23 +- .../models/layers/moe/backend/__init__.py | 79 +--- src/xorl/models/layers/moe/backend/native.py | 136 ------- .../layers/moe/backend/quack_moe_act.py | 28 -- .../layers/moe/backend/triton_moe_act.py | 28 -- src/xorl/models/layers/moe/experts.py | 19 +- src/xorl/models/layers/moe/moe_block.py | 96 +++-- src/xorl/models/module_utils.py | 94 ++++- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 84 ++-- .../qwen3_moe/modeling_qwen3_moe.py | 126 +++--- src/xorl/ops/moe/quack.py | 306 -------------- src/xorl/ops/moe/triton.py | 381 ------------------ src/xorl/trainers/trainer.py | 20 +- src/xorl/utils/count_flops.py | 20 +- src/xorl/utils/helper.py | 10 +- 25 files changed, 419 insertions(+), 1206 deletions(-) delete mode 100644 src/xorl/models/layers/moe/backend/quack_moe_act.py delete mode 100644 src/xorl/models/layers/moe/backend/triton_moe_act.py diff --git a/docs/src/content/docs/config-reference/local.md b/docs/src/content/docs/config-reference/local.md index 7314047b..7808a12d 100644 --- a/docs/src/content/docs/config-reference/local.md +++ b/docs/src/content/docs/config-reference/local.md @@ -158,10 +158,9 @@ Each entry in `datasets` (or `test_datasets`) is a dict: | Field | Default | Description | |---|---|---| | `enable_mixed_precision` | `true` | BF16 mixed-precision training. | -| `enable_gradient_checkpointing` | `true` | Activation recomputation to reduce memory. | +| `enable_gradient_checkpointing` | `true` | Enable activation recomputation to reduce memory. | +| `gradient_checkpointing_method` | `recompute_full_layer` | What to recompute in backward. Valid values: `recompute_full_layer` (recompute entire decoder layer, most memory-efficient), `recompute_before_dispatch` (recompute attn+router, keep dispatch+expert+combine, +25-34% throughput), `no_recompute` (no recomputation, max throughput, highest memory). See [gradient checkpointing guide](/training/local_training#gradient-checkpointing). | | `enable_reentrant` | `false` | Use reentrant gradient checkpointing. Default (non-reentrant) is generally preferred. | -| `recompute_modules` | `null` | Selective checkpointing by submodule: `[self_attn]`, `[mlp]`, or `[self_attn, mlp]`. `null` = whole-layer recompute. | -| `moe_checkpoint_method` | `null` | MoE-specific checkpoint: `null` (full recompute including EP communication), `moe_act` (recompute only gate/up activations, skip EP communication recompute — faster). | | `enable_full_shard` | `true` | FSDP2 full parameter sharding (ZeRO-3). Set `false` for ZeRO-2. | | `enable_forward_prefetch` | `true` | Prefetch next FSDP unit's parameters during forward pass. | | `enable_activation_offload` | `false` | Offload activations to CPU during forward pass. | diff --git a/docs/src/content/docs/moe/deepep.mdx b/docs/src/content/docs/moe/deepep.mdx index aca61095..254013a6 100644 --- a/docs/src/content/docs/moe/deepep.mdx +++ b/docs/src/content/docs/moe/deepep.mdx @@ -294,7 +294,7 @@ XORL_DEBUG_EP=1 torchrun ... -m xorl.cli.train config.yaml ``` **Shape mismatch errors during backward:** -Ensure `moe_checkpoint_method: moe_act` is set — R3 routing replay is required for EP + gradient checkpointing. See [MoE Routing Replay](/moe/overview/#routing-replay-r3). +With the default `gradient_checkpointing_method: recompute_full_layer`, R3 routing replay is used automatically to ensure correct expert assignments when alltoall is recomputed. If you see shape mismatches, verify routing replay is enabled. Using `recompute_before_dispatch` or `no_recompute` avoids alltoall recomputation entirely and does not require replay. See [MoE Routing Replay](/moe/overview/#routing-replay-r3). **No fallback to AllToAll:** If DeepEP fails to initialize, xorl does not fall back automatically. Set `ep_dispatch: alltoall` explicitly. diff --git a/docs/src/content/docs/moe/expert-parallelism.mdx b/docs/src/content/docs/moe/expert-parallelism.mdx index 4635dd4a..0dc390c1 100644 --- a/docs/src/content/docs/moe/expert-parallelism.mdx +++ b/docs/src/content/docs/moe/expert-parallelism.mdx @@ -123,7 +123,7 @@ train: expert_parallel_size: 4 ringattn_parallel_size: 4 enable_gradient_checkpointing: true - moe_checkpoint_method: moe_act + gradient_checkpointing_method: recompute_before_dispatch ``` **Qwen3-235B-A22B (EP=64, 8 nodes, DeepEP):** @@ -139,7 +139,7 @@ train: expert_parallel_size: 64 ulysses_parallel_size: 64 enable_gradient_checkpointing: true - moe_checkpoint_method: moe_act + gradient_checkpointing_method: recompute_before_dispatch init_device: meta load_weights_mode: all_ranks ``` diff --git a/docs/src/content/docs/moe/kernels.mdx b/docs/src/content/docs/moe/kernels.mdx index 8fade165..ad433d22 100644 --- a/docs/src/content/docs/moe/kernels.mdx +++ b/docs/src/content/docs/moe/kernels.mdx @@ -185,25 +185,53 @@ The `native` backend (`torch._grouped_mm`) compiles cleanly end-to-end. The `qua --- -## MoE Gradient Checkpointing (`moe_checkpoint_method`) +## MoE Gradient Checkpointing (`gradient_checkpointing_method`) -xorl provides two MoE-aware checkpoint methods: +`gradient_checkpointing_method` controls what is recomputed in the backward pass. For MoE models with EP, this has a large impact on throughput because `recompute_full_layer` recomputes AllToAll dispatch + combine on every backward layer. -| Method | What is recomputed | EP comm recomputed | Memory savings | -|---|---|---|---| -| `moe_act` (**default**) | Expert MLP activations only (gate×up product inside each expert) | **No** — dispatch/combine skipped via R3 | Moderate — recovers large expert FFN tensors | -| `full_recompute` | All activations: attention, router, expert FFN | Yes — full AllToAll re-executed | Maximum — stores nothing | -| `None` / not set | Nothing | N/A | None — all activations kept | +### Methods + +| Method | Backward recomputes | EP comm recomputed | Peak mem (32k) | Throughput | +|---|---|---|---|---| +| `recompute_full_layer` (default) | Entire decoder layer | **Yes** — full AllToAll re-executed | 37.5 GB (1N) / 21.2 GB (2N) | baseline | +| `recompute_before_dispatch` | Attn + router only; keeps dispatch + expert + combine | **No** | 54.8 GB (1N) / 46.6 GB (2N) | **+24.5% (1N) / +33.4% (2N)** | +| `no_recompute` | Nothing | **No** | highest | max throughput | + +*Benchmarks: Qwen3-Coder-30B-A3B, quack + alltoall, 32k seq, H100 80GB. 1N = 1-node EP8, 2N = 2-node EP16.* ```yaml +# Maximum throughput (short seq, memory permits): +train: + enable_gradient_checkpointing: true + gradient_checkpointing_method: no_recompute + +# Balanced memory/speed — recompute attn+router, keep dispatch+expert+combine: train: enable_gradient_checkpointing: true - moe_checkpoint_method: moe_act # strongly recommended for EP models + gradient_checkpointing_method: recompute_before_dispatch + +# Maximum memory savings (long seq, required for 128k): +train: + enable_gradient_checkpointing: true + gradient_checkpointing_method: recompute_full_layer ``` -**Why `moe_act` is the right default:** Expert FFN activations (gate×up intermediate tensor of shape `[tokens, intermediate_size]`) are the largest memory consumer in a MoE forward pass. `moe_act` drops these and recomputes them from stored inputs during backward, without re-running the expensive AllToAll dispatch. This gives most of the memory savings of full recompute at a fraction of the communication cost. +### How `recompute_before_dispatch` works (Megatron-style) + +In the expert MLP forward: `gate_output = x @ gate_proj`, `up_output = x @ up_proj`, `postact = silu(gate_output) * up_output`, `output = postact @ down_proj`. + +With `recompute_before_dispatch`, the attention and router outputs are recomputed in backward, but **MoE dispatch and combine are not recomputed**. This avoids re-executing AllToAll communication, which is the dominant cost at multi-node. + +The elementwise recompute (`postact = silu(gate_output) * up_output`) has ~0% overhead — confirmed by benchmarks showing `recompute_before_dispatch` within 0.2% of no-checkpointing speed. + +### Why the speedup grows with node count + +| Scale | recompute_full_layer | recompute_before_dispatch | Speedup | +|---|---|---|---| +| 1-node EP8 (NVLink) | 17,542 tok/s | 21,851 tok/s | +24.5% | +| 2-node EP16 (InfiniBand) | 20,586 tok/s | 27,470 tok/s | +33.4% | -**How it works in the Triton kernel:** The `group_gemm_same_nk_kernel` has an optional `STORE_ACTIVATIONS` mode. When `moe_act` is active, the kernel stores the post-gate×up intermediate activations to a separate buffer. The custom backward function recomputes only the down projection from these stored intermediates. +At 2-node, dispatch/combine crosses InfiniBand (~100 Gbps) instead of NVLink (~900 Gbps). `recompute_full_layer` recomputes both dispatch and combine in backward — doubling the IB communication. `recompute_before_dispatch` avoids this entirely. --- diff --git a/docs/src/content/docs/moe/overview.mdx b/docs/src/content/docs/moe/overview.mdx index 139256de..d6237e0e 100644 --- a/docs/src/content/docs/moe/overview.mdx +++ b/docs/src/content/docs/moe/overview.mdx @@ -137,7 +137,7 @@ train: expert_parallel_size: 4 ringattn_parallel_size: 4 enable_gradient_checkpointing: true - moe_checkpoint_method: moe_act + gradient_checkpointing_method: recompute_before_dispatch optimizer: muon muon_lr: 2e-4 lr: 1e-4 @@ -162,7 +162,7 @@ train: ulysses_parallel_size: 64 data_parallel_shard_size: 1 enable_gradient_checkpointing: true - moe_checkpoint_method: moe_act + gradient_checkpointing_method: recompute_before_dispatch enable_activation_offload: true init_device: meta load_weights_mode: all_ranks diff --git a/docs/src/content/docs/training/local_training.mdx b/docs/src/content/docs/training/local_training.mdx index df73804f..ba602512 100644 --- a/docs/src/content/docs/training/local_training.mdx +++ b/docs/src/content/docs/training/local_training.mdx @@ -244,16 +244,38 @@ Checkpoints are saved to `{output_dir}/checkpoints/global_step_{N}/`. ## Gradient Checkpointing +Controls activation recomputation during backward to trade compute for memory. + ```yaml train: enable_gradient_checkpointing: true - recompute_modules: [self_attn, mlp] # selective recompute + gradient_checkpointing_method: recompute_full_layer # default — recompute each decoder layer ``` -For MoE models, `moe_checkpoint_method: moe_act` recomputes only activations inside expert FFNs, skipping the EP dispatch (faster than full recompute): +### Available methods + +| Method | Backward recomputes | Peak memory | Speed | +|--------|---------------------|------------|-------| +| `recompute_full_layer` | entire decoder layer (attn + MoE dispatch + expert + combine) | **lowest** | baseline | +| `recompute_before_dispatch` | recompute attn+router, keep dispatch+expert+combine | moderate | **+25-34%** | +| `no_recompute` | nothing recomputed | highest | max throughput | + +### When to use each + +- **`recompute_full_layer`** (default): required for long sequences (128k+) where activations exceed GPU memory. +- **`recompute_before_dispatch`**: best tradeoff when memory allows. Recomputes attention and router but avoids recomputing MoE dispatch/combine, which is especially expensive at multi-node (InfiniBand). +24.5% at 1-node, +33.4% at 2-node. +- **`no_recompute`**: maximum throughput, highest memory usage. Use for short sequences where activations fit in GPU memory. + ```yaml +# Maximum throughput (short seq, memory permits): train: - moe_checkpoint_method: moe_act + enable_gradient_checkpointing: true + gradient_checkpointing_method: no_recompute + +# Balanced memory/speed — recompute attn+router, keep dispatch+expert+combine: +train: + enable_gradient_checkpointing: true + gradient_checkpointing_method: recompute_before_dispatch ``` ## Logging diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 7505f45a..899b21a6 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -702,20 +702,21 @@ def optimizer_kwargs(self) -> Dict[str, Any]: default=False, metadata={"help": "Use reentrant gradient checkpointing."}, ) - recompute_modules: Optional[List[str]] = field( + gradient_checkpointing_method: Optional[str] = field( default=None, metadata={ - "help": "Per-submodule selective checkpointing. Options: 'self_attn', 'mlp'. " - "When None, uses whole-layer checkpoint (legacy). " - "Example: ['self_attn'] checkpoints only attention, keeping MoE activations." - }, - ) - moe_checkpoint_method: Optional[str] = field( - default=None, - metadata={ - "help": "MoE checkpoint strategy. None = full recompute (default). " - "'moe_act' = recompute only gate/up activations in backward, " - "skip EP communication recomputation. Only effective with MoE models." + "help": ( + "Gradient checkpointing strategy. Controls the speed/memory tradeoff\n" + "by choosing what to recompute vs keep in backward.\n\n" + " 'recompute_full_layer' — recompute entire decoder layer including\n" + " dispatch + combine alltoall. Lowest memory.\n" + " Required for 128k+ seq. (default)\n" + " 'recompute_before_dispatch' — recompute attn + layernorm + router; keep\n" + " dispatch + expert + combine. +20% speed,\n" + " more memory than recompute_full_layer.\n" + " 'no_recompute' — no recomputation, max throughput. +34% speed\n" + " but highest memory, only fits short seq.\n" + ) }, ) @@ -726,9 +727,7 @@ def moe_recomputed(self) -> bool: Used to decide whether routing replay is needed with EP: replay is only required when the MoE forward (including EP all-to-all) is recomputed. """ - if self.recompute_modules is not None: - return "mlp" in self.recompute_modules and self.moe_checkpoint_method != "moe_act" - return True # legacy whole-layer checkpoint always recomputes MoE + return self.gradient_checkpointing_method in (None, "recompute_full_layer") enable_full_shard: bool = field( default=True, @@ -950,6 +949,16 @@ def moe_recomputed(self) -> bool: ) def __post_init__(self): + # Resolve gradient_checkpointing_method into internal fields used by + # the model and parallelization code. + gcm = self.gradient_checkpointing_method or "recompute_full_layer" + self.gradient_checkpointing_method = gcm + if gcm not in ("recompute_full_layer", "recompute_before_dispatch", "no_recompute"): + raise ValueError( + f"Unknown gradient_checkpointing_method: {gcm!r}. " + f"Choose from: recompute_full_layer, recompute_before_dispatch, no_recompute" + ) + if self.repo_commit is None: self.repo_commit = _detect_repo_commit() self.local_rank = int(os.getenv("LOCAL_RANK", "0")) @@ -1374,9 +1383,11 @@ def parse_args(rootclass: T) -> T: elif attr.default is MISSING: parser_kwargs["required"] = True else: - # Regular list handling + # Regular list handling. Optional[List[X]] uses nargs="*" so that + # an empty list can be passed (e.g. --train.gradient_checkpointing_method + # with no args). Non-optional lists keep nargs="+" (at least one item). parser_kwargs["type"] = attr_type.__args__[0] - parser_kwargs["nargs"] = "+" + parser_kwargs["nargs"] = "*" if is_optional else "+" if attr.default_factory is not MISSING: parser_kwargs["default"] = attr.default_factory() elif attr.default is MISSING: diff --git a/src/xorl/cli/direct_train.py b/src/xorl/cli/direct_train.py index 2ea3d810..e6bf5606 100644 --- a/src/xorl/cli/direct_train.py +++ b/src/xorl/cli/direct_train.py @@ -327,8 +327,7 @@ def main(): enable_compile=args.train.enable_compile, basic_modules=model._no_split_modules + args.model.basic_modules, enable_reentrant=args.train.enable_reentrant, - recompute_modules=args.train.recompute_modules, - moe_checkpoint_method=args.train.moe_checkpoint_method, + gradient_checkpointing_method=args.train.gradient_checkpointing_method, enable_forward_prefetch=args.train.enable_forward_prefetch, load_weights_mode=args.train.load_weights_mode, pp_schedule=args.train.pipeline_parallel_schedule if args.train.pipeline_parallel_size > 1 else None, @@ -434,9 +433,8 @@ def main(): config=model_config, global_batch_size=args.train.global_batch_size, empty_cache_steps=args.train.empty_cache_steps, - gc_enabled=args.train.enable_gradient_checkpointing, - recompute_modules=args.train.recompute_modules, - moe_checkpoint_method=args.train.moe_checkpoint_method, + gradient_checkpointing_enabled=args.train.enable_gradient_checkpointing, + gradient_checkpointing_method=args.train.gradient_checkpointing_method, cp_size=args.train.ulysses_parallel_size * args.train.ringattn_parallel_size, ) diff --git a/src/xorl/cli/train.py b/src/xorl/cli/train.py index 88b7cc78..1e0327f6 100644 --- a/src/xorl/cli/train.py +++ b/src/xorl/cli/train.py @@ -22,6 +22,7 @@ def _single_config_autotune(configs, *args, **kwargs): triton.autotune = _single_config_autotune + from xorl.arguments import Arguments, parse_args from xorl.trainers import Trainer diff --git a/src/xorl/distributed/torch_parallelize.py b/src/xorl/distributed/torch_parallelize.py index 4a08790a..e08ba949 100644 --- a/src/xorl/distributed/torch_parallelize.py +++ b/src/xorl/distributed/torch_parallelize.py @@ -413,8 +413,7 @@ def build_parallelize_model( # Extract gradient checkpointing kwargs before the loop use_reentrant = kwargs.pop("enable_reentrant", False) recompute_context_fn = kwargs.pop("recompute_context_fn", noop_context_fn) - recompute_modules = kwargs.pop("recompute_modules", None) - moe_checkpoint_method = kwargs.pop("moe_checkpoint_method", None) + grad_ckpt_method = kwargs.pop("gradient_checkpointing_method", None) # 4. Apply parallelism to each model part for i, model_part in enumerate(model_parts): @@ -425,17 +424,15 @@ def build_parallelize_model( if use_reentrant: torch.utils.checkpoint.CheckpointFunction = CheckpointFunction - gc_kwargs = {"use_reentrant": use_reentrant} + grad_ckpt_kwargs = {"use_reentrant": use_reentrant} # Skip context_fn when torch.compile is enabled -- dynamo can't trace # through the SkipFunctionVariable (noop_context_fn). Without context_fn, # checkpoint uses the same default behavior and dynamo can trace natively. if not enable_compile: - gc_kwargs["context_fn"] = recompute_context_fn if i == 0 else noop_context_fn - if recompute_modules is not None: - gc_kwargs["recompute_modules"] = recompute_modules - if moe_checkpoint_method is not None: - gc_kwargs["moe_checkpoint_method"] = moe_checkpoint_method - model_part.gradient_checkpointing_enable(gradient_checkpointing_kwargs=gc_kwargs) + grad_ckpt_kwargs["context_fn"] = recompute_context_fn if i == 0 else noop_context_fn + if grad_ckpt_method is not None: + grad_ckpt_kwargs["gradient_checkpointing_method"] = grad_ckpt_method + model_part.gradient_checkpointing_enable(gradient_checkpointing_kwargs=grad_ckpt_kwargs) # TP (if enabled) if ps.tp_enabled: @@ -551,23 +548,20 @@ def build_parallelize_model( if use_reentrant: torch.utils.checkpoint.CheckpointFunction = CheckpointFunction - gc_kwargs = {"use_reentrant": use_reentrant} + grad_ckpt_kwargs = {"use_reentrant": use_reentrant} # Only pass context_fn when torch.compile is NOT enabled AND a custom # recompute_context_fn is explicitly provided. Passing context_fn # (even noop_context_fn) triggers SAC which uses torch.compile internally, # causing unexpected Inductor compilation when compile is disabled. recompute_fn = kwargs.pop("recompute_context_fn", None) if recompute_fn is not None and not enable_compile: - gc_kwargs["context_fn"] = recompute_fn + grad_ckpt_kwargs["context_fn"] = recompute_fn - recompute_modules = kwargs.pop("recompute_modules", None) - moe_checkpoint_method = kwargs.pop("moe_checkpoint_method", None) - if recompute_modules is not None: - gc_kwargs["recompute_modules"] = recompute_modules - if moe_checkpoint_method is not None: - gc_kwargs["moe_checkpoint_method"] = moe_checkpoint_method + grad_ckpt_method = kwargs.pop("gradient_checkpointing_method", None) + if grad_ckpt_method is not None: + grad_ckpt_kwargs["gradient_checkpointing_method"] = grad_ckpt_method - model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=gc_kwargs) + model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=grad_ckpt_kwargs) if use_reentrant: # Reentrant checkpointing doesn't support kwargs. Wrap the checkpoint diff --git a/src/xorl/models/base.py b/src/xorl/models/base.py index bb266ea3..2d5b76a5 100644 --- a/src/xorl/models/base.py +++ b/src/xorl/models/base.py @@ -6,7 +6,6 @@ from torch import nn from ..utils import logging -from .layers.moe.experts import MoEExperts from .layers.moe.moe_block import MoEBlock from .layers.moe.routing_replay import RoutingReplay @@ -93,26 +92,18 @@ def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): gradient_checkpointing_kwargs = {"use_reentrant": False} # Pop selective checkpoint config before passing to torch checkpoint - recompute_modules = gradient_checkpointing_kwargs.pop("recompute_modules", None) - moe_checkpoint_method = gradient_checkpointing_kwargs.pop("moe_checkpoint_method", None) + grad_ckpt_method = gradient_checkpointing_kwargs.pop("gradient_checkpointing_method", None) - gc_func = partial(torch.utils.checkpoint.checkpoint, **gradient_checkpointing_kwargs) + grad_ckpt_func = partial(torch.utils.checkpoint.checkpoint, **gradient_checkpointing_kwargs) - moe_act = moe_checkpoint_method == "moe_act" for module in self.modules(): if hasattr(module, "gradient_checkpointing"): module.gradient_checkpointing = True - module._gradient_checkpointing_func = gc_func - module._recompute_modules = recompute_modules - module._moe_checkpoint_method = moe_checkpoint_method - if isinstance(module, MoEExperts): - module._moe_act = moe_act - - if recompute_modules is not None: - logger.info( - f"Selective checkpointing enabled: recompute_modules={recompute_modules}, " - f"moe_checkpoint_method={moe_checkpoint_method}" - ) + module._gradient_checkpointing_func = grad_ckpt_func + module._gradient_checkpointing_method = grad_ckpt_method + + if grad_ckpt_method is not None: + logger.info(f"Selective checkpointing enabled: gradient_checkpointing_method={grad_ckpt_method}") # Create routing replay instances for MoE blocks self.enable_routing_replay() diff --git a/src/xorl/models/layers/moe/backend/__init__.py b/src/xorl/models/layers/moe/backend/__init__.py index 271ef343..5b8f0796 100644 --- a/src/xorl/models/layers/moe/backend/__init__.py +++ b/src/xorl/models/layers/moe/backend/__init__.py @@ -60,10 +60,10 @@ try: from xorl.ops.moe.quack import QuackEPGroupGemm as _QuackEPGroupGemm - def _quack_ep_fused(permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size): + def _quack_ep_fused(permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores=None): gate_proj = gate_up_proj[..., :intermediate_size].contiguous() up_proj = gate_up_proj[..., intermediate_size:].contiguous() - return _QuackEPGroupGemm.apply(permute_tokens, cumsum, gate_proj, up_proj, down_proj) + return _QuackEPGroupGemm.apply(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores) EP_EXPERT_COMPUTE["quack"] = _quack_ep_fused except ImportError: @@ -73,10 +73,10 @@ def _quack_ep_fused(permute_tokens, cumsum, gate_up_proj, down_proj, intermediat try: from .native import native_ep_compute as _native_ep_compute - def _native_ep_fused(permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size): + def _native_ep_fused(permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores=None): gate_proj = gate_up_proj[..., :intermediate_size].contiguous() up_proj = gate_up_proj[..., intermediate_size:].contiguous() - return _native_ep_compute(permute_tokens, cumsum, gate_proj, up_proj, down_proj) + return _native_ep_compute(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores) EP_EXPERT_COMPUTE["native"] = _native_ep_fused except ImportError: @@ -186,75 +186,6 @@ def _native_ep_fused(permute_tokens, cumsum, gate_up_proj, down_proj, intermedia pass -# --------------------------------------------------------------------------- -# Moe_act registries: activation-recompute variants that drop gate_output/up_output -# from save_for_backward and recompute them in backward via local GEMMs. -# --------------------------------------------------------------------------- - -EP_EXPERT_COMPUTE_MOE_ACT: Dict[str, Callable] = {} - -# Triton EP moe_act compute -try: - from xorl.ops.moe.triton import TritonEPGroupGemmMoeAct - - EP_EXPERT_COMPUTE_MOE_ACT["triton"] = TritonEPGroupGemmMoeAct.apply -except ImportError: - pass - -# Quack EP moe_act compute — adapt fused interface -try: - from xorl.ops.moe.quack import QuackEPGroupGemmMoeAct as _QuackEPGroupGemmMoeAct - - def _quack_ep_moe_act_fused(permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size): - gate_proj = gate_up_proj[..., :intermediate_size].contiguous() - up_proj = gate_up_proj[..., intermediate_size:].contiguous() - return _QuackEPGroupGemmMoeAct.apply(permute_tokens, cumsum, gate_proj, up_proj, down_proj) - - EP_EXPERT_COMPUTE_MOE_ACT["quack"] = _quack_ep_moe_act_fused -except ImportError: - pass - -# Native EP moe_act compute — adapt fused interface -try: - from .native import native_ep_compute_moe_act as _native_ep_compute_moe_act - - def _native_ep_moe_act_fused(permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size): - gate_proj = gate_up_proj[..., :intermediate_size].contiguous() - up_proj = gate_up_proj[..., intermediate_size:].contiguous() - return _native_ep_compute_moe_act(permute_tokens, cumsum, gate_proj, up_proj, down_proj) - - EP_EXPERT_COMPUTE_MOE_ACT["native"] = _native_ep_moe_act_fused -except ImportError: - pass - - -MOE_EXPERT_BACKENDS_MOE_ACT: Dict[str, Callable] = {} - -# Triton local moe_act compute -try: - from .triton_moe_act import triton_expert_forward_moe_act - - MOE_EXPERT_BACKENDS_MOE_ACT["triton"] = triton_expert_forward_moe_act -except ImportError: - pass - -# Quack local moe_act compute -try: - from .quack_moe_act import quack_expert_forward_moe_act - - MOE_EXPERT_BACKENDS_MOE_ACT["quack"] = quack_expert_forward_moe_act -except ImportError: - pass - -# Native local moe_act compute -try: - from .native import native_expert_forward_moe_act - - MOE_EXPERT_BACKENDS_MOE_ACT["native"] = native_expert_forward_moe_act -except ImportError: - pass - - __all__ = [ "MOE_EXPERT_BACKENDS", "EP_EXPERT_COMPUTE", @@ -262,6 +193,4 @@ def _native_ep_moe_act_fused(permute_tokens, cumsum, gate_up_proj, down_proj, in "EP_COMBINE", "EP_EXPERT_COMPUTE_LORA", "MOE_EXPERT_BACKENDS_LORA", - "EP_EXPERT_COMPUTE_MOE_ACT", - "MOE_EXPERT_BACKENDS_MOE_ACT", ] diff --git a/src/xorl/models/layers/moe/backend/native.py b/src/xorl/models/layers/moe/backend/native.py index 297c3769..118fb395 100644 --- a/src/xorl/models/layers/moe/backend/native.py +++ b/src/xorl/models/layers/moe/backend/native.py @@ -177,76 +177,6 @@ def _run_experts_grouped_mm( _run_experts_compiled = torch.compile(_run_experts_grouped_mm, fullgraph=True) -def _gate_up_swiglu( - x: torch.Tensor, - gate_proj: torch.Tensor, - up_proj: torch.Tensor, - offsets: torch.Tensor, -) -> torch.Tensor: - """Gate+Up SwiGLU computation: ``silu(x @ gate) * (x @ up)``. - - Factored out so it can be wrapped in ``torch.utils.checkpoint.checkpoint`` - for the moe_act variant — gate_output and up_output are NOT saved for - backward, they are recomputed from ``x``, ``gate_proj``, ``up_proj``. - - Fully compilable: only ``torch._grouped_mm``, ``F.silu``, and elementwise - multiply — no ``@torch.compiler.disable`` helpers. - """ - compute_dtype = torch.bfloat16 - x_bf16 = x.to(compute_dtype) - - gate_out = F.silu(torch._grouped_mm(x_bf16, gate_proj.to(compute_dtype), offs=offsets)) - up_out = torch._grouped_mm(x_bf16, up_proj.to(compute_dtype), offs=offsets) - - return gate_out * up_out - - -# Compiled variant for use inside checkpoint — fuses SiLU + elementwise multiply. -# torch.compile traces through checkpoint(use_reentrant=False) natively since -# PyTorch 2.4 (dynamo inlines the function and generates the recomputation graph). -_gate_up_swiglu_compiled = torch.compile(_gate_up_swiglu, fullgraph=True) - - -def _run_experts_moe_act( - gate_proj: torch.Tensor, - up_proj: torch.Tensor, - down_proj: torch.Tensor, - x: torch.Tensor, - padded_counts: torch.Tensor, - expert_scores: torch.Tensor | None = None, -) -> torch.Tensor: - """Run MoE experts with moe_act: gate+up SwiGLU is checkpointed so - gate_output and up_output are recomputed in backward instead of saved. - - Uses ``checkpoint(_gate_up_swiglu_compiled, ...)`` so the inner GEMMs + - activation are compiled (fuses SiLU + mul), and the checkpoint ensures - gate_out/up_out are recomputed in backward rather than saved. - """ - offsets = torch.cumsum(padded_counts, dim=0, dtype=torch.int32) - - # Checkpoint gate+up: in backward, recomputes silu(x @ gate) * (x @ up) - # instead of saving gate_out + up_out tensors. The compiled function is - # called twice (forward + recompute in backward) — this is intentional. - h = torch.utils.checkpoint.checkpoint( - _gate_up_swiglu_compiled, - x, - gate_proj, - up_proj, - offsets, - use_reentrant=False, - ) - if expert_scores is not None: - h = h * expert_scores.to(h.dtype).unsqueeze(-1) - - out = torch._grouped_mm( - h, - down_proj.to(torch.bfloat16), - offs=offsets, - ).to(x.dtype) - - return out - - # --------------------------------------------------------------------------- # Shared token preparation / scatter helper # --------------------------------------------------------------------------- @@ -657,69 +587,3 @@ def native_ep_compute_lora( ) return _unpad(out_padded, counts, padded_counts, num_local_experts, permute_tokens.shape[0]) - - -# --------------------------------------------------------------------------- -# moe_act variants: gate+up checkpointed, recomputed in backward -# --------------------------------------------------------------------------- - - -def native_ep_compute_moe_act( - permute_tokens: torch.Tensor, - cumsum: torch.Tensor, - gate_proj: torch.Tensor, - up_proj: torch.Tensor, - down_proj: torch.Tensor, - expert_scores: torch.Tensor | None = None, -) -> torch.Tensor: - """EP expert compute with moe_act using ``torch._grouped_mm``. - - Same interface as ``native_ep_compute``, but gate+up SwiGLU is - checkpointed — gate_output and up_output are recomputed in backward - instead of saved. Avoids recomputing EP all-to-all communication. - """ - if permute_tokens.shape[0] == 0: - return permute_tokens - - num_local_experts = gate_proj.shape[0] - counts = _cumsum_to_counts(cumsum, num_local_experts) - - padded_tokens, padded_counts = _pad_to_alignment(permute_tokens, counts, num_local_experts) - expert_scores_padded = None - if expert_scores is not None: - expert_scores_padded = padded_tokens.new_zeros(padded_tokens.shape[0]) - expert_scores_padded[ - _compute_pad_indices( - counts, padded_counts, num_local_experts, permute_tokens.shape[0], padded_tokens.device - ) - ] = expert_scores.to(padded_tokens.dtype) - out_padded = _run_experts_moe_act(gate_proj, up_proj, down_proj, padded_tokens, padded_counts, expert_scores_padded) - - return _unpad(out_padded, counts, padded_counts, num_local_experts, permute_tokens.shape[0]) - - -def native_expert_forward_moe_act( - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - gate_proj: torch.Tensor, - up_proj: torch.Tensor, - down_proj: torch.Tensor, - num_experts: int, - **kwargs, -) -> torch.Tensor: - """Local native forward with moe_act (gate+up checkpointed). - - Same as ``native_expert_forward`` but uses ``_run_experts_moe_act`` - instead of the compiled grouped GEMM. - """ - return _native_expert_forward_impl( - hidden_states, - routing_weights, - selected_experts, - gate_proj, - up_proj, - down_proj, - num_experts, - _run_experts_moe_act, - ) diff --git a/src/xorl/models/layers/moe/backend/quack_moe_act.py b/src/xorl/models/layers/moe/backend/quack_moe_act.py deleted file mode 100644 index a5766d3d..00000000 --- a/src/xorl/models/layers/moe/backend/quack_moe_act.py +++ /dev/null @@ -1,28 +0,0 @@ -"""Quack MoE expert backend — moe_act variant (activation recompute).""" - -import torch - -from xorl.ops.moe.quack import quack_moe_forward_moe_act - - -def quack_expert_forward_moe_act( - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - gate_proj: torch.Tensor, - up_proj: torch.Tensor, - down_proj: torch.Tensor, - num_experts: int, - **kwargs, -) -> torch.Tensor: - """Forward pass using quack group GEMM with moe_act (activation recompute).""" - return quack_moe_forward_moe_act( - module=None, - num_experts=num_experts, - routing_weights=routing_weights, - selected_experts=selected_experts, - hidden_states=hidden_states, - gate_proj=gate_proj, - up_proj=up_proj, - down_proj=down_proj, - ) diff --git a/src/xorl/models/layers/moe/backend/triton_moe_act.py b/src/xorl/models/layers/moe/backend/triton_moe_act.py deleted file mode 100644 index e3f34aac..00000000 --- a/src/xorl/models/layers/moe/backend/triton_moe_act.py +++ /dev/null @@ -1,28 +0,0 @@ -"""Triton MoE expert backend — moe_act variant (activation recompute).""" - -import torch - -from xorl.ops.moe.triton import triton_moe_forward_moe_act - - -def triton_expert_forward_moe_act( - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - gate_proj: torch.Tensor, - up_proj: torch.Tensor, - down_proj: torch.Tensor, - num_experts: int, - **kwargs, -) -> torch.Tensor: - """Forward pass using Triton group GEMM with moe_act (activation recompute).""" - return triton_moe_forward_moe_act( - module=None, - num_experts=num_experts, - routing_weights=routing_weights, - selected_experts=selected_experts, - hidden_states=hidden_states, - gate_proj=gate_proj, - up_proj=up_proj, - down_proj=down_proj, - ) diff --git a/src/xorl/models/layers/moe/experts.py b/src/xorl/models/layers/moe/experts.py index bf5eca52..86100c44 100644 --- a/src/xorl/models/layers/moe/experts.py +++ b/src/xorl/models/layers/moe/experts.py @@ -11,9 +11,7 @@ EP_COMBINE, EP_DISPATCH, EP_EXPERT_COMPUTE, - EP_EXPERT_COMPUTE_MOE_ACT, MOE_EXPERT_BACKENDS, - MOE_EXPERT_BACKENDS_MOE_ACT, ) from .common import split_gate_up_proj @@ -74,9 +72,6 @@ def __init__( ) self.act_fn = ACT2FN[hidden_act] - # Set by gradient_checkpointing_enable when moe_checkpoint_method="moe_act" - self._moe_act: bool = False - # EP dispatch strategy: "alltoall" (default) or "deepep" (NVLink-optimized) self.ep_dispatch: str = "alltoall" self.deepep_buffer_size_gb: float = 2.0 @@ -112,8 +107,6 @@ def forward( When Expert Parallelism is enabled, all backends (triton/native/quack) use the unified dispatch → compute → combine path via ``_ep_forward()``. """ - _moe_act = self._moe_act - if self.moe_implementation == "eager": fn = MOE_EXPERT_BACKENDS[self.moe_implementation] assert expert_idx is not None @@ -137,10 +130,7 @@ def forward( # Local single-GPU path gate_proj = self.gate_proj.contiguous() up_proj = self.up_proj.contiguous() - if _moe_act and self.moe_implementation in MOE_EXPERT_BACKENDS_MOE_ACT: - fn = MOE_EXPERT_BACKENDS_MOE_ACT[self.moe_implementation] - else: - fn = MOE_EXPERT_BACKENDS[self.moe_implementation] + fn = MOE_EXPERT_BACKENDS[self.moe_implementation] return fn( hidden_states, @@ -182,12 +172,7 @@ def _ep_forward( dispatch_fn = EP_DISPATCH[self.ep_dispatch] combine_fn = EP_COMBINE[self.ep_dispatch] - # Select moe_act compute variant when available - _moe_act = self._moe_act - if _moe_act and self.moe_implementation in EP_EXPERT_COMPUTE_MOE_ACT: - compute_fn = EP_EXPERT_COMPUTE_MOE_ACT[self.moe_implementation] - else: - compute_fn = EP_EXPERT_COMPUTE[self.moe_implementation] + compute_fn = EP_EXPERT_COMPUTE[self.moe_implementation] # Step 1: Dispatch tokens to expert-owning ranks dispatch_kwargs = self._build_dispatch_kwargs(hidden_states, routing_weights, selected_experts, parallel_state) diff --git a/src/xorl/models/layers/moe/moe_block.py b/src/xorl/models/layers/moe/moe_block.py index 13abe2af..c33acc90 100644 --- a/src/xorl/models/layers/moe/moe_block.py +++ b/src/xorl/models/layers/moe/moe_block.py @@ -128,20 +128,22 @@ def _regather_routing(self, router_logits, cached_experts, input_dtype): routing_weights = routing_weights.to(input_dtype) return cached_experts, routing_weights - def forward(self, hidden_states: torch.Tensor): - """Forward pass. + def route(self, hidden_states: torch.Tensor): + """Compute routing weights and expert assignments. + + Extracts the full routing logic: gate projection (with optional fp32 + upcast), routing replay stage handling, ``_regather_routing``, and + ``train_router`` detach. Args: - hidden_states: ``(batch, seq_len, hidden_dim)``. + hidden_states: ``(num_tokens, hidden_dim)`` — already flattened. Returns: - Tuple of ``(output, router_logits)`` where: - - output: ``(batch, seq_len, hidden_dim)`` + Tuple of ``(routing_weights, selected_experts, router_logits)`` where: + - routing_weights: ``(num_tokens, top_k)`` + - selected_experts: ``(num_tokens, top_k)`` - router_logits: ``(num_tokens, num_experts)`` """ - batch_size, sequence_length, hidden_dim = hidden_states.shape - hidden_states = hidden_states.view(-1, hidden_dim) - # Route (optionally upcast to fp32 for numerical alignment with SGLang) if getattr(self, "config", None) is not None and getattr(self.config, "_router_fp32", False): router_logits = F.linear(hidden_states.float(), self.gate.weight.float()) @@ -163,7 +165,6 @@ def forward(self, hidden_states: torch.Tensor): replay = self._routing_replay if stage is not None and replay is not None: - cached_weights = None if stage == "record": # Determine expert selection without creating autograd nodes with torch.no_grad(): @@ -171,19 +172,33 @@ def forward(self, hidden_states: torch.Tensor): replay.record(selected_experts) elif stage == "replay_forward": selected_experts = replay.pop_forward() - cached_weights = replay.pop_forward_weights() elif stage == "replay_backward": selected_experts = replay.pop_backward() - cached_weights = replay.pop_backward_weights() - if cached_weights is not None: - # Use pre-populated weights from inference (R3 weight replay) - routing_weights = cached_weights.to(hidden_states.dtype) - else: - # Uniform autograd path: softmax -> gather(cached indices) -> normalize - selected_experts, routing_weights = self._regather_routing( - router_logits, selected_experts, hidden_states.dtype - ) + # Uniform autograd path: softmax -> gather(cached indices) -> normalize. + # Both forward and recompute go through the same ops so non-reentrant + # checkpoint sees identical saved-tensor counts. + selected_experts, routing_weights = self._regather_routing( + router_logits, selected_experts, hidden_states.dtype + ) + + if stage == "record": + # Cache weights for replay_backward. During checkpoint recompute, + # router_logits may differ (non-deterministic attention), so the + # regathered weights above could be wrong. We override them below + # during replay_backward with these cached values. + replay.record_weights(routing_weights) + elif stage == "replay_backward": + # Override with cached weights to ensure deterministic EP dispatch. + # The _regather_routing call above still runs (for autograd graph + # structure), but its output is replaced with the recorded values. + cached_weights = replay.pop_backward_weights() + if cached_weights is not None: + routing_weights = cached_weights.to(hidden_states.dtype) + elif stage == "replay_forward": + cached_weights = replay.pop_forward_weights() + if cached_weights is not None: + routing_weights = cached_weights.to(hidden_states.dtype) else: # No replay active: use standard router routing_weights, selected_experts = self.router(router_logits, hidden_states.dtype) @@ -198,13 +213,48 @@ def forward(self, hidden_states: torch.Tensor): if not self.train_router: routing_weights = routing_weights.detach() - # Expert computation + return routing_weights, selected_experts, router_logits + + def forward_experts_only(self, hidden_states, routing_weights, selected_experts): + """Run expert computation with pre-computed routing weights. + + Args: + hidden_states: ``(batch, seq_len, hidden_dim)``. + routing_weights: ``(num_tokens, top_k)``. + selected_experts: ``(num_tokens, top_k)``. + + Returns: + ``(batch, seq_len, hidden_dim)`` expert output. + """ + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + hidden_states = self.experts(hidden_states, routing_weights, selected_experts) + hidden_states = hidden_states.reshape(batch_size, sequence_length, hidden_dim) + return hidden_states + + def forward(self, hidden_states: torch.Tensor): + """Forward pass. + + Args: + hidden_states: ``(batch, seq_len, hidden_dim)``. + + Returns: + Tuple of ``(output, router_logits)`` where: + - output: ``(batch, seq_len, hidden_dim)`` + - router_logits: ``(num_tokens, num_experts)`` + """ + batch_size, sequence_length, hidden_dim = hidden_states.shape + flat_hidden_states = hidden_states.view(-1, hidden_dim) + + routing_weights, selected_experts, router_logits = self.route(flat_hidden_states) + + # Expert computation (already flat — call experts directly to avoid extra reshape) if self.moe_implementation == "eager": - final_hidden_states = self._eager_forward(hidden_states, routing_weights, selected_experts) + final_hidden_states = self._eager_forward(flat_hidden_states, routing_weights, selected_experts) else: - final_hidden_states = self.experts(hidden_states, routing_weights, selected_experts) - + final_hidden_states = self.experts(flat_hidden_states, routing_weights, selected_experts) final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) + return final_hidden_states, router_logits def _eager_forward( diff --git a/src/xorl/models/module_utils.py b/src/xorl/models/module_utils.py index 00753aac..05940356 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -1605,6 +1605,98 @@ class GradientCheckpointingLayer(nn.Module): gradient_checkpointing = False def __call__(self, *args, **kwargs): - if self.gradient_checkpointing and self.training: + if ( + self.gradient_checkpointing + and self.training + and getattr(self, "_gradient_checkpointing_method", "recompute_full_layer") == "recompute_full_layer" + ): return self._gradient_checkpointing_func(partial(super().__call__, **kwargs), *args) return super().__call__(*args, **kwargs) + + +class MoEGradientCheckpointingLayer(nn.Module): + """Base class for MoE decoder layers with selective gradient checkpointing. + + Subclasses implement ``_pre_mlp_forward(hidden_states, **kwargs)`` which runs + layernorm → attention → layernorm and returns ``(hidden_states, residual)``. + + The checkpointing branching is handled automatically: + + - ``recompute_full_layer``: outer checkpoint wraps the entire layer + (handled by ``GradientCheckpointingLayer`` on the model-level loop). + - ``recompute_before_dispatch``: this class checkpoints attn+layernorm+router + via ``_pre_dispatch_forward``, then runs dispatch+expert+combine outside. + - ``no_recompute``: no checkpointing, runs normally. + + Works with both alltoall and DeepEP backends since dispatch/combine sit + outside the checkpoint boundary. + """ + + gradient_checkpointing = False + + def _pre_mlp_forward(self, hidden_states, **kwargs): + """Layernorm → attention → layernorm. Override per model. + + Returns: + ``(hidden_states, residual)`` + """ + raise NotImplementedError + + def _pre_dispatch_forward(self, hidden_states, **kwargs): + """Layernorm → attention → layernorm → router. + + Composed from ``_pre_mlp_forward`` + ``self.mlp.route()``. + Override only if the model needs custom routing (rare). + """ + hidden_states, residual = self._pre_mlp_forward(hidden_states, **kwargs) + # route() expects flattened (num_tokens, hidden_dim); hidden_states is 3-D here. + orig_shape = hidden_states.shape + flat = hidden_states.view(-1, orig_shape[-1]) + routing_weights, selected_experts, router_logits = self.mlp.route(flat) + return hidden_states, residual, routing_weights, selected_experts, router_logits + + def _moe_forward(self, hidden_states, output_router_logits=False, **kwargs): + """Forward with selective checkpointing. Called by subclass ``forward()``. + + Args: + hidden_states: ``(batch, seq_len, hidden_dim)``. + output_router_logits: Whether to include router logits in output. + **kwargs: Forwarded to ``_pre_mlp_forward`` (attention_mask, position_embeddings, etc.) + + Returns: + Tuple of ``(hidden_states, ...)`` with optional router_logits. + """ + from xorl.distributed.moe.deepep import sync_pending_combine + from xorl.models.layers.moe.moe_block import MoEBlock + + _selective = ( + self.training + and getattr(self, "gradient_checkpointing", False) + and getattr(self, "_gradient_checkpointing_method", "recompute_full_layer") != "recompute_full_layer" + ) + _is_moe = isinstance(self.mlp, MoEBlock) + + if _selective and _is_moe: + moe_input, residual, routing_weights, selected_experts, router_logits = self._gradient_checkpointing_func( + self._pre_dispatch_forward, hidden_states, **kwargs + ) + hidden_states = self.mlp.forward_experts_only(moe_input, routing_weights, selected_experts) + elif _selective: + hidden_states, residual = self._gradient_checkpointing_func(self._pre_mlp_forward, hidden_states, **kwargs) + hidden_states = self.mlp(hidden_states) + router_logits = None + else: + hidden_states, residual = self._pre_mlp_forward(hidden_states, **kwargs) + hidden_states = self.mlp(hidden_states) + if isinstance(hidden_states, tuple): + hidden_states, router_logits = hidden_states + else: + router_logits = None + + sync_pending_combine() + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + if output_router_logits: + outputs += (router_logits,) + return outputs diff --git a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py index a3a7fe62..52a54b97 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -4,7 +4,6 @@ import torch from torch import nn -from xorl.distributed.moe.deepep import sync_pending_combine from xorl.distributed.parallel_state import get_parallel_state from xorl.distributed.sequence_parallel.strategy import get_cp_strategy from xorl.models.base import XorlPreTrainedModel @@ -24,6 +23,7 @@ get_rmsnorm_mode, native_zero_centered_rms_norm, ) +from xorl.models.module_utils import MoEGradientCheckpointingLayer from xorl.models.outputs import MoeCausalLMOutput, MoeModelOutput from xorl.models.transformers.qwen3_5_moe import parallelize from xorl.models.transformers.qwen3_5_moe.checkpoint_handler import Qwen3_5MoeCheckpointHandler @@ -235,14 +235,21 @@ def __init__(self, config, moe_implementation="triton"): self.shared_expert = Qwen3_5MoeMLP(config, intermediate_size=config.shared_expert_intermediate_size) self.shared_expert_gate = nn.Linear(config.hidden_size, 1, bias=False) + def _shared_expert(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Shared expert: MLP + sigmoid gate.""" + flat = hidden_states.view(-1, hidden_states.size(-1)) + out = self.shared_expert(flat) + out = torch.sigmoid(self.shared_expert_gate(flat)) * out + return out.view_as(hidden_states) + + def forward_experts_only(self, hidden_states, routing_weights, selected_experts): + """Sparse experts + shared expert with pre-computed routing.""" + expert_output = super().forward_experts_only(hidden_states, routing_weights, selected_experts) + return expert_output + self._shared_expert(hidden_states) + def forward(self, hidden_states: torch.Tensor): expert_output, router_logits = super().forward(hidden_states) - batch_size, sequence_length, hidden_dim = hidden_states.shape - flat_hidden_states = hidden_states.view(-1, hidden_dim) - shared_expert_output = self.shared_expert(flat_hidden_states) - shared_expert_output = torch.sigmoid(self.shared_expert_gate(flat_hidden_states)) * shared_expert_output - shared_expert_output = shared_expert_output.view(batch_size, sequence_length, hidden_dim) - return expert_output + shared_expert_output, router_logits + return expert_output + self._shared_expert(hidden_states), router_logits QWEN3_5_MOE_CLASSES = { @@ -253,7 +260,7 @@ def forward(self, hidden_states: torch.Tensor): } -class Qwen3_5MoeDecoderLayer(nn.Module): +class Qwen3_5MoeDecoderLayer(MoEGradientCheckpointingLayer): def __init__(self, config: Qwen3_5MoeConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size @@ -287,18 +294,16 @@ def __init__(self, config: Qwen3_5MoeConfig, layer_idx: int): else: self.mlp = Qwen3_5MoeMLP(config, intermediate_size=config.intermediate_size) - def forward( + def _pre_mlp_forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values=None, - use_cache: bool | None = False, - output_attentions: Optional[bool] = False, - output_router_logits: Optional[bool] = False, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, **kwargs: Unpack[AttentionKwargs], - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Layernorm → attention → layernorm.""" residual = hidden_states hidden_states = self.input_layernorm(hidden_states) @@ -315,15 +320,15 @@ def forward( linear_mask = attention_mask if attention_mask is not None and attention_mask.dim() == 2 else None if cp_context is not None: linear_mask = None - hidden_states, self_attn_weights, _ = self.linear_attn( + hidden_states, _, _ = self.linear_attn( hidden_states=hidden_states, attention_mask=linear_mask, past_key_values=past_key_values, - use_cache=use_cache, + use_cache=False, **linear_kwargs, ) else: - hidden_states, self_attn_weights = self.self_attn( + hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, @@ -332,20 +337,29 @@ def forward( **kwargs, ) hidden_states, residual = self.post_attention_layernorm(hidden_states, residual=residual, prenorm=True) - hidden_states = self.mlp(hidden_states) - if isinstance(hidden_states, tuple): - hidden_states, router_logits = hidden_states - else: - router_logits = None - sync_pending_combine() - hidden_states = residual + hidden_states + return hidden_states, residual - outputs = (hidden_states,) - if output_attentions: - outputs += (self_attn_weights,) - if output_router_logits: - outputs += (router_logits,) - return outputs + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values=None, + use_cache: bool | None = False, + output_attentions: Optional[bool] = False, + output_router_logits: Optional[bool] = False, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs: Unpack[AttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + return self._moe_forward( + hidden_states, + output_router_logits=output_router_logits, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + position_embeddings=position_embeddings, + **kwargs, + ) class Qwen3_5MoePreTrainedModel(XorlPreTrainedModel): @@ -503,7 +517,14 @@ def forward( if decoder_layer is None: continue layer_mask = linear_attn_mask if decoder_layer.layer_type == "linear_attention" else causal_mask - if self.gradient_checkpointing and self.training: + # When selective checkpointing is enabled, the decoder layer handles + # its own sub-checkpointing — skip the outer checkpoint. + _use_outer_checkpoint = ( + self.gradient_checkpointing + and self.training + and getattr(self, "_gradient_checkpointing_method", "recompute_full_layer") == "recompute_full_layer" + ) + if _use_outer_checkpoint: layer_outputs = self._gradient_checkpointing_func( decoder_layer.__call__, hidden_states, @@ -530,7 +551,8 @@ def forward( ) hidden_states = layer_outputs[0] if output_attentions: - all_self_attns += (layer_outputs[1],) + # _moe_forward does not produce attention weights; use None placeholder. + all_self_attns += (None,) if output_router_logits: all_router_logits += (layer_outputs[-1],) diff --git a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py index 33a61fe1..5e11651e 100644 --- a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py +++ b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py @@ -12,12 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple, Unpack +from typing import Optional, Unpack import torch from torch import nn -from xorl.distributed.moe.deepep import sync_pending_combine from xorl.distributed.parallel_state import get_parallel_state from xorl.distributed.sequence_parallel.strategy import get_cp_strategy from xorl.models.base import XorlPreTrainedModel @@ -35,6 +34,7 @@ update_causal_mask, ) from xorl.models.layers.moe import MoEBlock, MoEExperts +from xorl.models.module_utils import MoEGradientCheckpointingLayer from xorl.models.outputs import MoeCausalLMOutput, MoeModelOutput from xorl.models.transformers.qwen3_moe import parallelize from xorl.models.transformers.qwen3_moe.checkpoint_handler import Qwen3MoeCheckpointHandler @@ -215,15 +215,14 @@ def __init__(self, config): } -class Qwen3MoeDecoderLayer(nn.Module): +class Qwen3MoeDecoderLayer(MoEGradientCheckpointingLayer): def __init__(self, config: Qwen3MoeConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size + self.gradient_checkpointing = False # set by gradient_checkpointing_enable self.self_attn = Qwen3MoeAttention(config, layer_idx) - self.mlp = Qwen3MoeMLP(config) - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -235,78 +234,36 @@ def __init__(self, config: Qwen3MoeConfig, layer_idx: int): else: self.mlp = Qwen3MoeMLP(config, intermediate_size=config.intermediate_size) - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - output_attentions: Optional[bool] = False, - output_router_logits: Optional[bool] = False, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - **kwargs: Unpack[AttentionKwargs], - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: - _selective = ( - self.training - and getattr(self, "gradient_checkpointing", False) - and getattr(self, "_recompute_modules", None) is not None - ) + def _pre_mlp_forward(self, hidden_states, attention_mask=None, position_embeddings=None, **kwargs): residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual=residual, prenorm=True) + return hidden_states, residual - # Self Attention - if _selective and "self_attn" in self._recompute_modules: - # MultiHeadAttention.forward positional order: hidden_states, position_embeddings, attention_mask - hidden_states, self_attn_weights = self._gradient_checkpointing_func( - self.self_attn.__call__, - hidden_states, - position_embeddings, - attention_mask, - **kwargs, - ) - else: - hidden_states, self_attn_weights = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - position_embeddings=position_embeddings, - **kwargs, - ) - # Fully Connected - hidden_states, residual = self.post_attention_layernorm( + def forward( + self, + hidden_states, + attention_mask=None, + position_ids=None, + output_attentions=False, + output_router_logits=False, + position_embeddings=None, + **kwargs, + ): + return self._moe_forward( hidden_states, - residual=residual, - prenorm=True, + output_router_logits=output_router_logits, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + **kwargs, ) - if _selective and "mlp" in self._recompute_modules: - hidden_states = self._gradient_checkpointing_func( - self.mlp.__call__, - hidden_states, - ) - else: - hidden_states = self.mlp(hidden_states) - - if isinstance(hidden_states, tuple): - hidden_states, router_logits = hidden_states - else: - router_logits = None - - # Sync any pending async DeepEP combine before reading MoE output. - # No-op when async combine is disabled or non-DeepEP dispatch. - sync_pending_combine() - - hidden_states = residual + hidden_states - - outputs = (hidden_states,) - - if output_attentions: - outputs += (self_attn_weights,) - - if output_router_logits: - outputs += (router_logits,) - - return outputs - class Qwen3MoePreTrainedModel(XorlPreTrainedModel): config_class = Qwen3MoeConfig @@ -481,16 +438,15 @@ def forward( all_self_attns = () if output_attentions else None all_router_logits = () if output_router_logits else None + _grad_ckpt_method = getattr(self, "_gradient_checkpointing_method", None) or "recompute_full_layer" + _grad_ckpt_active = self.gradient_checkpointing and self.training + for decoder_layer in self.layers: if decoder_layer is None: # PP: pruned layer continue - # When selective checkpointing is enabled (_recompute_modules is set), - # the decoder layer handles its own sub-checkpointing — skip the outer checkpoint. - _use_outer_checkpoint = ( - self.gradient_checkpointing and self.training and getattr(self, "_recompute_modules", None) is None - ) - if _use_outer_checkpoint: + if _grad_ckpt_active and _grad_ckpt_method == "recompute_full_layer": + # Recompute entire layer in backward (including dispatch + combine) layer_outputs = self._gradient_checkpointing_func( decoder_layer.__call__, hidden_states, @@ -501,7 +457,20 @@ def forward( position_embeddings, **kwargs, ) + elif _grad_ckpt_active and _grad_ckpt_method == "recompute_before_dispatch": + # Decoder layer handles checkpoint internally via _pre_dispatch_forward. + # Dispatch + combine run outside checkpoint (alltoall not recomputed). + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_router_logits=output_router_logits, + position_embeddings=position_embeddings, + **kwargs, + ) else: + # no_recompute or gc disabled: run decoder layer normally layer_outputs = decoder_layer( hidden_states, attention_mask=causal_mask, @@ -515,7 +484,8 @@ def forward( hidden_states = layer_outputs[0] if output_attentions: - all_self_attns += (layer_outputs[1],) + # _moe_forward does not produce attention weights; use None placeholder. + all_self_attns += (None,) if output_router_logits: all_router_logits += (layer_outputs[-1],) diff --git a/src/xorl/ops/moe/quack.py b/src/xorl/ops/moe/quack.py index 83a5dab7..9237c14e 100644 --- a/src/xorl/ops/moe/quack.py +++ b/src/xorl/ops/moe/quack.py @@ -24,136 +24,6 @@ def _scatter_and_cumsum(hidden_states: torch.Tensor, expert_index: torch.Tensor, return scatter_output, scatter_index, cumsum_t -class QuackEPGroupGemmMoeAct(torch.autograd.Function): - """EP expert MLP with moe_act: saves only inputs + weights, recomputes - gate_output and up_output in backward (2 local GEMMs, no EP communication).""" - - @staticmethod - def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores=None): - max_M = permute_tokens.shape[0] - cu_seqlens = cumsum_to_cu_seqlens(cumsum) - ctx.has_expert_scores = expert_scores is not None - - gate_output = quack_group_gemm_same_nk( - a=permute_tokens, b=gate_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=False, cu_seqlens_m=cu_seqlens - ) - up_output = quack_group_gemm_same_nk( - a=permute_tokens, b=up_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=False, cu_seqlens_m=cu_seqlens - ) - - gate_activation = torch.ops.aten.silu(gate_output) - gated_output = gate_activation * up_output - if expert_scores is not None: - gated_output = gated_output * expert_scores.to(gated_output.dtype).unsqueeze(-1) - del gate_activation - - down_output = quack_group_gemm_same_nk( - a=gated_output, b=down_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=False, cu_seqlens_m=cu_seqlens - ) - del gated_output - - # moe_act: save only 5 tensors (drop gate_output, up_output) - if expert_scores is None: - expert_scores = permute_tokens.new_ones(permute_tokens.shape[0]) - ctx.save_for_backward(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores) - return down_output - - @staticmethod - def backward(ctx, grad_output): - permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores = ctx.saved_tensors - max_M = grad_output.shape[0] - cu_seqlens_m = cumsum_to_cu_seqlens(cumsum) - - # Recompute gate_output and up_output from saved inputs + weights - gate_output = quack_group_gemm_same_nk( - a=permute_tokens, b=gate_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=False, cu_seqlens_m=cu_seqlens_m - ) - up_output = quack_group_gemm_same_nk( - a=permute_tokens, b=up_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=False, cu_seqlens_m=cu_seqlens_m - ) - - # Rest identical to QuackEPGroupGemm.backward - gate_activation = torch.ops.aten.silu(gate_output) - gated_output = gate_activation * up_output - expert_scores_dtype = expert_scores.dtype - expert_scores = expert_scores.to(gated_output.dtype) - gated_weighted = gated_output * expert_scores.unsqueeze(-1) - - # dgrad FC2 - grad_gated_weighted = quack_group_gemm_same_nk( - a=grad_output, b=down_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=True, cu_seqlens_m=cu_seqlens_m - ) - - # wgrad FC2 - grad_down_proj = None - if down_proj.requires_grad: - grad_down_proj = torch.empty_like(down_proj) - quack_group_gemm_same_mn( - a=gated_weighted, - b=grad_output, - c=grad_down_proj, - cumsum_K=cumsum, - max_K=max_M, - transpose_a=True, - transpose_b=False, - cu_seqlens_k=cu_seqlens_m, - ) - grad_expert_scores = None - if ctx.has_expert_scores: - grad_expert_scores = (grad_gated_weighted * gated_output).sum(dim=-1).to(expert_scores_dtype) - del gated_output, gated_weighted - - grad_gated_output = grad_gated_weighted * expert_scores.unsqueeze(-1) - del grad_gated_weighted - - # Activation backward - grad_up_output = gate_activation * grad_gated_output - grad_gate_activation = grad_gated_output * up_output - del grad_gated_output, gate_activation, up_output - grad_gate_output = torch.ops.aten.silu_backward(grad_gate_activation, gate_output) - del grad_gate_activation, gate_output - - # dgrad FC1: in-place add - grad_permute_tokens = quack_group_gemm_same_nk( - a=grad_gate_output, b=gate_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=True, cu_seqlens_m=cu_seqlens_m - ) - grad_permute_tokens += quack_group_gemm_same_nk( - a=grad_up_output, b=up_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=True, cu_seqlens_m=cu_seqlens_m - ) - - # wgrad FC1 - grad_gate_proj = None - if gate_proj.requires_grad: - grad_gate_proj = torch.empty_like(gate_proj) - quack_group_gemm_same_mn( - a=permute_tokens, - b=grad_gate_output, - c=grad_gate_proj, - cumsum_K=cumsum, - max_K=max_M, - transpose_a=True, - transpose_b=False, - cu_seqlens_k=cu_seqlens_m, - ) - del grad_gate_output - grad_up_proj = None - if up_proj.requires_grad: - grad_up_proj = torch.empty_like(up_proj) - quack_group_gemm_same_mn( - a=permute_tokens, - b=grad_up_output, - c=grad_up_proj, - cumsum_K=cumsum, - max_K=max_M, - transpose_a=True, - transpose_b=False, - cu_seqlens_k=cu_seqlens_m, - ) - del grad_up_output - - return grad_permute_tokens, None, grad_gate_proj, grad_up_proj, grad_down_proj, grad_expert_scores - - class QuackMoeExpertsFunction(torch.autograd.Function): """Memory-optimized: separate gate/up GEMMs, recompute cheap intermediates, explicit del for dead tensors, in-place add for dgrad.""" @@ -679,179 +549,3 @@ def quack_moe_forward( return QuackMoeExpertsFunction.apply( num_experts, routing_weights, selected_experts, hidden_states, gate_proj, up_proj, down_proj, gate_up_weight ) - - -class QuackMoeExpertsFunctionMoeAct(torch.autograd.Function): - """Local MoE expert computation with moe_act: drops gate_output/up_output - from save_for_backward and recomputes them in backward.""" - - @staticmethod - def forward( - ctx, num_experts, gate_weights, expert_index, hidden_states, gate_proj, up_proj, down_proj, gate_up_weight=None - ): - scatter_output, scatter_index, cumsum_t = _scatter_and_cumsum(hidden_states, expert_index, num_experts) - max_M = scatter_output.shape[0] - cu_seqlens = cumsum_to_cu_seqlens(cumsum_t) - - gate_output = quack_group_gemm_same_nk( - a=scatter_output, b=gate_proj, cumsum_M=cumsum_t, max_M=max_M, transpose_b=False, cu_seqlens_m=cu_seqlens - ) - up_output = quack_group_gemm_same_nk( - a=scatter_output, b=up_proj, cumsum_M=cumsum_t, max_M=max_M, transpose_b=False, cu_seqlens_m=cu_seqlens - ) - del scatter_output - - gate_activation = torch.ops.aten.silu(gate_output) - gated_activation = gate_activation * up_output - del gate_activation - - scattered_gate_weight = torch.empty_like(gate_weights.reshape(-1, 1)) - scattered_gate_weight[scatter_index.flatten()] = gate_weights.reshape(-1, 1) - gated_weighted = gated_activation * scattered_gate_weight - del gated_activation - - down_output = quack_group_gemm_same_nk( - a=gated_weighted, b=down_proj, cumsum_M=cumsum_t, max_M=max_M, transpose_b=False, cu_seqlens_m=cu_seqlens - ) - del gated_weighted - output = moe_gather(down_output, scatter_index).reshape(hidden_states.shape) - del down_output - - # moe_act: save 8 tensors (drop gate_output, up_output vs 10 in standard) - ctx.save_for_backward( - gate_weights, - gate_proj, - up_proj, - down_proj, - hidden_states, - scatter_index, - cumsum_t, - scattered_gate_weight, - ) - return output - - @staticmethod - def backward(ctx, grad_output): - ( - gate_weights, - gate_proj, - up_proj, - down_proj, - hidden_states, - scatter_index, - cumsum_t, - scattered_gate_weight, - ) = ctx.saved_tensors - grad_output = grad_output.view(-1, grad_output.shape[-1]) - max_M = grad_output.shape[0] - cu_seqlens_m = cumsum_to_cu_seqlens(cumsum_t) - - # Recompute scatter_output, gate_output, up_output from saved inputs + weights - scatter_output = moe_scatter(hidden_states, scatter_index) - gate_output = quack_group_gemm_same_nk( - a=scatter_output, b=gate_proj, cumsum_M=cumsum_t, max_M=max_M, transpose_b=False, cu_seqlens_m=cu_seqlens_m - ) - up_output = quack_group_gemm_same_nk( - a=scatter_output, b=up_proj, cumsum_M=cumsum_t, max_M=max_M, transpose_b=False, cu_seqlens_m=cu_seqlens_m - ) - - gate_activation = torch.ops.aten.silu(gate_output) - gated_activation = gate_activation * up_output - gated_weighted = gated_activation * scattered_gate_weight - - grad_down_output = moe_scatter(grad_output, scatter_index) - - # dgrad FC2 - grad_gated_weighted = quack_group_gemm_same_nk( - a=grad_down_output, b=down_proj, cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, cu_seqlens_m=cu_seqlens_m - ) - - # wgrad FC2 - grad_down_proj = None - if down_proj.requires_grad: - grad_down_proj = torch.empty_like(down_proj) - quack_group_gemm_same_mn( - a=gated_weighted, - b=grad_down_output, - c=grad_down_proj, - cumsum_K=cumsum_t, - max_K=max_M, - transpose_a=True, - transpose_b=False, - cu_seqlens_k=cu_seqlens_m, - ) - del grad_down_output, gated_weighted - - # Routing weight gradient - grad_gated_activation = grad_gated_weighted * scattered_gate_weight - grad_gate_weight = torch.sum(gated_activation * grad_gated_weighted, dim=-1)[scatter_index.flatten()] - grad_gate_weight = grad_gate_weight.reshape(gate_weights.shape) - del gated_activation, grad_gated_weighted - - # Activation backward - grad_up_output = gate_activation * grad_gated_activation - grad_gate_activation = grad_gated_activation * up_output - del grad_gated_activation, gate_activation, up_output - grad_gate_output = torch.ops.aten.silu_backward(grad_gate_activation, gate_output) - del grad_gate_activation, gate_output - - # dgrad FC1: in-place add - grad_scatter_output = quack_group_gemm_same_nk( - a=grad_gate_output, b=gate_proj, cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, cu_seqlens_m=cu_seqlens_m - ) - grad_scatter_output += quack_group_gemm_same_nk( - a=grad_up_output, b=up_proj, cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, cu_seqlens_m=cu_seqlens_m - ) - - # wgrad FC1 - grad_gate_proj = None - if gate_proj.requires_grad: - grad_gate_proj = torch.empty_like(gate_proj) - quack_group_gemm_same_mn( - a=scatter_output, - b=grad_gate_output, - c=grad_gate_proj, - cumsum_K=cumsum_t, - max_K=max_M, - transpose_a=True, - transpose_b=False, - cu_seqlens_k=cu_seqlens_m, - ) - del grad_gate_output - grad_up_proj = None - if up_proj.requires_grad: - grad_up_proj = torch.empty_like(up_proj) - quack_group_gemm_same_mn( - a=scatter_output, - b=grad_up_output, - c=grad_up_proj, - cumsum_K=cumsum_t, - max_K=max_M, - transpose_a=True, - transpose_b=False, - cu_seqlens_k=cu_seqlens_m, - ) - del grad_up_output, scatter_output - - grad_hidden_states = moe_gather(grad_scatter_output, scatter_index).reshape(hidden_states.shape) - return None, grad_gate_weight, None, grad_hidden_states, grad_gate_proj, grad_up_proj, grad_down_proj, None - - -def quack_moe_forward_moe_act( - module: torch.nn.Module, - num_experts: int, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - hidden_states: torch.Tensor, - gate_proj: torch.Tensor, - up_proj: torch.Tensor, - down_proj: torch.Tensor, - gate_up_weight: torch.Tensor = None, - **kwargs, -): - """Forward pass for MoE experts using quack group GEMM with moe_act - (activation recompute, no EP recompute).""" - del module - return QuackMoeExpertsFunctionMoeAct.apply( - num_experts, routing_weights, selected_experts, hidden_states, gate_proj, up_proj, down_proj, gate_up_weight - ) diff --git a/src/xorl/ops/moe/triton.py b/src/xorl/ops/moe/triton.py index f7c5f854..f06c62c1 100644 --- a/src/xorl/ops/moe/triton.py +++ b/src/xorl/ops/moe/triton.py @@ -151,148 +151,6 @@ def backward(ctx, grad_output): ) -class TritonEPGroupGemmMoeAct(torch.autograd.Function): - """EP expert MLP with moe_act and fused gate+up GEMM. - - Saves only inputs + weights (no activation outputs), recomputes via - fused GEMM in backward. Zero extra memory from weight copies. - """ - - @staticmethod - def forward(ctx, permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores=None): - max_M = permute_tokens.shape[0] - I = intermediate_size - ctx.has_expert_scores = expert_scores is not None - - gate_up_output = group_gemm_same_nk( - a=permute_tokens, - b=gate_up_proj, - cumsum_M=cumsum, - max_M=max_M, - transpose_a=False, - transpose_b=False, - ) - - gate_activation = torch.ops.aten.silu(gate_up_output[..., :I]) - gated_output = gate_activation * gate_up_output[..., I:] - if expert_scores is not None: - gated_output = gated_output * expert_scores.to(gated_output.dtype).unsqueeze(-1) - del gate_activation, gate_up_output - - down_output = group_gemm_same_nk( - a=gated_output, - b=down_proj, - cumsum_M=cumsum, - max_M=max_M, - transpose_a=False, - transpose_b=False, - ) - del gated_output - - if expert_scores is None: - expert_scores = permute_tokens.new_ones(permute_tokens.shape[0]) - ctx.save_for_backward(permute_tokens, cumsum, gate_up_proj, down_proj, expert_scores) - ctx.intermediate_size = I - return down_output - - @staticmethod - def backward(ctx, grad_output): - permute_tokens, cumsum, gate_up_proj, down_proj, expert_scores = ctx.saved_tensors - I = ctx.intermediate_size - max_M = grad_output.shape[0] - - # Recompute via fused GEMM - gate_up_output = group_gemm_same_nk( - a=permute_tokens, - b=gate_up_proj, - cumsum_M=cumsum, - max_M=max_M, - transpose_a=False, - transpose_b=False, - ) - gate_output = gate_up_output[..., :I] - up_output = gate_up_output[..., I:] - - gate_activation = torch.ops.aten.silu(gate_output) - gated_output = gate_activation * up_output - expert_scores_dtype = expert_scores.dtype - expert_scores = expert_scores.to(gated_output.dtype) - gated_weighted = gated_output * expert_scores.unsqueeze(-1) - - # dgrad FC2 - grad_gated_weighted = group_gemm_same_nk( - a=grad_output, - b=down_proj, - cumsum_M=cumsum, - max_M=max_M, - transpose_b=True, - ) - - # wgrad FC2 - grad_down_proj = None - if down_proj.requires_grad: - grad_down_proj = torch.empty_like(down_proj) - group_gemm_same_mn( - a=gated_weighted, - b=grad_output, - c=grad_down_proj, - cumsum_K=cumsum, - max_K=max_M, - transpose_a=True, - transpose_b=False, - ) - grad_expert_scores = None - if ctx.has_expert_scores: - grad_expert_scores = (grad_gated_weighted * gated_output).sum(dim=-1).to(expert_scores_dtype) - del gated_output, gated_weighted - - grad_gated_output = grad_gated_weighted * expert_scores.unsqueeze(-1) - del grad_gated_weighted - - # Activation backward - grad_up_output = gate_activation * grad_gated_output - grad_gate_activation = grad_gated_output * up_output - del grad_gated_output, gate_activation, up_output, gate_up_output - - grad_gate_output = torch.ops.aten.silu_backward(grad_gate_activation, gate_output) - del grad_gate_activation, gate_output - - # Fused dgrad FC1: cat grads → single GEMM with gate_up_proj (no .contiguous() copies) - grad_gate_up_act = torch.cat([grad_gate_output, grad_up_output], dim=-1) - del grad_gate_output, grad_up_output - grad_permute_tokens = group_gemm_same_nk( - a=grad_gate_up_act, - b=gate_up_proj, - cumsum_M=cumsum, - max_M=max_M, - transpose_b=True, - ) - - # Fused wgrad FC1: single GEMM produces grad_gate_up_proj directly - grad_gate_up_proj = None - if gate_up_proj.requires_grad: - grad_gate_up_proj = torch.empty_like(gate_up_proj) - group_gemm_same_mn( - a=permute_tokens, - b=grad_gate_up_act, - c=grad_gate_up_proj, - cumsum_K=cumsum, - max_K=max_M, - transpose_a=True, - transpose_b=False, - ) - del grad_gate_up_act - - return ( - grad_permute_tokens, - None, # cumsum - grad_gate_up_proj, - grad_down_proj, - None, # intermediate_size - grad_expert_scores, - ) - - class TritonMoeExpertsFunction(torch.autograd.Function): """MoE expert computation with custom autograd for efficient backward pass. @@ -545,242 +403,3 @@ def triton_moe_forward( down_proj, gate_up_weight, ) - - -class TritonMoeExpertsFunctionMoeAct(torch.autograd.Function): - """Local MoE expert computation with moe_act: drops gate_output/up_output - from save_for_backward and recomputes them in backward.""" - - @staticmethod - def forward( - ctx, - num_experts, - gate_weights, - expert_index, - hidden_states, - gate_proj, - up_proj, - down_proj, - gate_up_weight=None, - ): - splits = expert_histogram(expert_index, num_experts) - cumsum_t = torch.cumsum(splits, dim=0) - scatter_index = moe_index_compute(expert_index, cumsum_t) - scatter_output = moe_scatter(hidden_states, scatter_index) - max_M = scatter_output.shape[0] - - gate_output = group_gemm_same_nk( - a=scatter_output, - b=gate_proj, - cumsum_M=cumsum_t, - max_M=max_M, - transpose_a=False, - transpose_b=False, - ) - up_output = group_gemm_same_nk( - a=scatter_output, - b=up_proj, - cumsum_M=cumsum_t, - max_M=max_M, - transpose_a=False, - transpose_b=False, - ) - del scatter_output - - gate_activation = torch.ops.aten.silu(gate_output) - gated_activation = gate_activation * up_output - del gate_activation - - reshaped_gate_weight = gate_weights.reshape(-1, 1) - scattered_gate_weight = torch.empty_like(reshaped_gate_weight) - scattered_gate_weight[scatter_index.flatten()] = reshaped_gate_weight - gated_weighted = gated_activation * scattered_gate_weight - del gated_activation - - down_output = group_gemm_same_nk( - a=gated_weighted, - b=down_proj, - cumsum_M=cumsum_t, - max_M=max_M, - transpose_a=False, - transpose_b=False, - ) - del gated_weighted - - output = moe_gather(down_output, scatter_index).reshape(hidden_states.shape) - del down_output - - # moe_act: save 8 tensors (drop gate_output, up_output vs 10 in standard) - ctx.save_for_backward( - gate_weights, - gate_proj, - up_proj, - down_proj, - hidden_states, - scatter_index, - cumsum_t, - scattered_gate_weight, - ) - - return output - - @staticmethod - def backward(ctx, grad_output): - ( - gate_weights, - gate_proj, - up_proj, - down_proj, - hidden_states, - scatter_index, - cumsum_t, - scattered_gate_weight, - ) = ctx.saved_tensors - grad_output = grad_output.view(-1, grad_output.shape[-1]) - max_M = grad_output.shape[0] - - # Recompute scatter_output, gate_output, up_output from saved inputs + weights - scatter_output = moe_scatter(hidden_states, scatter_index) - gate_output = group_gemm_same_nk( - a=scatter_output, - b=gate_proj, - cumsum_M=cumsum_t, - max_M=max_M, - transpose_a=False, - transpose_b=False, - ) - up_output = group_gemm_same_nk( - a=scatter_output, - b=up_proj, - cumsum_M=cumsum_t, - max_M=max_M, - transpose_a=False, - transpose_b=False, - ) - - gate_activation = torch.ops.aten.silu(gate_output) - gated_activation = gate_activation * up_output - gated_weighted = gated_activation * scattered_gate_weight - - # Scatter grad to expert-sorted layout - grad_down_output = moe_scatter(grad_output, scatter_index) - - # FC2 dgrad - grad_gated_weighted = group_gemm_same_nk( - a=grad_down_output, - b=down_proj, - cumsum_M=cumsum_t, - max_M=max_M, - transpose_b=True, - ) - - # FC2 wgrad - grad_down_proj = None - if down_proj.requires_grad: - grad_down_proj = torch.empty_like(down_proj) - group_gemm_same_mn( - a=gated_weighted, - b=grad_down_output, - c=grad_down_proj, - cumsum_K=cumsum_t, - max_K=max_M, - transpose_a=True, - transpose_b=False, - ) - del grad_down_output, gated_weighted - - # Routing weight gradient - grad_gated_activation = grad_gated_weighted * scattered_gate_weight - grad_gate_weight = torch.sum(gated_activation * grad_gated_weighted, dim=-1)[scatter_index.flatten()] - grad_gate_weight = grad_gate_weight.reshape(gate_weights.shape) - del gated_activation, grad_gated_weighted - - # Activation backward - grad_up_output = gate_activation * grad_gated_activation - grad_gate_activation = grad_gated_activation * up_output - del grad_gated_activation, gate_activation, up_output - grad_gate_output = torch.ops.aten.silu_backward(grad_gate_activation, gate_output) - del grad_gate_activation, gate_output - - # FC1 dgrad - grad_scatter_output = group_gemm_same_nk( - a=grad_gate_output, - b=gate_proj, - cumsum_M=cumsum_t, - max_M=max_M, - transpose_b=True, - ) - grad_scatter_output += group_gemm_same_nk( - a=grad_up_output, - b=up_proj, - cumsum_M=cumsum_t, - max_M=max_M, - transpose_b=True, - ) - - # FC1 wgrad - grad_gate_proj = None - if gate_proj.requires_grad: - grad_gate_proj = torch.empty_like(gate_proj) - group_gemm_same_mn( - a=scatter_output, - b=grad_gate_output, - c=grad_gate_proj, - cumsum_K=cumsum_t, - max_K=max_M, - transpose_a=True, - transpose_b=False, - ) - del grad_gate_output - grad_up_proj = None - if up_proj.requires_grad: - grad_up_proj = torch.empty_like(up_proj) - group_gemm_same_mn( - a=scatter_output, - b=grad_up_output, - c=grad_up_proj, - cumsum_K=cumsum_t, - max_K=max_M, - transpose_a=True, - transpose_b=False, - ) - del grad_up_output, scatter_output - - grad_hidden_states = moe_gather(grad_scatter_output, scatter_index).reshape(hidden_states.shape) - - return ( - None, # num_experts - grad_gate_weight, - None, # expert_index - grad_hidden_states, - grad_gate_proj, - grad_up_proj, - grad_down_proj, - None, # gate_up_weight - ) - - -def triton_moe_forward_moe_act( - module: torch.nn.Module, - num_experts: int, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - hidden_states: torch.Tensor, - gate_proj: torch.Tensor, - up_proj: torch.Tensor, - down_proj: torch.Tensor, - gate_up_weight: torch.Tensor = None, - **kwargs, -): - """Forward pass for MoE experts using Triton group GEMM with moe_act - (activation recompute, no EP recompute).""" - return TritonMoeExpertsFunctionMoeAct.apply( - num_experts, - routing_weights, - selected_experts, - hidden_states, - gate_proj, - up_proj, - down_proj, - gate_up_weight, - ) diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index b19013c9..6467cfda 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -518,8 +518,7 @@ def _parallelize(self) -> None: enable_compile=args.train.enable_compile, basic_modules=self.model._no_split_modules + args.model.basic_modules, enable_reentrant=args.train.enable_reentrant, - recompute_modules=args.train.recompute_modules, - moe_checkpoint_method=args.train.moe_checkpoint_method, + gradient_checkpointing_method=args.train.gradient_checkpointing_method, enable_forward_prefetch=args.train.enable_forward_prefetch, load_weights_mode=args.train.load_weights_mode, pp_schedule=args.train.pipeline_parallel_schedule if args.train.pipeline_parallel_size > 1 else None, @@ -641,9 +640,8 @@ def _setup_observability(self) -> None: config=self.model_config, global_batch_size=args.train.global_batch_size, empty_cache_steps=args.train.empty_cache_steps, - gc_enabled=args.train.enable_gradient_checkpointing, - recompute_modules=args.train.recompute_modules, - moe_checkpoint_method=args.train.moe_checkpoint_method, + gradient_checkpointing_enabled=args.train.enable_gradient_checkpointing, + gradient_checkpointing_method=args.train.gradient_checkpointing_method, cp_size=args.train.ulysses_parallel_size * args.train.ringattn_parallel_size, ) self._maybe_log_startup_metrics( @@ -1036,11 +1034,12 @@ def _maybe_log( use_tqdm: bool, tqdm_bar: Optional[Any], ) -> None: - """Log metrics to stdout / tqdm / wandb.""" + """Log metrics to stdout / tqdm / wandb / structured JSON.""" args = self.args tflops_per_gpu = train_metrics.get("efficiency/flops_achieved(T)", 0) / args.train.world_size mfu = train_metrics.get("efficiency/mfu", 0) tokens_per_sec = train_metrics.get("efficiency/tokens_per_second(K)", 0) * 1e3 + gpu_mem_gb = train_metrics.get("memory/gpu_allocated(GB)", 0) if use_tqdm and tqdm_bar is not None: tqdm_bar.set_postfix_str(f"loss={total_loss:.2f} gn={grad_norm:.2f} lr={lr:.1e} tok/s={tokens_per_sec:.0f}") @@ -1051,7 +1050,8 @@ def _maybe_log( f"[STEP {self.state.global_step}/{max_steps_str}] " f"loss={total_loss:.4f} grad_norm={grad_norm:.4f} lr={lr:.6e} " f"tflops={tflops_per_gpu:.1f} mfu={mfu:.4f} " - f"tokens_per_sec={tokens_per_sec:.0f} time={delta_time:.3f}s" + f"tokens_per_sec={tokens_per_sec:.0f} time={delta_time:.3f}s " + f"peak_mem={gpu_mem_gb:.1f}GB" ) if ( @@ -1061,12 +1061,16 @@ def _maybe_log( ): import wandb # noqa: PLC0415 + epoch_fraction = self.state.epoch + ( + (self.state.global_step - self.state.epoch * self.train_steps_per_epoch) + / max(self.train_steps_per_epoch, 1) + ) train_metrics.update( { "training/loss": total_loss, "training/grad_norm": grad_norm, "training/lr": lr, - "training/epoch": self.state.epoch, + "training/epoch": epoch_fraction, "training/step_time": delta_time, "training/samples_seen": self.state.global_step * args.train.global_batch_size, } diff --git a/src/xorl/utils/count_flops.py b/src/xorl/utils/count_flops.py index cd946afa..a2b39987 100644 --- a/src/xorl/utils/count_flops.py +++ b/src/xorl/utils/count_flops.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional +from typing import Dict, Optional from transformers import PretrainedConfig @@ -20,10 +20,9 @@ def _attention_score_elements(batch_seqlens: List[int], causal: bool) -> int: return sum(seqlen * seqlen for seqlen in batch_seqlens) -def _gc_multipliers( - gc_enabled: bool, - recompute_modules: Optional[List[str]], - moe_checkpoint_method: Optional[str], +def _grad_ckpt_multipliers( + gradient_checkpointing_enabled: bool, + gradient_checkpointing_method: Optional[str] = None, ) -> Dict[str, int]: """Logical training FLOPs multipliers aligned with common benchmark reporting. @@ -44,7 +43,7 @@ def _gc_multipliers( down - MoE down_proj multiplier dense_mlp - Dense MLP (non-MoE layers) multiplier """ - del gc_enabled, recompute_modules, moe_checkpoint_method + del gradient_checkpointing_enabled, gradient_checkpointing_method return dict(attn_linear=6, attn_qkv=12, router=6, gate=6, up=6, down=6, dense_mlp=6) @@ -92,12 +91,11 @@ class XorlFlopsCounter: def __init__( self, config: PretrainedConfig, - gc_enabled: bool = False, - recompute_modules: Optional[List[str]] = None, - moe_checkpoint_method: Optional[str] = None, + gradient_checkpointing_enabled: bool = False, + gradient_checkpointing_method: Optional[str] = None, cp_size: int = 1, ): - self._m = _gc_multipliers(gc_enabled, recompute_modules, moe_checkpoint_method) + self._m = _grad_ckpt_multipliers(gradient_checkpointing_enabled, gradient_checkpointing_method) # CP correction: each rank computes on local_tokens = total/cp_size. # The formula uses local tokens → per-rank FLOPs. Multiply by cp_size # to report total-batch FLOPs so that EnvironMeter's @@ -200,7 +198,7 @@ def _estimate_qwen3_moe_flops(self, tokens_sum, batch_seqlens, delta_time): attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim) embed_lm_N = vocab_size * hidden_size # lm_head only; embedding is a lookup (0 FLOPs) - # Per-layer FLOPs with GC-corrected multipliers + # Per-layer FLOPs with gradient-checkpoint-corrected multipliers per_layer_flops = ( m["router"] * router_N * tokens_sum + m["gate"] * gate_up_N * tokens_sum diff --git a/src/xorl/utils/helper.py b/src/xorl/utils/helper.py index 32bf93c5..af098e95 100644 --- a/src/xorl/utils/helper.py +++ b/src/xorl/utils/helper.py @@ -93,9 +93,8 @@ def __init__( global_batch_size: int, empty_cache_steps: int = 500, gc_steps: int = 0, - gc_enabled: bool = False, - recompute_modules=None, - moe_checkpoint_method=None, + gradient_checkpointing_enabled: bool = False, + gradient_checkpointing_method=None, cp_size: int = 1, ) -> None: self.config = config @@ -109,9 +108,8 @@ def __init__( self.estimate_flops = XorlFlopsCounter( config, - gc_enabled=gc_enabled, - recompute_modules=recompute_modules, - moe_checkpoint_method=moe_checkpoint_method, + gradient_checkpointing_enabled=gradient_checkpointing_enabled, + gradient_checkpointing_method=gradient_checkpointing_method, cp_size=cp_size, ).estimate_flops From 9a5f5a065c122768e312e0fc6dd77683d69aa7f4 Mon Sep 17 00:00:00 2001 From: AMAN SINGHAL <65446479+AmanSinghal927@users.noreply.github.com> Date: Wed, 15 Apr 2026 16:58:44 -0400 Subject: [PATCH 30/41] Fix NCCL collective mismatch in packing bin computation causing distributed training hang (#136) * barrier bug fix * Update auto.py * Remove unused torch.distributed import from packing.py * trigger CI (cherry picked from commit 2e75acf14024f713b618114f6fe90b18de88ce1d) --- src/xorl/data/prepare/packing.py | 38 +++++++++++++------------------- 1 file changed, 15 insertions(+), 23 deletions(-) diff --git a/src/xorl/data/prepare/packing.py b/src/xorl/data/prepare/packing.py index 8d1f1b6e..ea894e9b 100644 --- a/src/xorl/data/prepare/packing.py +++ b/src/xorl/data/prepare/packing.py @@ -13,7 +13,6 @@ import numba import numpy as np -import torch.distributed as dist from datasets import Dataset as HFDataset from torch.utils.data import Dataset from transformers import PreTrainedTokenizer @@ -458,10 +457,16 @@ def _save_bins_cache(self, bins: List[List[int]]) -> None: LOG.info(f"Saved bins cache to {cache_path}") def _load_or_compute_bins(self) -> List[List[int]]: - """Load cached bins or compute new ones with distributed coordination.""" - # Get rank from environment variable (set by torchrun/distributed launcher) + """Load cached bins or compute new ones. + + Distributed synchronization (barrier) is handled by the caller + (``prepare_datasets``), NOT here. ``prepare_datasets`` calls + ``_load_datasets`` on rank 0 first, then barriers, then on + non-zero ranks. Adding a barrier inside this function would + cause an NCCL collective mismatch because different ranks call + this function at different times. + """ rank = int(os.environ.get("RANK", "0")) - is_distributed = dist.is_available() and dist.is_initialized() # Try to load from cache first cached_bins = self._load_cached_bins() @@ -470,33 +475,20 @@ def _load_or_compute_bins(self) -> List[List[int]]: cache_path = self._get_bins_cache_path() - # Use rank 0 to compute and save, others wait and load if rank == 0: - # Rank 0 computes new bins + # Rank 0 computes new bins and saves to cache. + # Non-zero ranks will load from cache after the caller's + # barrier ensures the cache file is visible. LOG.info("Computing new bins...") bins = self._compute_bins() - # Save to cache for future use if cache_path is not None: self._save_bins_cache(bins) - if is_distributed: - dist.barrier() else: - # Other ranks wait for rank 0 to finish + # Non-zero ranks: try loading from cache (written by rank 0). + # The caller's barrier guarantees rank 0 has finished writing + # before non-zero ranks reach this point. if cache_path is not None: - LOG.info(f"Rank {rank} waiting for rank 0 to compute bins...") - if is_distributed: - dist.barrier() - else: - max_wait_time = 3600 - wait_interval = 5 - waited = 0 - while not cache_path.exists() and waited < max_wait_time: - time.sleep(wait_interval) - waited += wait_interval - if not cache_path.exists(): - raise RuntimeError(f"Rank {rank} timed out waiting for rank 0 to compute bins") - bins = None max_retries = 10 for attempt in range(max_retries): From 2b033287145ecebbfca95edb81a8f78aa2e65f6b Mon Sep 17 00:00:00 2001 From: Zhongzhu Zhou Date: Fri, 17 Apr 2026 01:59:03 +0800 Subject: [PATCH 31/41] [Perf] Fused SiLU+mul kernels in Triton MoE EP group GEMM (#113) * - Replace manual gate/up split + SiLU + mul with fused silu_and_mul kernel in TritonEPGroupGemm forward/backward (fewer intermediates, one kernel) - Fix quack/native EP adapters to forward expert_scores to downstream (was silently dropped, causing incorrect routing score application) - Fix quack/native moe_act adapters to accept expert_scores param (was TypeError on 6-arg call from Experts._ep_forward) - Add adapter-level regression tests for expert_scores forwarding * Fix lint: remove moe_act refs from test, fix ruff formatting moe_act was removed in PR #85. Update test to only check EP_EXPERT_COMPUTE adapters (quack/native). Fix assert formatting for ruff. Made-with: Cursor * Fix ruff lint: remove unused var, fix import order, format - Remove unused `I = ctx.intermediate_size` in TritonEPGroupGemm.backward - Move import_utils import to top-level in test file - Fix import sorting and trailing blank lines Made-with: Cursor (cherry picked from commit 7989ab38a6fd7a16c1a88918ab717650ecec6f07) --- src/xorl/ops/moe/triton.py | 28 +--- tests/ops/test_ep_adapter_wrappers.py | 181 ++++++++++++++++++++++++++ 2 files changed, 188 insertions(+), 21 deletions(-) create mode 100644 tests/ops/test_ep_adapter_wrappers.py diff --git a/src/xorl/ops/moe/triton.py b/src/xorl/ops/moe/triton.py index f06c62c1..d4324bd2 100644 --- a/src/xorl/ops/moe/triton.py +++ b/src/xorl/ops/moe/triton.py @@ -1,5 +1,6 @@ import torch +from xorl.ops.fused_silu_and_mul import silu_and_mul, silu_and_mul_backward from xorl.utils.import_utils import is_fused_moe_available @@ -36,14 +37,11 @@ def forward(ctx, permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_s transpose_a=False, transpose_b=False, ) - gate_output = gate_up_output[..., :I] - up_output = gate_up_output[..., I:] - gate_activation = torch.ops.aten.silu(gate_output) - gated_output = gate_activation * up_output + # Fused SiLU+mul: silu(gate) * up in a single Triton kernel + gated_output = silu_and_mul(gate_up_output) if expert_scores is not None: gated_output = gated_output * expert_scores.to(gated_output.dtype).unsqueeze(-1) - del gate_activation down_output = group_gemm_same_nk( a=gated_output, @@ -65,14 +63,9 @@ def forward(ctx, permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_s @staticmethod def backward(ctx, grad_output): permute_tokens, cumsum, gate_up_proj, down_proj, gate_up_output, expert_scores = ctx.saved_tensors - I = ctx.intermediate_size max_M = grad_output.shape[0] - gate_output = gate_up_output[..., :I] - up_output = gate_up_output[..., I:] - - gate_activation = torch.ops.aten.silu(gate_output) - gated_output = gate_activation * up_output + gated_output = silu_and_mul(gate_up_output) expert_scores_dtype = expert_scores.dtype expert_scores = expert_scores.to(gated_output.dtype) gated_weighted = gated_output * expert_scores.unsqueeze(-1) @@ -107,17 +100,10 @@ def backward(ctx, grad_output): grad_gated_output = grad_gated_weighted * expert_scores.unsqueeze(-1) del grad_gated_weighted - # Activation backward - grad_up_output = gate_activation * grad_gated_output - grad_gate_activation = grad_gated_output * up_output - del grad_gated_output, gate_activation, up_output, gate_up_output - - grad_gate_output = torch.ops.aten.silu_backward(grad_gate_activation, gate_output) - del grad_gate_activation, gate_output + # Fused activation backward: single Triton kernel for SiLU+mul gradient + grad_gate_up_act = silu_and_mul_backward(grad_gated_output, gate_up_output) + del grad_gated_output, gate_up_output - # Fused dgrad FC1: cat grads → single GEMM with gate_up_proj (no .contiguous() copies) - grad_gate_up_act = torch.cat([grad_gate_output, grad_up_output], dim=-1) - del grad_gate_output, grad_up_output grad_permute_tokens = group_gemm_same_nk( a=grad_gate_up_act, b=gate_up_proj, diff --git a/tests/ops/test_ep_adapter_wrappers.py b/tests/ops/test_ep_adapter_wrappers.py new file mode 100644 index 00000000..94e90053 --- /dev/null +++ b/tests/ops/test_ep_adapter_wrappers.py @@ -0,0 +1,181 @@ +"""Tests that EP adapter wrappers in backend/__init__.py correctly forward expert_scores. + +Bug: _quack_ep_fused / _native_ep_fused accepted expert_scores but silently dropped it. + +These tests monkeypatch the downstream implementations to verify the adapters pass +expert_scores through, and compare output to a naive reference. +""" + +import importlib.util +import sys +import types +from pathlib import Path + +import pytest +import torch +import torch.nn.functional as F + +from xorl.utils import import_utils + + +pytestmark = pytest.mark.cpu + +_MODULE_PATHS = { + "xorl.ops.moe.triton": Path(__file__).resolve().parents[2] / "src/xorl/ops/moe/triton.py", + "xorl.ops.moe.quack": Path(__file__).resolve().parents[2] / "src/xorl/ops/moe/quack.py", +} + +_BACKEND_INIT_PATH = Path(__file__).resolve().parents[2] / "src/xorl/models/layers/moe/backend/__init__.py" + + +def _counts_from_cumsum(cumsum: torch.Tensor) -> list[int]: + counts = torch.empty_like(cumsum) + counts[0] = cumsum[0] + counts[1:] = cumsum[1:] - cumsum[:-1] + return counts.tolist() + + +def _naive_group_gemm_same_nk(a, b, cumsum_M, max_M, transpose_a=False, transpose_b=False, **kwargs): + del max_M, kwargs + assert not transpose_a + outputs = [] + start = 0 + for expert_idx, count in enumerate(_counts_from_cumsum(cumsum_M)): + end = start + count + weight = b[expert_idx] + if transpose_b: + outputs.append(a[start:end] @ weight.transpose(0, 1)) + else: + outputs.append(a[start:end] @ weight) + start = end + return torch.cat(outputs, dim=0) + + +def _naive_group_gemm_same_mn(a, b, c, cumsum_K, max_K, transpose_a=False, transpose_b=False, **kwargs): + del max_K, kwargs + start = 0 + for expert_idx, count in enumerate(_counts_from_cumsum(cumsum_K)): + end = start + count + lhs = a[start:end].transpose(0, 1) if transpose_a else a[start:end] + rhs = b[start:end].transpose(0, 1) if transpose_b else b[start:end] + c[expert_idx].copy_(lhs @ rhs) + start = end + return c + + +def reference_ep_forward(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores): + """Naive reference: per-expert matmul with SiLU + optional score scaling.""" + outputs = [] + start = 0 + for expert_idx, count in enumerate(_counts_from_cumsum(cumsum)): + end = start + count + x = permute_tokens[start:end] + h = F.silu(x @ gate_proj[expert_idx]) * (x @ up_proj[expert_idx]) + if expert_scores is not None: + h = h * expert_scores[start:end].to(h.dtype).unsqueeze(-1) + outputs.append(h @ down_proj[expert_idx]) + start = end + return torch.cat(outputs, dim=0) + + +def _patch_kernels_and_load_backend(monkeypatch, backend_type: str): + """Patch kernel modules and load backend/__init__.py to get the adapter wrappers.""" + moe_stub = types.ModuleType("xorl.ops.group_gemm.kernel.moe") + moe_stub.expert_histogram = None + moe_stub.moe_gather = None + moe_stub.moe_index_compute = None + moe_stub.moe_scatter = None + monkeypatch.setattr(import_utils, "is_fused_moe_available", lambda: True) + sys.modules.pop("xorl.ops.group_gemm.kernel.moe", None) + sys.modules.pop("xorl.ops.group_gemm.kernel.group_gemm", None) + sys.modules.pop("xorl.ops.group_gemm.kernel.quack", None) + monkeypatch.setitem(sys.modules, "xorl.ops.group_gemm.kernel.moe", moe_stub) + + if backend_type == "triton": + group_gemm_stub = types.ModuleType("xorl.ops.group_gemm.kernel.group_gemm") + group_gemm_stub.group_gemm_same_nk = _naive_group_gemm_same_nk + group_gemm_stub.group_gemm_same_mn = _naive_group_gemm_same_mn + monkeypatch.setitem(sys.modules, "xorl.ops.group_gemm.kernel.group_gemm", group_gemm_stub) + module_name = "xorl.ops.moe.triton" + else: + quack_stub = types.ModuleType("xorl.ops.group_gemm.kernel.quack") + quack_stub.cumsum_to_cu_seqlens = lambda cumsum: cumsum + quack_stub.quack_group_gemm_same_nk = _naive_group_gemm_same_nk + quack_stub.quack_group_gemm_same_mn = _naive_group_gemm_same_mn + monkeypatch.setitem(sys.modules, "xorl.ops.group_gemm.kernel.quack", quack_stub) + module_name = "xorl.ops.moe.quack" + + spec = importlib.util.spec_from_file_location(f"adapter_test_{backend_type}", _MODULE_PATHS[module_name]) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + monkeypatch.setitem(sys.modules, module_name, module) + + return module + + +def _make_test_data(dtype=torch.float32): + torch.manual_seed(42) + num_local_experts = 2 + hidden_dim = 8 + intermediate_size = 12 + counts = torch.tensor([3, 2]) + cumsum = torch.cumsum(counts, dim=0) + num_tokens = int(cumsum[-1].item()) + + permute_tokens = torch.randn(num_tokens, hidden_dim, dtype=dtype) + gate_proj = torch.randn(num_local_experts, hidden_dim, intermediate_size, dtype=dtype) + up_proj = torch.randn(num_local_experts, hidden_dim, intermediate_size, dtype=dtype) + down_proj = torch.randn(num_local_experts, intermediate_size, hidden_dim, dtype=dtype) + gate_up_proj = torch.cat([gate_proj, up_proj], dim=-1) + expert_scores = torch.rand(num_tokens, dtype=dtype) + + return permute_tokens, cumsum, gate_proj, up_proj, down_proj, gate_up_proj, intermediate_size, expert_scores + + +@pytest.mark.parametrize( + ("backend_type", "class_name"), + [ + pytest.param("triton", "TritonEPGroupGemm", id="triton-ep"), + pytest.param("quack", "QuackEPGroupGemm", id="quack-ep"), + ], +) +def test_adapter_forwards_expert_scores(monkeypatch, backend_type, class_name): + """Verify that EP adapter wrappers accept and forward expert_scores.""" + try: + kernel_module = _patch_kernels_and_load_backend(monkeypatch, backend_type) + except ImportError as exc: + pytest.skip(f"{backend_type} unavailable: {exc}") + + ( + permute_tokens, + cumsum, + gate_proj, + up_proj, + down_proj, + gate_up_proj, + intermediate_size, + expert_scores, + ) = _make_test_data() + + fn_cls = getattr(kernel_module, class_name) + output = fn_cls.apply(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores) + + ref = reference_ep_forward(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores) + torch.testing.assert_close(output, ref) + + +def test_adapter_source_forwards_expert_scores(): + """Regression test: verify backend/__init__.py adapter wrappers pass expert_scores. + + Parses the source to confirm the apply() / compute() calls include expert_scores. + This catches silent-drop bugs without needing to load the full module. + """ + source = _BACKEND_INIT_PATH.read_text() + + assert "_QuackEPGroupGemm.apply(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores)" in source, ( + "_quack_ep_fused does not forward expert_scores to _QuackEPGroupGemm.apply()" + ) + assert "_native_ep_compute(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores)" in source, ( + "_native_ep_fused does not forward expert_scores to _native_ep_compute()" + ) From 5d8ec39fc27450359f89e14e302e81ae5e488d44 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Thu, 16 Apr 2026 11:36:04 -0700 Subject: [PATCH 32/41] Add Llama 3 model support (#125) Implement Llama (covers Llama 2, 3, 3.1, 3.2, 3.3) as a new model family following the established Qwen3 pattern. Key differences from Qwen3: - No QK normalization (replaced with nn.Identity) - No sliding window attention - GQA with 8 KV heads (vs Qwen3's 32 MHA heads) - mlp_bias config support - Llama 3 defaults (128K vocab, rope_theta=500K, rms_norm_eps=1e-5) Verified on 8x H100: - Correctness: cosine sim 0.998 vs HF, 100% token agreement - FSDP2: 69K tok/s, 30.42 GB peak VRAM - TP4: 10.91->7.43 loss, 23 GB peak VRAM - PP2: 10.91->8.90 loss, 37.67 GB peak VRAM - LoRA (rank 16): 10.91->9.98 loss, 21.42 GB peak VRAM - QLoRA NVFP4: 13.82->10.67 loss, 21.56 GB peak VRAM (cherry picked from commit 6f3c5814cd81b7be9bd67490439a89342561ad45) --- .../local/dummy/configs/full/llama3_8b.yaml | 60 +++ .../dummy/configs/full/llama3_8b_pp2.yaml | 53 +++ .../dummy/configs/full/llama3_8b_tp4.yaml | 54 +++ .../dummy/configs/lora/llama3_8b_lora.yaml | 54 +++ .../configs/qlora/llama3_8b_qlora_nvfp4.yaml | 56 +++ src/xorl/models/transformers/__init__.py | 2 + .../models/transformers/llama3/__init__.py | 0 .../transformers/llama3/checkpoint_handler.py | 170 +++++++++ .../llama3/configuration_llama3.py | 78 ++++ .../transformers/llama3/modeling_llama3.py | 350 ++++++++++++++++++ .../models/transformers/llama3/parallelize.py | 37 ++ 11 files changed, 914 insertions(+) create mode 100644 examples/local/dummy/configs/full/llama3_8b.yaml create mode 100644 examples/local/dummy/configs/full/llama3_8b_pp2.yaml create mode 100644 examples/local/dummy/configs/full/llama3_8b_tp4.yaml create mode 100644 examples/local/dummy/configs/lora/llama3_8b_lora.yaml create mode 100644 examples/local/dummy/configs/qlora/llama3_8b_qlora_nvfp4.yaml create mode 100644 src/xorl/models/transformers/llama3/__init__.py create mode 100644 src/xorl/models/transformers/llama3/checkpoint_handler.py create mode 100644 src/xorl/models/transformers/llama3/configuration_llama3.py create mode 100644 src/xorl/models/transformers/llama3/modeling_llama3.py create mode 100644 src/xorl/models/transformers/llama3/parallelize.py diff --git a/examples/local/dummy/configs/full/llama3_8b.yaml b/examples/local/dummy/configs/full/llama3_8b.yaml new file mode 100644 index 00000000..76c7a5b6 --- /dev/null +++ b/examples/local/dummy/configs/full/llama3_8b.yaml @@ -0,0 +1,60 @@ +model: + model_path: meta-llama/Llama-3.1-8B + attn_implementation: flash_attention_3 + +data: + datasets: + - path: dummy + type: tokenized + max_seq_len: 8000 + select_columns: [input_ids, labels] + dataset_prepared_path: last_prepared_dataset + sample_packing_method: sequential + sample_packing_sequence_len: 32000 + dataloader_num_workers: 4 + dataloader_prefetch_factor: 2 + dataloader_pin_memory: true + + +train: + output_dir: outputs/Llama-3.1-8B-local-test + data_parallel_mode: fsdp2 + ulysses_parallel_size: 1 + data_parallel_replicate_size: 1 + data_parallel_shard_size: 8 + expert_parallel_size: 1 + tensor_parallel_size: 1 + + num_train_epochs: 5 + + # batch size + micro_batch_size: 1 + gradient_accumulation_steps: 1 + + optimizer: adamw + # lr and warmup + lr: 1.0e-5 + lr_warmup_ratio: 0.005 + lr_decay_style: cosine + lr_decay_ratio: 1.0 + weight_decay: 0.01 + + max_grad_norm: 1.0 + enable_mixed_precision: true + enable_gradient_checkpointing: true + enable_full_shard: true + enable_activation_offload: false + init_device: meta + load_weights_mode: broadcast + enable_full_determinism: false + empty_cache_steps: 500 + ckpt_manager: dcp + save_steps: 1000 + save_hf_weights: true + + max_steps: 20 + + # wandb + use_wandb: false + wandb_project: local-training + wandb_name: Llama-3.1-8B-local-test diff --git a/examples/local/dummy/configs/full/llama3_8b_pp2.yaml b/examples/local/dummy/configs/full/llama3_8b_pp2.yaml new file mode 100644 index 00000000..b4f6d869 --- /dev/null +++ b/examples/local/dummy/configs/full/llama3_8b_pp2.yaml @@ -0,0 +1,53 @@ +model: + model_path: meta-llama/Llama-3.1-8B + attn_implementation: flash_attention_3 + +data: + datasets: + - path: dummy + type: tokenized + max_seq_len: 8000 + select_columns: [input_ids, labels] + dataset_prepared_path: last_prepared_dataset + sample_packing_method: sequential + sample_packing_sequence_len: 8000 + dataloader_num_workers: 4 + dataloader_prefetch_factor: 2 + dataloader_pin_memory: true + +train: + output_dir: outputs/Llama-3.1-8B-pp2 + data_parallel_mode: fsdp2 + ulysses_parallel_size: 1 + pipeline_parallel_size: 2 + data_parallel_replicate_size: 1 + data_parallel_shard_size: 4 + expert_parallel_size: 1 + + num_train_epochs: 5 + max_steps: 20 + + micro_batch_size: 1 + gradient_accumulation_steps: 4 + + optimizer: adamw + lr: 1.0e-4 + lr_warmup_ratio: 0.01 + lr_decay_style: cosine + lr_decay_ratio: 1.0 + weight_decay: 0.01 + + max_grad_norm: 5.0 + enable_mixed_precision: true + enable_gradient_checkpointing: true + enable_full_shard: true + enable_activation_offload: false + init_device: meta + load_weights_mode: broadcast + enable_full_determinism: false + empty_cache_steps: 500 + ckpt_manager: dcp + save_steps: 0 + save_hf_weights: false + + use_wandb: false diff --git a/examples/local/dummy/configs/full/llama3_8b_tp4.yaml b/examples/local/dummy/configs/full/llama3_8b_tp4.yaml new file mode 100644 index 00000000..e990cf43 --- /dev/null +++ b/examples/local/dummy/configs/full/llama3_8b_tp4.yaml @@ -0,0 +1,54 @@ +model: + model_path: meta-llama/Llama-3.1-8B + attn_implementation: flash_attention_3 + +data: + datasets: + - path: dummy + type: tokenized + max_seq_len: 8000 + select_columns: [input_ids, labels] + dataset_prepared_path: last_prepared_dataset + sample_packing_method: sequential + sample_packing_sequence_len: 32000 + dataloader_num_workers: 4 + dataloader_prefetch_factor: 2 + dataloader_pin_memory: true + +train: + output_dir: outputs/Llama-3.1-8B-tp4 + data_parallel_mode: fsdp2 + tensor_parallel_size: 4 + ulysses_parallel_size: 1 + ringattn_parallel_size: 1 + data_parallel_replicate_size: 1 + data_parallel_shard_size: 2 + expert_parallel_size: 1 + + num_train_epochs: 5 + max_steps: 20 + + micro_batch_size: 1 + gradient_accumulation_steps: 1 + + optimizer: adamw + lr: 1.0e-5 + lr_warmup_ratio: 0.005 + lr_decay_style: cosine + lr_decay_ratio: 1.0 + weight_decay: 0.01 + + max_grad_norm: 1.0 + enable_mixed_precision: true + enable_gradient_checkpointing: true + enable_full_shard: true + enable_activation_offload: false + init_device: meta + load_weights_mode: broadcast + enable_full_determinism: false + empty_cache_steps: 500 + ckpt_manager: dcp + save_steps: 0 + save_hf_weights: false + + use_wandb: false diff --git a/examples/local/dummy/configs/lora/llama3_8b_lora.yaml b/examples/local/dummy/configs/lora/llama3_8b_lora.yaml new file mode 100644 index 00000000..cf82f128 --- /dev/null +++ b/examples/local/dummy/configs/lora/llama3_8b_lora.yaml @@ -0,0 +1,54 @@ +model: + model_path: meta-llama/Llama-3.1-8B + attn_implementation: flash_attention_3 + +data: + datasets: + - path: dummy + type: tokenized + max_seq_len: 8000 + select_columns: [input_ids, labels] + dataset_prepared_path: last_prepared_dataset + sample_packing_method: sequential + sample_packing_sequence_len: 32000 + dataloader_num_workers: 4 + dataloader_prefetch_factor: 2 + dataloader_pin_memory: true + +train: + output_dir: outputs/llama3-8b-lora + data_parallel_mode: fsdp2 + ulysses_parallel_size: 1 + data_parallel_replicate_size: 1 + data_parallel_shard_size: 8 + + num_train_epochs: 1 + max_steps: 20 + + micro_batch_size: 1 + gradient_accumulation_steps: 1 + + optimizer: adamw + lr: 1.0e-4 + lr_warmup_ratio: 0.1 + lr_decay_style: cosine + weight_decay: 0.01 + + max_grad_norm: 1.0 + enable_mixed_precision: true + enable_gradient_checkpointing: true + enable_full_shard: true + init_device: meta + load_weights_mode: broadcast + ckpt_manager: dcp + save_steps: 0 + save_hf_weights: false + + use_wandb: false + +lora: + enable_lora: true + lora_rank: 16 + lora_alpha: 16 + lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] + save_lora_only: true diff --git a/examples/local/dummy/configs/qlora/llama3_8b_qlora_nvfp4.yaml b/examples/local/dummy/configs/qlora/llama3_8b_qlora_nvfp4.yaml new file mode 100644 index 00000000..c7891aa8 --- /dev/null +++ b/examples/local/dummy/configs/qlora/llama3_8b_qlora_nvfp4.yaml @@ -0,0 +1,56 @@ +model: + model_path: meta-llama/Llama-3.1-8B + attn_implementation: flash_attention_3 + +data: + datasets: + - path: dummy + type: tokenized + max_seq_len: 8000 + select_columns: [input_ids, labels] + dataset_prepared_path: last_prepared_dataset + sample_packing_method: sequential + sample_packing_sequence_len: 32000 + dataloader_num_workers: 4 + dataloader_prefetch_factor: 2 + dataloader_pin_memory: true + +train: + output_dir: outputs/llama3-8b-qlora-nvfp4 + data_parallel_mode: fsdp2 + ulysses_parallel_size: 1 + data_parallel_replicate_size: 1 + data_parallel_shard_size: 8 + tensor_parallel_size: 1 + + num_train_epochs: 1 + max_steps: 20 + + micro_batch_size: 1 + gradient_accumulation_steps: 1 + + optimizer: adamw + lr: 1.0e-4 + lr_warmup_ratio: 0.1 + lr_decay_style: cosine + weight_decay: 0.01 + + max_grad_norm: 1.0 + enable_mixed_precision: true + enable_gradient_checkpointing: true + enable_full_shard: true + init_device: meta + load_weights_mode: broadcast + ckpt_manager: dcp + save_steps: 0 + save_hf_weights: false + + use_wandb: false + +lora: + enable_qlora: true + lora_rank: 16 + lora_alpha: 16 + lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] + quant_format: nvfp4 + quant_group_size: 16 diff --git a/src/xorl/models/transformers/__init__.py b/src/xorl/models/transformers/__init__.py index 830d6933..84d4d83f 100644 --- a/src/xorl/models/transformers/__init__.py +++ b/src/xorl/models/transformers/__init__.py @@ -1,4 +1,5 @@ from . import ( + llama3, qwen3, qwen3_5, qwen3_5_moe, @@ -7,6 +8,7 @@ __all__ = [ + "llama3", "qwen3", "qwen3_5", "qwen3_moe", diff --git a/src/xorl/models/transformers/llama3/__init__.py b/src/xorl/models/transformers/llama3/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/xorl/models/transformers/llama3/checkpoint_handler.py b/src/xorl/models/transformers/llama3/checkpoint_handler.py new file mode 100644 index 00000000..e4c2751a --- /dev/null +++ b/src/xorl/models/transformers/llama3/checkpoint_handler.py @@ -0,0 +1,170 @@ +"""Checkpoint handler for Llama 3 models.""" + +import warnings +from typing import Callable, List, Optional, Set, Tuple + +import torch +import torch.nn as nn + +from ...checkpoint_handlers.base import CheckpointHandler +from ...checkpoint_handlers.buffers import ( + DENSE_DOWN_PROJ_PATTERN, + DENSE_GATE_UP_PATTERN, + FP8_AUX_SUFFIX_PATTERN, + OPROJ_WEIGHT_PATTERN, + QKV_PROJ_PATTERN, + QUANT_AUX_SUFFIX_PATTERN, + GateUpMergeBuffer, + QKVMergeBuffer, + QLoRAWeightBuffer, +) + + +class Llama3CheckpointHandler(CheckpointHandler): + """Checkpoint handler for Llama 3 models. + + Load: merge gate_proj + up_proj -> gate_up_proj + merge q_proj + k_proj + v_proj -> qkv_proj + Save: split gate_up_proj -> gate_proj + up_proj + split qkv_proj -> q_proj + k_proj + v_proj + + When ``is_prequantized=True`` and a model is provided, quantized weights are + loaded inline via QLoRAWeightBuffer (single-pass I/O). When model is not + provided, falls back to skipping quantized keys (deferred loading). + """ + + def __init__( + self, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + is_prequantized: bool = False, + exclude_modules: Optional[Set[str]] = None, + model: Optional[nn.Module] = None, + ): + self._gate_up_buffer = GateUpMergeBuffer() + self._qkv_buffer = QKVMergeBuffer() + self._q_dim = num_attention_heads * head_dim + self._kv_dim = num_key_value_heads * head_dim + self._is_prequantized = is_prequantized + self._exclude_modules = exclude_modules or set() + self._qlora_buffer: Optional[QLoRAWeightBuffer] = None + if is_prequantized and model is not None: + self._qlora_buffer = QLoRAWeightBuffer(model) + + def get_skip_key_fn(self) -> Optional[Callable[[str], bool]]: + if not self._is_prequantized: + return None + + if self._qlora_buffer is not None: + return None + + exclude_modules = self._exclude_modules + + def _should_skip(key: str) -> bool: + if exclude_modules: + module_fqn = key.rsplit(".", 1)[0] if "." in key else key + module_short_name = module_fqn.rsplit(".", 1)[-1] + if module_short_name in exclude_modules: + return False + + if QUANT_AUX_SUFFIX_PATTERN.search(key): + return True + if FP8_AUX_SUFFIX_PATTERN.search(key): + return True + if key.endswith(".weight"): + if ( + QKV_PROJ_PATTERN.match(key) + or DENSE_GATE_UP_PATTERN.match(key) + or OPROJ_WEIGHT_PATTERN.match(key) + or DENSE_DOWN_PROJ_PATTERN.match(key) + ): + return True + return False + + return _should_skip + + def _is_excluded_module(self, key: str) -> bool: + if not self._exclude_modules: + return False + module_fqn = key.rsplit(".", 1)[0] if "." in key else key + module_short_name = module_fqn.rsplit(".", 1)[-1] + return module_short_name in self._exclude_modules + + def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: + # Drop input_scale (unused by our quantization) + if key.endswith(".input_scale"): + return [] + + # QLoRA buffer: route quantized keys for inline loading + if self._qlora_buffer is not None and not self._is_excluded_module(key): + result = self._qlora_buffer.try_consume(key, tensor) + if result is not None: + return result + + # Pre-quantized without buffer (deferred path): drop quantized keys + if self._is_prequantized and self._qlora_buffer is None and not self._is_excluded_module(key): + if QUANT_AUX_SUFFIX_PATTERN.search(key): + return [] + if FP8_AUX_SUFFIX_PATTERN.search(key): + return [] + if key.endswith(".weight"): + if OPROJ_WEIGHT_PATTERN.match(key) or DENSE_DOWN_PROJ_PATTERN.match(key): + return [] + + # QKV merge + if self._is_prequantized and key.endswith(".weight"): + if self._qkv_buffer.is_qkv_key(key): + return [] + else: + qkv_result = self._qkv_buffer.add(key, tensor) + if qkv_result is not None: + return [qkv_result] + if self._qkv_buffer.is_qkv_key(key): + return [] + + # Gate/up merge + if self._is_prequantized and key.endswith(".weight"): + if self._gate_up_buffer.is_gate_up_key(key): + return [] + else: + merge_result = self._gate_up_buffer.add(key, tensor) + if merge_result is not None: + return [merge_result] + if self._gate_up_buffer.is_gate_up_key(key): + return [] + + return [(key, tensor)] + + def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: + pending_gu = self._gate_up_buffer.get_pending() + if pending_gu: + warnings.warn(f"Incomplete gate/up merge pairs after loading: {pending_gu}") + pending_qkv = self._qkv_buffer.get_pending() + if pending_qkv: + warnings.warn(f"Incomplete QKV merge groups after loading: {pending_qkv}") + if self._qlora_buffer is not None: + self._qlora_buffer.set_inline_metadata() + return [] + + def on_save_weight(self, param_name: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: + # Split gate_up_proj -> gate_proj + up_proj + if ".gate_up_proj." in param_name: + prefix, suffix = param_name.rsplit(".gate_up_proj.", 1) + half = tensor.shape[0] // 2 + return [ + (f"{prefix}.gate_proj.{suffix}", tensor[:half]), + (f"{prefix}.up_proj.{suffix}", tensor[half:]), + ] + + # Split qkv_proj -> q_proj + k_proj + v_proj + if ".qkv_proj." in param_name: + prefix, suffix = param_name.rsplit(".qkv_proj.", 1) + q, k, v = tensor.split([self._q_dim, self._kv_dim, self._kv_dim], dim=0) + return [ + (f"{prefix}.q_proj.{suffix}", q), + (f"{prefix}.k_proj.{suffix}", k), + (f"{prefix}.v_proj.{suffix}", v), + ] + + return [(param_name, tensor)] diff --git a/src/xorl/models/transformers/llama3/configuration_llama3.py b/src/xorl/models/transformers/llama3/configuration_llama3.py new file mode 100644 index 00000000..520f925b --- /dev/null +++ b/src/xorl/models/transformers/llama3/configuration_llama3.py @@ -0,0 +1,78 @@ +"""Llama 3 model configuration.""" + +from transformers.configuration_utils import PretrainedConfig + +from ....utils import logging +from .parallelize import TP_PLAN + + +logger = logging.get_logger(__name__) + + +class Llama3Config(PretrainedConfig): + r""" + Configuration class for the Llama 3 dense model. + + Covers Llama 3, 3.1, 3.2, and 3.3 model variants. + """ + + model_type = "xorl_llama3" + + base_model_tp_plan = TP_PLAN + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + def __init__( + self, + vocab_size=128256, + hidden_size=4096, + intermediate_size=14336, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=8, + head_dim=128, + hidden_act="silu", + max_position_embeddings=131072, + initializer_range=0.02, + rms_norm_eps=1e-5, + use_cache=True, + tie_word_embeddings=False, + rope_theta=500000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + mlp_bias=False, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.mlp_bias = mlp_bias + + # BC: if there is a 'type' field, move it to 'rope_type'. + if self.rope_scaling is not None and "type" in self.rope_scaling: + self.rope_scaling["rope_type"] = self.rope_scaling["type"] + + super().__init__( + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +__all__ = ["Llama3Config"] diff --git a/src/xorl/models/transformers/llama3/modeling_llama3.py b/src/xorl/models/transformers/llama3/modeling_llama3.py new file mode 100644 index 00000000..ee1a3b19 --- /dev/null +++ b/src/xorl/models/transformers/llama3/modeling_llama3.py @@ -0,0 +1,350 @@ +from typing import Optional, Tuple, Unpack + +import torch +from torch import nn + +from xorl.distributed.parallel_state import get_parallel_state +from xorl.distributed.sequence_parallel.strategy import get_cp_strategy +from xorl.models.base import XorlPreTrainedModel +from xorl.models.checkpoint_handlers.buffers import ( + detect_prequantized_block_fp8_checkpoint, + detect_prequantized_checkpoint, + get_prequantized_exclude_modules, +) +from xorl.models.layers import ACT2FN, RMSNorm, RotaryEmbedding +from xorl.models.layers.attention import ( + AttentionKwargs, + MultiHeadAttention, + is_flash_attention, + update_causal_mask, +) +from xorl.models.module_utils import GradientCheckpointingLayer +from xorl.models.outputs import BaseModelOutput, CausalLMOutput +from xorl.models.transformers.llama3 import parallelize +from xorl.models.transformers.llama3.checkpoint_handler import Llama3CheckpointHandler +from xorl.models.transformers.llama3.configuration_llama3 import Llama3Config +from xorl.ops.fused_silu_and_mul import fused_silu_and_mul +from xorl.utils import logging + + +logger = logging.get_logger(__name__) + + +class LlamaMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + bias = getattr(config, "mlp_bias", False) + self.gate_up_proj = nn.Linear(self.hidden_size, 2 * self.intermediate_size, bias=bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias) + self.act_fn = ACT2FN[config.hidden_act] + self._use_fused_silu = config.hidden_act == "silu" and not getattr(config, "_activation_native", False) + self._bias = bias + + def unfuse_for_tp(self): + """Replace fused gate_up_proj with separate gate_proj and up_proj for tensor parallelism.""" + device = self.gate_up_proj.weight.device + dtype = self.gate_up_proj.weight.dtype + self.gate_proj = nn.Linear( + self.hidden_size, self.intermediate_size, bias=self._bias, device=device, dtype=dtype + ) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=self._bias, device=device, dtype=dtype) + del self.gate_up_proj + + def forward(self, x): + if hasattr(self, "gate_up_proj"): + if self._use_fused_silu: + x = fused_silu_and_mul(self.gate_up_proj(x)) + else: + gate, up = self.gate_up_proj(x).chunk(2, dim=-1) + x = self.act_fn(gate) * up + else: + x = self.act_fn(self.gate_proj(x)) * self.up_proj(x) + return self.down_proj(x) + + +class LlamaAttention(MultiHeadAttention): + """Llama attention -- standard GQA without QK normalization.""" + + def __init__(self, config, layer_idx: int): + super().__init__(config, layer_idx) + # Llama does not use QK normalization; replace with identity so the + # base class _project_qkv works without override. + self.q_norm = nn.Identity() + self.k_norm = nn.Identity() + + def _init_sliding_window(self, config): + return None + + +class LlamaDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_idx: int): + super().__init__() + self.self_attn = LlamaAttention(config=config, layer_idx=layer_idx) + self.mlp = LlamaMLP(config) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = False, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs: Unpack[AttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + **kwargs, + ) + # Fully Connected + hidden_states, residual = self.post_attention_layernorm( + hidden_states, + residual=residual, + prenorm=True, + ) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + if output_attentions: + outputs += (self_attn_weights,) + + return outputs + + +class LlamaPreTrainedModel(XorlPreTrainedModel): + config_class = Llama3Config + base_model_prefix = "model" + _no_split_modules = ["LlamaDecoderLayer"] + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, RMSNorm): + module.weight.data.fill_(1.0) + elif isinstance(module, RotaryEmbedding): + inv_freq, module.attention_scaling = module.rope_init_fn(module.config, module.inv_freq.device) + module.inv_freq.copy_(inv_freq) + module.original_inv_freq = module.inv_freq + + def get_checkpoint_handler(self, **kwargs): + if getattr(self, "_unfused_for_tp", False): + return None + + weights_path = kwargs.get("weights_path", None) + is_prequantized = detect_prequantized_checkpoint(weights_path) + if not is_prequantized: + is_prequantized = detect_prequantized_block_fp8_checkpoint(weights_path) + + exclude_modules = getattr(self, "_qlora_exclude_modules", None) + if exclude_modules is None: + exclude_modules = get_prequantized_exclude_modules(weights_path) if is_prequantized else set() + + head_dim = getattr(self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads) + return Llama3CheckpointHandler( + num_attention_heads=self.config.num_attention_heads, + num_key_value_heads=self.config.num_key_value_heads, + head_dim=head_dim, + is_prequantized=is_prequantized, + exclude_modules=exclude_modules, + model=self if is_prequantized else None, + ) + + +class LlamaModel(LlamaPreTrainedModel): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. + Each layer is a [`LlamaDecoderLayer`]. + """ + + def __init__(self, config): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [LlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = RotaryEmbedding(config=config) + + self.gradient_checkpointing = False + self._skip_causal_mask = is_flash_attention(config._attn_implementation) + + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + **flash_attn_kwargs: Unpack[AttentionKwargs], + ) -> BaseModelOutput: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + # PP support: when embed_tokens is None, input is already hidden_states + if self.embed_tokens is not None: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + hidden_states = inputs_embeds + else: + hidden_states = input_ids if inputs_embeds is None else inputs_embeds + + if position_ids is None: + position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0) + + if self._skip_causal_mask: + causal_mask = None + else: + cache_position = torch.arange(hidden_states.shape[1], device=hidden_states.device) + causal_mask = update_causal_mask( + self.config._attn_implementation, + attention_mask, + hidden_states, + cache_position, + is_training=self.training, + output_attentions=output_attentions, + ) + + # Create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + # SP strategy handles slicing (sync: slice, async: keep full-length) + ps = get_parallel_state() + position_embeddings = get_cp_strategy(num_kv_heads=self.config.num_key_value_heads).prepare_position_embeddings( + position_embeddings, + dim=1, + sp_group=ps.sp_group, + num_kv_heads=self.config.num_key_value_heads, + ) + + # Decoder layers + all_self_attns = () if output_attentions else None + + for decoder_layer in self.layers: + if decoder_layer is None: # PP: pruned layer + continue + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + output_attentions=output_attentions, + position_embeddings=position_embeddings, + **flash_attn_kwargs, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + # PP support: norm may be None on non-last stages + hidden_states = self.norm(hidden_states) if self.norm is not None else hidden_states + + return BaseModelOutput( + last_hidden_state=hidden_states, + attentions=all_self_attns, + ) + + +class LlamaForCausalLM(LlamaPreTrainedModel): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + + _tp_plan = parallelize.MODEL_TP_PLAN + + def __init__(self, config): + super().__init__(config) + self.model = LlamaModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + self.post_init() + + def unfuse_for_tp(self): + """Unfuse all fused projections for tensor parallelism compatibility.""" + parallelize.unfuse_for_tp(self) + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def get_pp_module_config(self): + """Return PP module config for pipeline_module_split.""" + return { + "input_fqns": ["model.embed_tokens"], + "layer_prefix": "model.layers", + "output_fqns": ["model.norm", "lm_head"], + "always_keep_fqns": ["model.rotary_emb"], + "num_layers": self.config.num_hidden_layers, + } + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> CausalLMOutput: + outputs: BaseModelOutput = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + last_hidden_state = outputs.last_hidden_state + + return CausalLMOutput(last_hidden_state=last_hidden_state) + + +ModelClass = LlamaForCausalLM + +__all__ = ["LlamaForCausalLM", "LlamaModel", "LlamaPreTrainedModel"] diff --git a/src/xorl/models/transformers/llama3/parallelize.py b/src/xorl/models/transformers/llama3/parallelize.py new file mode 100644 index 00000000..72a39976 --- /dev/null +++ b/src/xorl/models/transformers/llama3/parallelize.py @@ -0,0 +1,37 @@ +"""Parallelization plan and utilities for Llama 3 models.""" + +# TP plan for the base model (Llama3Model). +# Keys use wildcard patterns relative to the base model prefix. +TP_PLAN = { + "embed_tokens": "embedding", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", +} + +# TP plan for top-level modules on the CausalLM wrapper. +MODEL_TP_PLAN = { + "lm_head": "colwise", +} + + +def unfuse_for_tp(model): + """Unfuse fused projections for tensor parallelism compatibility. + + Splits ``qkv_proj`` -> ``q_proj / k_proj / v_proj`` in attention, + and ``gate_up_proj`` -> ``gate_proj / up_proj`` in MLP, for every + decoder layer. + + After unfusing, checkpoint keys from HuggingFace already match + the model's parameter names -- no merging handler is needed. + """ + for layer in model.model.layers: + layer.self_attn.unfuse_for_tp() + layer.mlp.unfuse_for_tp() + model._unfused_for_tp = True + # Override HF config's TP plan with our plan for unfused projections. + model.config.base_model_tp_plan = TP_PLAN From a160a7df3da31901643bd597db676b8c8db4cb69 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Thu, 16 Apr 2026 11:38:10 -0700 Subject: [PATCH 33/41] feat(distributed): unified eFSDP gradient scaling with factor=1.0 (#83) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(distributed): implement torchtitan eFSDP design for EP+FSDP2 Aligns the EP+FSDP2 implementation with torchtitan's eFSDP design, which separates the DP-shard dimension into two named sub-meshes and guarantees correct gradient handling regardless of optimizer choice. Key changes: **torch_parallelize.py — _build_ep_param_groups()** Add this helper and call it at the end of parallelize_model_fsdp2() when EP is enabled. Previously _ep_param_groups was only set by the EP-aware optimizer builder, so ep_fsdp2_clip_grad_norm silently fell back to the non-EP path when using a different optimizer or before optimizer init. Now it is always set right after FSDP wrapping. Classification: _skip_fsdp params -> ep group; DTensors on ep_fsdp mesh -> ep group; everything else -> non_ep group. **parallel_state.py — torchtitan-aligned mesh properties** Add dp_shard_mod_ep_mesh/size and dp_shard_in_ep_mesh/size using torchtitan's naming convention: dp_shard_mod_ep = full DP shard mesh (all ep_size x ep_fsdp ranks) FSDP mesh for non-expert params dp_shard_in_ep = sub-mesh within one EP group (ep_fsdp ranks) FSDP mesh for expert params **trainer.py — peak memory logging** Always print [PEAK_MEMORY] peak_alloc_gb=... to stdout in _finalize(). Parseable by benchmark scripts. **scripts/benchmark_ep_fsdp_30b.py — eFSDP benchmark** New script testing EP=1/2/4/8 on 8 GPUs with a 30B-proxy MoE config. Reports tok/s, TFLOPs/GPU, peak GPU memory, grad norm, speedup per EP. Supports --model-size full for actual Qwen3-30B-A3B runs on a cluster. * fix(distributed): align eFSDP gradient division with torchtitan design Expert gradients were being divided by ep_size twice: once via set_gradient_divide_factor(ep_size) in torch_parallelize.py and again via mul_(1/ep_size) in clip_grad_norm.py. This caused ep_size² over-division of expert gradients, destabilising training at large EP. torchtitan's design (disable_fsdp_gradient_division) sets the divide factor to 1.0 for ALL modules including experts, relying on the loss normalisation for uniform gradient handling. xorl now matches this. Also adds EP=8+eFSDP=2 benchmark config/k8s manifests and fixes the EP=16 reference config (correct dataset_prepared_path on shared PVC, non-login shell to avoid .bashrc env issues in K8s pods). * fix(moe): FP32 scatter_add in MoE combine for EP-size-independent precision BF16 scatter_add_ in the MoE combine step produces order-dependent rounding when accumulating top-K expert outputs per token. Different EP sizes create different accumulation orders (different GPU placements), causing routing divergence that cascades through 48 MoE layers. Changes: - utils.py: FP32 scatter_add in alltoall combine (eliminates 100% of EP-size divergence for alltoall dispatch, verified bit-identical) - deepep.py: FP32 scatter_add + index_add in DeepEP combine/backward (reduces divergence; residual diff from DeepEP's two-stage architecture) - torch_parallelize.py: improved comments explaining why factor=1.0 is correct for both expert and non-expert FSDP gradient division - clip_grad_norm.py: add per-group (expert vs non-expert) gradient norm logging for debugging gradient scale issues - trainer.py: add diagnostic tools (_save_moe_diagnostics, _save_gradient_diagnostics, _save_logprob_diagnostics) for comparing MoE routing indices, hidden states, and gradients across EP configs Verification: - alltoall + FP32 scatter_add: 0.0000% output diff, 100% routing match across EP4/EP8/EP16 at all 48 layers (bit-identical) - DeepEP + FP32 scatter_add: 0.29% residual diff at layer 0 from NVSHMEM pre-accumulation of partial sums (architectural limitation) - Expert gradient norms identical across EP configs (5.647 vs 5.650) - No measurable throughput overhead from FP32 scatter_add * feat(moe): configurable DeepEP combine dtype (DEEPEP_COMBINE_DTYPE) Add support for FP16/FP32 combine precision in DeepEP dispatch, controlled by the DEEPEP_COMBINE_DTYPE env var ("bf16", "fp16", "fp32"). Changes: - deepep.py: add _get_combine_dtype() cached helper; cast scatter_add output to the configured combine dtype before buffer.combine() - deepep.py: cast combined_x back to model dtype after combine receives - trainer.py: refactor MoE diagnostics to use hooks on the training forward pass instead of a separate forward (avoids BF16/FP16 combine dtype conflict that caused hangs) Benchmark results (235B, 8-node, EP8+eFSDP8): BF16 combine: 313 TFLOPs, 18.5s/step (baseline) FP16 combine: 305 TFLOPs, 19.0s/step (+2.7%) FP32 combine: 306 TFLOPs, 18.9s/step (+2.2%) Precision (30B, 2-node, layer 0 output diff EP8 vs EP16): BF16: 0.292% (16382 tokens differ) FP16: 0.091% (16382 tokens, 3.2× better) FP32: 0.000015% (2 tokens, 19500× better) Requires DeepEP built with FP16/FP32 combine kernel support (nv_half + float template instantiation in SWITCH_TYPES macro). * refactor(moe): remove DEEPEP_COMBINE_DTYPE env var Stick with BF16 combine for DeepEP. The FP16/FP32 combine options showed minimal benefit at 235B scale (FP16: +24% overhead on EP16, FP32: crashed EP16) while the FP32 scatter_add already handles precision for the alltoall path. The DeepEP combine still uses FP32 scatter_add on the expert GPU before casting to BF16 for NVSHMEM transfer. * Add benchmark configs, k8s manifests, diagnostics arg, and eFSDP grad tests - Add save_step_diagnostics training arg for correctness verification - Add Qwen3-235B 8-node and Qwen3-Coder-30B 2-node benchmark configs (EP/eFSDP variants: alltoall, deepep, fp32, correctness, long) - Add corresponding k8s manifests with NCCL IB and nccl node-group tolerations - Update existing configs for single-step diagnostics runs - Add eFSDP gradient scaling tests * cleanup: remove diagnostics, debug logging, and excess benchmark configs - Remove save_step_diagnostics arg and all diagnostic methods from trainer - Remove debug gradient scale logging from clip_grad_norm - Remove test scripts and benchmark script - Remove diagnostic/long/fp32/alltoall config variants (keep core EP configs) - Revert debug settings in base configs (max_steps back to 50) * fix: merge resolution, train_router default, lint fixes - Add missing weighted scatter line dropped during merge - Default MoEBlock.train_router=False (matches args default, fixes deepep) - Move peak memory logging to debug level - Remove unused get_lm_head_weight import (ruff) * style: remove redundant FSDPModule import (ruff) * Fix DeepEP double-scoring and remove dead PEAK_MEMORY log deepep.py: Remove score multiplication from _FusedUnpermuteAndCombine.forward(). Router weights are already applied by the expert compute function (triton/native backends multiply by expert_scores). The previous code double-applied scores (output * scores^2) and the backward path didn't account for the factor, creating a gradient mismatch. Keep FP32 accumulation improvement. trainer.py: Remove [PEAK_MEMORY] debug log that was gated behind DEBUG level but the default is INFO, making it dead code. * Address review: train_router default, double-scoring, stale env var 1. MoEBlock constructor default back to train_router=True so direct callers (tests) still exercise router gradient propagation. Config- level default remains False via getattr(config, "train_router", False). 2. Remove score multiplication from _FusedUnpermuteAndCombine.forward() — expert compute already applies scores, this was double-counting. Backward didn't account for the factor either (gradient mismatch). 3. Remove dead [PEAK_MEMORY] debug log (gated behind DEBUG, default is INFO). 4. Remove DEEPEP_COMBINE_DTYPE from k8s manifests — env var was removed in this branch, so these configs no longer represent distinct settings. * Fix LoRA EP path: apply router scores after compute LoRA EP compute functions (triton/native) don't accept expert_scores, unlike the non-LoRA versions which multiply by scores internally. After removing the score multiply from _FusedUnpermuteAndCombine, the LoRA path had no score application at all. Fix: apply scores in MoEExpertsLoRA._ep_forward() after compute returns, matching the non-LoRA path where scores are applied inside the compute function. Both paths now apply scores exactly once. * Remove user-specific paths from benchmark configs and k8s manifests Replace hardcoded /home/qywu paths with generic /workspace paths matching the existing manifest conventions. Remove dataset_prepared_path (not used in other configs) and use relative output_dir paths. * Add review fixes: tests, ruff, _skip_fsdp guard, rename ep_outside → ep_intranode - Fix ruff PLC0415: add noqa to DTensor import in _build_ep_param_groups - Add _skip_fsdp + ep_fsdp_size > 1 assertion (missing eFSDP gradient sync) - Add TestEPLoRARouterScores: verify score application for alltoall/deepep contexts, no-scores identity, and gradient flow - Add test_ep_clip_grad_norm.py: 16 tests covering param classification, norm computation, clipping, no double-division, _skip_fsdp end-to-end - Rename ep_outside → ep_intranode (inverted semantics, default True): clearer name for "EP all-to-all stays within the node (NVLink)" * Apply ruff-format to new test files * Block all LoRA experts (not just QLoRA) with ep_fsdp_size > 1 LoRA + EP + eFSDP gradient interaction has not been validated. Guard both regular LoRA (MoEExpertsLoRA) and QLoRA (_skip_fsdp) expert modules when ep_fsdp_size > 1 until correctness is verified. (cherry picked from commit d4c08c06f960a47c093e5ae64afb785b17eb8dfc) --- .../content/docs/config-reference/local.md | 2 +- .../docs/parallelism/expert_parallelism.mdx | 16 +- src/xorl/arguments.py | 6 +- src/xorl/distributed/fsdp2/clip_grad_norm.py | 12 +- src/xorl/distributed/moe/deepep.py | 15 +- src/xorl/distributed/parallel_state.py | 58 ++- src/xorl/distributed/torch_parallelize.py | 107 ++++- src/xorl/models/layers/moe/lora.py | 7 + src/xorl/trainers/trainer.py | 1 + tests/distributed/test_ep_clip_grad_norm.py | 391 ++++++++++++++++++ tests/distributed/test_parallel_state.py | 16 +- tests/models/test_moe_experts_lora.py | 164 ++++++++ 12 files changed, 745 insertions(+), 50 deletions(-) create mode 100644 tests/distributed/test_ep_clip_grad_norm.py diff --git a/docs/src/content/docs/config-reference/local.md b/docs/src/content/docs/config-reference/local.md index 7808a12d..803564e2 100644 --- a/docs/src/content/docs/config-reference/local.md +++ b/docs/src/content/docs/config-reference/local.md @@ -112,7 +112,7 @@ Each entry in `datasets` (or `test_datasets`) is a dict: | `ringattn_parallel_size` | `1` | Ring Attention degree. | | `cp_fsdp_mode` | `all` | How context parallelism interacts with FSDP: `all` (both Ulysses+Ring), `ulysses_only`, `ring_only`, `none`. | | `reshard_after_forward` | `null` | FSDP2 reshard after forward. `true` = save memory, `false` = save communication (used for PP by default). `null` = auto. | -| `ep_outside` | `false` | Place EP outside the EP-FSDP mesh. | +| `ep_intranode` | `true` | Keep EP all-to-all within the node (NVLink). Set `false` to place EP across nodes. | :::caution[Field interactions] - `data_parallel_mode: fsdp2` **requires** `init_device: meta` diff --git a/docs/src/content/docs/parallelism/expert_parallelism.mdx b/docs/src/content/docs/parallelism/expert_parallelism.mdx index 52263818..4ad2a7ea 100644 --- a/docs/src/content/docs/parallelism/expert_parallelism.mdx +++ b/docs/src/content/docs/parallelism/expert_parallelism.mdx @@ -404,19 +404,19 @@ With PP, a 3-D mesh `(_pp_ep, ep, ep_fsdp)` is created so that each PP stage get own EP process groups. Slicing by `["ep"]` automatically returns the correct per-stage submesh for the calling rank. -### `ep_outside` flag +### `ep_intranode` flag -By default (`ep_outside=False`), the rank assignment within each `[ep_size, ep_fsdp_size]` -block interleaves EP and FSDP ranks so that consecutive ranks belong to the same EP group. -Setting `ep_outside=True` transposes the layout so consecutive ranks belong to the same -FSDP group instead. +By default (`ep_intranode=True`), consecutive ranks belong to the same EP group, +keeping EP all-to-all traffic within the node (fast NVLink). Setting `ep_intranode=False` +transposes the layout so consecutive ranks belong to the same FSDP group instead, +placing EP across nodes (IB/RDMA). `init_ep_mesh_matrix()` implements both layouts: ```python -# ep_outside=False (default): FSDP-major, EP-minor +# ep_intranode=True (default): EP ranks are consecutive (intra-node) mesh = arange(ep_size * ep_fsdp_size).view(ep_fsdp_size, ep_size).T -# ep_outside=True: EP-major, FSDP-minor +# ep_intranode=False: EP spans across nodes mesh = arange(ep_size * ep_fsdp_size).view(ep_size, ep_fsdp_size) ``` @@ -833,7 +833,7 @@ ranks. | Parameter | Type | Default | Description | |---|---|---|---| | `expert_parallel_size` | `int` | `1` | Number of EP ranks. Must divide `num_experts` and `ranks_per_pp_stage`. | -| `ep_outside` | `bool` | `False` | When `True`, consecutive ranks are in the same EP-FSDP group rather than the same EP group. | +| `ep_intranode` | `bool` | `True` | When `True` (default), EP all-to-all uses consecutive ranks (intra-node NVLink). When `False`, EP spans across nodes. | | `data_parallel_shard_size` | `int` | `-1` | FSDP shard size. Together with `expert_parallel_size`, determines `ep_fsdp_size = ranks_per_stage / ep_size`. | | `pipeline_parallel_size` | `int` | `1` | PP degree. EP groups are confined within each PP stage. | | `ringattn_parallel_size` | `int` | `1` | Ring attention CP size. Can be folded into EP-FSDP axis via `ep_fsdp_size`. | diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 899b21a6..27c8a10d 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -793,9 +793,9 @@ def moe_recomputed(self) -> bool: default=1, metadata={"help": "Expert parallel size."}, ) - ep_outside: bool = field( - default=False, - metadata={"help": "Enable expert parallelism outside in ep-fsdp."}, + ep_intranode: bool = field( + default=True, + metadata={"help": "Place EP all-to-all within the node (NVLink). When False, EP spans across nodes."}, ) ulysses_parallel_size: int = field( default=1, diff --git a/src/xorl/distributed/fsdp2/clip_grad_norm.py b/src/xorl/distributed/fsdp2/clip_grad_norm.py index 785a99d7..ef695722 100644 --- a/src/xorl/distributed/fsdp2/clip_grad_norm.py +++ b/src/xorl/distributed/fsdp2/clip_grad_norm.py @@ -92,13 +92,11 @@ def ep_fsdp2_clip_grad_norm( p for p in model._ep_param_groups.get("non_ep", []) if p.grad is not None ] - # Average FSDP-sharded EP gradients across EP ranks - if ps.ep_enabled and ps.ep_size > 1 and ep_fsdp_params: - scale = 1.0 / float(ps.ep_size) - for q in ep_fsdp_params: - if q.grad is not None: - q.grad.detach().mul_(scale) - # Note: ep_local_params (e.g., QLoRAMoeExperts LoRA) are NOT scaled — + # Note: torchtitan eFSDP design disables FSDP's automatic gradient division for ALL + # modules (including experts) via disable_fsdp_gradient_division. Gradient normalisation + # is handled uniformly by the loss (gradient_accumulate_loss). No additional EP-specific + # scaling is needed here — doing so would double-divide expert grads. + # ep_local_params (_skip_fsdp, e.g. QLoRAMoeExperts LoRA) are also left unscaled: # each rank trains its own unique local experts independently. # Compute and reduce non-EP norms across FSDP group diff --git a/src/xorl/distributed/moe/deepep.py b/src/xorl/distributed/moe/deepep.py index beb31f2c..100d834e 100644 --- a/src/xorl/distributed/moe/deepep.py +++ b/src/xorl/distributed/moe/deepep.py @@ -325,14 +325,15 @@ def backward(ctx, grad_expert_input, grad_cumsum, grad_dispatch_ctx): buffer = ctx.buffer handle = ctx.handle - # Step 1: Scatter gradient from expert order → recv order + # Step 1: Scatter gradient from expert order → recv order (FP32 for precision) grad_recv_x = torch.zeros( ctx.num_recv_tokens, ctx.hidden_dim, - dtype=grad_expert_input.dtype, + dtype=torch.float32, device=grad_expert_input.device, ) - grad_recv_x.index_add_(0, permuted_indices, grad_expert_input) + grad_recv_x.index_add_(0, permuted_indices, grad_expert_input.float()) + grad_recv_x = grad_recv_x.to(grad_expert_input.dtype) # Step 2: Combine to reverse dispatch previous_event = EventOverlap(EventHandle()) @@ -385,17 +386,21 @@ def forward( device=device, ) else: + # Scores are already applied by the expert compute function + # (triton/native backends multiply by expert_scores). Do NOT + # re-apply here — that would double-count router weights. gather_output = torch.zeros( dispatch_ctx.num_recv_tokens, hidden_dim, - dtype=dtype, + dtype=torch.float32, device=device, ) idx_2d = dispatch_ctx.permuted_indices.unsqueeze(1).expand(-1, hidden_dim) _CHUNK = 4096 for _i in range(0, expert_output.shape[0], _CHUNK): _end = min(_i + _CHUNK, expert_output.shape[0]) - gather_output.scatter_add_(0, idx_2d[_i:_end], expert_output[_i:_end]) + gather_output.scatter_add_(0, idx_2d[_i:_end], expert_output[_i:_end].float()) + gather_output = gather_output.to(dtype) # Step 2: Combine previous_event = EventOverlap(EventHandle()) diff --git a/src/xorl/distributed/parallel_state.py b/src/xorl/distributed/parallel_state.py index e005c530..3f9c7dd4 100644 --- a/src/xorl/distributed/parallel_state.py +++ b/src/xorl/distributed/parallel_state.py @@ -39,15 +39,16 @@ def _inner(self: "ParallelState", *args, **kwargs): return _inner -def init_ep_mesh_matrix(ep_size: int, ep_fsdp_size: int, ep_outside: bool = False) -> "DeviceMesh": +def init_ep_mesh_matrix(ep_size: int, ep_fsdp_size: int, ep_intranode: bool = True) -> "DeviceMesh": """ Initialize the device mesh matrix for the EP. Args: ep_size (int): The size of the EP. ep_fsdp_size (int): The size of the EP-FSDP. - ep_outside (bool): Whether the EP is outside in ep-fsdp group. + ep_intranode (bool): When True (default), EP all-to-all uses consecutive + ranks (intra-node NVLink). When False, EP spans across nodes. """ - if ep_outside: + if not ep_intranode: with torch.device("cpu"): mesh = torch.arange(math.prod((ep_size, ep_fsdp_size)), dtype=torch.int).view(ep_size, ep_fsdp_size) else: @@ -404,6 +405,47 @@ def ep_enabled(self) -> bool: def ep_rank(self) -> int: return self.ep_fsdp_device_mesh.get_local_rank("ep") + # --- Torchtitan eFSDP naming (aligns with torchtitan's 6-D mesh design) --- + + @property + @requires_mesh + def dp_shard_mod_ep_mesh(self) -> "DeviceMesh": + """FSDP mesh for non-expert params: all DP-shard ranks (ep × ep_fsdp). + + Torchtitan names this ``dp_shard_mod_ep``: the full shard dimension that + *contains* EP groups. Non-expert params are sharded across all ranks in + this mesh, so every GPU holds a shard of every weight. + """ + return self.fsdp_mesh + + @property + def dp_shard_mod_ep_size(self) -> int: + """Total ranks participating in non-expert FSDP (= ep_size × ep_fsdp_size).""" + return self.fsdp_size + + @property + @requires_mesh + def dp_shard_in_ep_mesh(self) -> "DeviceMesh": + """FSDP mesh for expert params: ranks *within* the same EP group. + + Torchtitan names this ``dp_shard_in_ep``: the inner shard dimension that + operates *inside* a single EP group. Expert params are FSDP-sharded only + across these ranks; different EP groups hold different expert weights. + + When EP is disabled, falls back to the full FSDP mesh (all ranks). + """ + if self.ep_enabled and self.ep_fsdp_device_mesh is not None: + return self.ep_fsdp_device_mesh["ep_fsdp"] + return self.fsdp_mesh + + @property + def dp_shard_in_ep_size(self) -> int: + """Ranks per EP group used for expert FSDP (= total_dp / ep_size).""" + if self.ep_enabled: + world_size = dist.get_world_size() if dist.is_initialized() else 1 + return (world_size // self.pp_size) // self.ep_size + return self.fsdp_size + # ------------------------------ SP (sequence parallel) ------------------------------ # @property def sp_group(self) -> Optional["ProcessGroup"]: @@ -533,7 +575,7 @@ def init_parallel_state( dp_mode: Literal["none", "ddp", "fsdp2"] = "fsdp2", device_type: str = None, cp_fsdp_mode: Literal["all", "ulysses_only", "ring_only", "none"] = "all", - ep_outside: bool = False, + ep_intranode: bool = True, ) -> None: """ Initializes global parallel state. @@ -643,10 +685,10 @@ def _safe_flatten(mesh, dim_names, flat_name): for pp_stage in range(pp_size): stage_start = pp_stage * ranks_per_stage stage_ranks = torch.arange(stage_start, stage_start + ranks_per_stage, dtype=torch.int) - if ep_outside: - pp_ep_mesh[pp_stage] = stage_ranks.view(ep_size, ep_fsdp_size) - else: + if ep_intranode: pp_ep_mesh[pp_stage] = stage_ranks.view(ep_fsdp_size, ep_size).transpose(0, 1) + else: + pp_ep_mesh[pp_stage] = stage_ranks.view(ep_size, ep_fsdp_size) ep_fsdp_device_mesh = DeviceMesh( device_type=device_type, @@ -654,7 +696,7 @@ def _safe_flatten(mesh, dim_names, flat_name): mesh_dim_names=("_pp_ep", "ep", "ep_fsdp"), ) else: - mesh = init_ep_mesh_matrix(ep_size=ep_size, ep_fsdp_size=ep_fsdp_size, ep_outside=ep_outside) + mesh = init_ep_mesh_matrix(ep_size=ep_size, ep_fsdp_size=ep_fsdp_size, ep_intranode=ep_intranode) ep_fsdp_device_mesh = DeviceMesh( device_type=device_type, mesh=mesh, diff --git a/src/xorl/distributed/torch_parallelize.py b/src/xorl/distributed/torch_parallelize.py index e08ba949..e92dbca1 100644 --- a/src/xorl/distributed/torch_parallelize.py +++ b/src/xorl/distributed/torch_parallelize.py @@ -220,13 +220,18 @@ def _experts_shard_placement_fn(param): layer_mod._fsdp_modules = [] # ep enabled and this layer contains the expert module if parallel_state.ep_enabled and experts_mod is not None and not getattr(experts_mod, "_skip_fsdp", False): + # Block LoRA experts with eFSDP > 1: LoRA + EP + eFSDP gradient + # interaction has not been validated. Until it is, require ep_fsdp_size=1. + from xorl.lora.modules.base import LoraModule # noqa: PLC0415 + + if isinstance(experts_mod, LoraModule): + assert parallel_state.dp_shard_in_ep_size <= 1, ( + f"LoRA expert modules are not yet supported with ep_fsdp_size > 1 " + f"(got {parallel_state.dp_shard_in_ep_size}). " + f"Use ep_fsdp_size=1 for LoRA + EP training." + ) # shard expert fully_shard(experts_mod, **expert_fsdp_kwargs) - if hasattr(experts_mod, "set_gradient_divide_factor"): - # average EP grads across EP ranks - experts_mod.set_gradient_divide_factor(parallel_state.ep_size) - # mark so the global divide-factor reset below doesn't override this - experts_mod._is_ep_fsdp = True layer_mod._fsdp_modules.append(experts_mod) # shard module that needs to ignore mixed precision control if mp_ignored_classes: @@ -244,6 +249,19 @@ def _experts_shard_placement_fn(param): # (each EP rank has different local expert LoRA params — they should not be all-gathered globally) layer_fsdp_kwargs = dict(fsdp_kwargs) if parallel_state.ep_enabled and experts_mod is not None and getattr(experts_mod, "_skip_fsdp", False): + # _skip_fsdp params are excluded from FSDP, so they receive no automatic + # gradient all-reduce across eFSDP replicas. When ep_fsdp_size > 1 the + # same expert weights are replicated on each eFSDP rank but process + # different data (separate all-to-all per EP group), producing divergent + # gradients. Until an explicit eFSDP gradient sync is added for these + # params, block the unsupported combination. + assert parallel_state.dp_shard_in_ep_size <= 1, ( + f"_skip_fsdp expert modules (e.g. QLoRA) are not yet supported with " + f"ep_fsdp_size > 1 (got {parallel_state.dp_shard_in_ep_size}). " + f"Per-expert LoRA params have no gradient sync across eFSDP replicas, " + f"which would cause weight divergence. Use ep_fsdp_size=1 or remove " + f"_skip_fsdp (wrap experts with FSDP instead)." + ) layer_fsdp_kwargs["ignored_params"] = set(experts_mod.parameters()) fully_shard(layer_mod, **layer_fsdp_kwargs) layer_mod._fsdp_modules.append(layer_mod) @@ -297,13 +315,19 @@ def _experts_shard_placement_fn(param): prefetch_modules = prev_block._fsdp_modules current_block.set_modules_to_backward_prefetch(list(reversed(prefetch_modules))) - # Disable FSDP's automatic gradient averaging — we normalize gradients manually - # (gradient_accumulate_loss for non-PP; explicit grad.mul_(1/gvt) for PP). - # Skip EP-sharded expert modules: they manage their own divide factor (ep_size) - # for averaging expert gradients across EP ranks. + # Disable FSDP's automatic gradient averaging for ALL modules (including EP experts). + # Gradient normalisation is handled uniformly by the loss (gradient_accumulate_loss + # for non-PP; explicit grad.mul_(1/gvt) for PP). + # + # factor=1.0 is correct for BOTH expert and non-expert modules because: + # - Non-expert: SUM across fsdp_mesh (all 64 ranks), each contributing grad scaled + # by local_valid_tokens/global_valid_tokens = 1/world_size → correct mean. + # - Expert: each eFSDP rank already accumulated gradients from ALL GPUs in its EP + # group (via all-to-all dispatch). SUM across eFSDP replicas covers all data + # replicas, so the total = (1/world_size) * all tokens routed to this expert → correct. for module in model.modules(): - if isinstance(module, FSDPModule) and not getattr(module, "_is_ep_fsdp", False): + if isinstance(module, FSDPModule): module.set_gradient_divide_factor(1.0) # Handle meta initialization for FSDP2 (fallback if pre-load not done) @@ -334,12 +358,75 @@ def _experts_shard_placement_fn(param): logger.info_rank0("Every rank reading weights from disk independently...") all_ranks_load_weights(model, weights_path, weight_device, dtensor_factory=distribute_tensor) + # Build EP param groups now (torchtitan eFSDP design: track groups at wrap time, + # not just at optimizer build time, so grad clipping works regardless of optimizer + # choice and before the optimizer is constructed). + if parallel_state.ep_enabled: + _build_ep_param_groups(model) + # Register grad norm clipping method for FSDP2 model.clip_grad_norm_ = types.MethodType(clip_grad_norm, model) return model +def _build_ep_param_groups(model: "nn.Module") -> None: + """Classify post-FSDP parameters into EP and non-EP groups. + + Implements the torchtitan eFSDP design: after ``fully_shard()`` wraps expert + modules with the ep_fsdp sub-mesh and non-expert modules with the full DP mesh, + we classify every parameter by whether it lives on the ep_fsdp mesh. The + resulting ``model._ep_param_groups`` dict is consumed by + ``ep_fsdp2_clip_grad_norm`` to apply the correct reduction and scaling to each + group. + + Classification rules: + - ``_skip_fsdp`` expert params (plain tensors, e.g. QLoRA LoRA per-expert): + EP group — each rank holds unique expert weights, no cross-rank reduction. + - DTensor whose device mesh dim names include ``"ep_fsdp"``: + EP group — sharded within the ep_fsdp sub-mesh. + - Everything else: + non-EP group — standard FSDP sharding across the full DP mesh. + + Note: The optimizer builder (``build_ep_fsdp2_optimizer``) rebuilds this dict + to apply additional ``requires_grad`` filtering and Muon/AdamW classification. + Setting it here ensures grad clipping is correct even when using a non-EP-aware + optimizer or before the optimizer is constructed. + """ + try: + from torch.distributed._tensor import DTensor # noqa: PLC0415 + except ImportError: + DTensor = None + + # _skip_fsdp expert modules hold plain tensor params (not wrapped by FSDP) + skip_fsdp_param_ids: set = set() + for m in model.modules(): + if getattr(m, "_skip_fsdp", False): + for p in m.parameters(): + skip_fsdp_param_ids.add(id(p)) + + ep_params = [] + non_ep_params = [] + for p in model.parameters(): + if id(p) in skip_fsdp_param_ids: + ep_params.append(p) + continue + if DTensor is not None and isinstance(p, DTensor): + mesh = getattr(p, "device_mesh", None) + names = getattr(mesh, "mesh_dim_names", ()) if mesh is not None else () + if "ep_fsdp" in names: + ep_params.append(p) + continue + non_ep_params.append(p) + + model._ep_param_groups = {"ep": ep_params, "non_ep": non_ep_params} + logger.info_rank0( + f"eFSDP param groups (torchtitan design): " + f"{len(ep_params)} EP params (ep_fsdp mesh), " + f"{len(non_ep_params)} non-EP params (full DP mesh)" + ) + + def build_parallelize_model( model: "nn.Module", weights_path: Optional[str] = None, diff --git a/src/xorl/models/layers/moe/lora.py b/src/xorl/models/layers/moe/lora.py index fad35392..837ad78b 100644 --- a/src/xorl/models/layers/moe/lora.py +++ b/src/xorl/models/layers/moe/lora.py @@ -285,6 +285,13 @@ def _ep_forward( self.scaling, ) + # Apply router scores — LoRA compute functions don't accept + # expert_scores, so apply them here (matches the non-LoRA path + # where scores are applied inside the compute function). + expert_scores = getattr(ctx, "expert_scores", getattr(ctx, "permuted_scores", None)) + if expert_scores is not None: + expert_output = expert_output * expert_scores.unsqueeze(1).to(expert_output.dtype) + # Step 3: Combine expert outputs back to original ranks combine_kwargs = self._build_combine_kwargs(expert_output, ctx, dispatch_kwargs, parallel_state) return combine_fn(**combine_kwargs) diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index 6467cfda..dc58df5c 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -848,6 +848,7 @@ def train_step(self, micro_batches: List[Dict[str, Any]]) -> Tuple[float, float] RoutingReplay.clear_all() self._sync_sp_gradients() + grad_norm = self._clip_and_step() self._maybe_merge_lora() total_loss, grad_norm = self._reduce_metrics(total_loss, grad_norm) diff --git a/tests/distributed/test_ep_clip_grad_norm.py b/tests/distributed/test_ep_clip_grad_norm.py new file mode 100644 index 00000000..41792fbd --- /dev/null +++ b/tests/distributed/test_ep_clip_grad_norm.py @@ -0,0 +1,391 @@ +"""Tests for EP-aware gradient clipping (clip_grad_norm / ep_fsdp2_clip_grad_norm). + +Verifies: +1. _build_ep_param_groups correctly classifies _skip_fsdp and plain params. +2. ep_fsdp2_clip_grad_norm computes the correct total norm from all three + parameter groups (non-EP, EP-FSDP, EP-local) and clips uniformly. +3. No double-division of EP gradients (the bug this branch fixes). +4. inf-norm path works correctly. + +All tests run single-rank: process groups are None so all_reduce is skipped, +letting us verify the local math without distributed infrastructure. +""" + +import math +from unittest.mock import MagicMock, patch + +import pytest +import torch +import torch.nn as nn + +from xorl.distributed.fsdp2.clip_grad_norm import ( + clip_grad_norm, + ep_fsdp2_clip_grad_norm, +) +from xorl.distributed.torch_parallelize import _build_ep_param_groups + + +pytestmark = [pytest.mark.cpu] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_param(*shape, grad=None): + """Create a plain parameter with an optional gradient.""" + p = nn.Parameter(torch.randn(*shape)) + if grad is not None: + p.grad = grad + return p + + +def _mock_parallel_state(ep_enabled=True): + """Return a mock parallel state with all groups set to None (single-rank).""" + ps = MagicMock() + ps.ep_enabled = ep_enabled + ps.fsdp_group = None + ps.ep_group = None + ps.tp_enabled = False + ps.tp_group = None + ps.ep_fsdp_device_mesh = None + return ps + + +def _l2_norm(*params): + """Compute the expected L2 norm across multiple params' gradients.""" + total = 0.0 + for p in params: + if p.grad is not None: + total += p.grad.detach().to(torch.float32).norm(2).item() ** 2 + return math.sqrt(total) + + +# --------------------------------------------------------------------------- +# 1. _build_ep_param_groups classification +# --------------------------------------------------------------------------- + + +class TestBuildEPParamGroups: + """Test that _build_ep_param_groups classifies params correctly.""" + + def test_skip_fsdp_params_classified_as_ep(self): + """Params from _skip_fsdp modules go into the EP group.""" + model = nn.Module() + # Regular submodule + regular = nn.Linear(8, 8) + model.add_module("regular", regular) + # _skip_fsdp submodule (e.g. QLoRAMoeExperts) + expert = nn.Linear(8, 8) + expert._skip_fsdp = True + model.add_module("expert", expert) + + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): + _build_ep_param_groups(model) + + assert hasattr(model, "_ep_param_groups") + ep_ids = {id(p) for p in model._ep_param_groups["ep"]} + non_ep_ids = {id(p) for p in model._ep_param_groups["non_ep"]} + + # Expert params in EP group + for p in expert.parameters(): + assert id(p) in ep_ids + # Regular params in non-EP group + for p in regular.parameters(): + assert id(p) in non_ep_ids + + def test_no_skip_fsdp_all_non_ep(self): + """Without _skip_fsdp modules, all plain params go to non-EP.""" + model = nn.Module() + model.add_module("linear", nn.Linear(8, 8)) + + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): + _build_ep_param_groups(model) + + assert len(model._ep_param_groups["ep"]) == 0 + assert len(model._ep_param_groups["non_ep"]) == 2 # weight + bias + + def test_nested_skip_fsdp_params_all_classified(self): + """All params inside nested _skip_fsdp modules are classified as EP.""" + model = nn.Module() + # Mimic QLoRAMoeExperts: a _skip_fsdp module with multiple sub-params + expert = nn.Module() + expert._skip_fsdp = True + expert.lora_A = nn.Parameter(torch.randn(4, 32, 8)) + expert.lora_B = nn.Parameter(torch.randn(4, 8, 64)) + expert.base_weight = nn.Parameter(torch.randn(4, 32, 64), requires_grad=False) + model.add_module("expert", expert) + model.add_module("non_expert", nn.Linear(32, 32)) + + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): + _build_ep_param_groups(model) + + ep_ids = {id(p) for p in model._ep_param_groups["ep"]} + assert id(expert.lora_A) in ep_ids + assert id(expert.lora_B) in ep_ids + assert id(expert.base_weight) in ep_ids # even frozen params are classified + # non-expert params not in EP group + non_ep_ids = {id(p) for p in model._ep_param_groups["non_ep"]} + for p in model.non_expert.parameters(): + assert id(p) in non_ep_ids + + +# --------------------------------------------------------------------------- +# 2. ep_fsdp2_clip_grad_norm: norm computation and clipping +# --------------------------------------------------------------------------- + + +class TestEPFSDP2ClipGradNorm: + """Test norm computation and gradient clipping logic.""" + + def _setup_model(self, ep_grads, non_ep_grads): + """Create a model with _ep_param_groups populated from given gradient tensors. + + Args: + ep_grads: list of gradient tensors for EP-local params. + non_ep_grads: list of gradient tensors for non-EP params. + + Returns: + (model, ep_params, non_ep_params) + """ + ep_params = [] + for g in ep_grads: + p = _make_param(*g.shape, grad=g) + ep_params.append(p) + + non_ep_params = [] + for g in non_ep_grads: + p = _make_param(*g.shape, grad=g) + non_ep_params.append(p) + + model = MagicMock() + model._ep_param_groups = {"ep": ep_params, "non_ep": non_ep_params} + return model, ep_params, non_ep_params + + def test_l2_norm_single_group(self): + """Total L2 norm is correct when only non-EP params have grads.""" + g = torch.tensor([3.0, 4.0]) # norm = 5 + model, _, non_ep = self._setup_model(ep_grads=[], non_ep_grads=[g]) + + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): + total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=100.0) + + assert total_norm.item() == pytest.approx(5.0, abs=1e-5) + + def test_l2_norm_combined_groups(self): + """Total L2 norm combines EP-local and non-EP norms correctly.""" + ep_g = torch.tensor([3.0, 0.0]) # norm = 3 + non_ep_g = torch.tensor([0.0, 4.0]) # norm = 4 + # combined: sqrt(9 + 16) = 5 + model, ep_params, non_ep_params = self._setup_model(ep_grads=[ep_g], non_ep_grads=[non_ep_g]) + + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): + total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=100.0) + + expected = _l2_norm(*ep_params, *non_ep_params) + assert total_norm.item() == pytest.approx(expected, abs=1e-5) + + def test_clipping_reduces_gradients(self): + """When total_norm > max_norm, all gradients are scaled down uniformly.""" + ep_g = torch.tensor([6.0, 0.0]) + non_ep_g = torch.tensor([0.0, 8.0]) + # total norm = sqrt(36 + 64) = 10 + model, ep_params, non_ep_params = self._setup_model(ep_grads=[ep_g], non_ep_grads=[non_ep_g]) + max_norm = 5.0 # clip factor = 5/10 = 0.5 + + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): + total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=max_norm) + + assert total_norm.item() == pytest.approx(10.0, abs=1e-5) + # Both EP and non-EP grads should be scaled by 0.5 + torch.testing.assert_close(ep_params[0].grad, torch.tensor([3.0, 0.0])) + torch.testing.assert_close(non_ep_params[0].grad, torch.tensor([0.0, 4.0])) + + def test_no_clipping_below_max(self): + """When total_norm <= max_norm, gradients are unchanged.""" + g = torch.tensor([3.0, 4.0]) # norm = 5 + model, _, non_ep = self._setup_model(ep_grads=[], non_ep_grads=[g]) + + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): + total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=10.0) + + # Gradients unchanged + torch.testing.assert_close(non_ep[0].grad, torch.tensor([3.0, 4.0])) + + def test_ep_grads_not_double_scaled(self): + """EP gradients are NOT divided by ep_size — the double-division fix.""" + ep_g = torch.tensor([6.0, 8.0]) # norm = 10 + model, ep_params, _ = self._setup_model(ep_grads=[ep_g], non_ep_grads=[]) + + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): + total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=100.0) + + # Gradients should be completely unchanged (no scaling, no division) + torch.testing.assert_close(ep_params[0].grad, torch.tensor([6.0, 8.0])) + assert total_norm.item() == pytest.approx(10.0, abs=1e-5) + + def test_inf_norm(self): + """Inf-norm returns the max absolute gradient value.""" + ep_g = torch.tensor([3.0, -7.0]) + non_ep_g = torch.tensor([5.0, 2.0]) + model, _, _ = self._setup_model(ep_grads=[ep_g], non_ep_grads=[non_ep_g]) + + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): + total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=100.0, norm_type=float("inf")) + + assert total_norm.item() == pytest.approx(7.0, abs=1e-5) + + def test_inf_norm_clips_correctly(self): + """Inf-norm clipping scales gradients when max element exceeds max_norm.""" + ep_g = torch.tensor([3.0, -10.0]) + non_ep_g = torch.tensor([5.0, 2.0]) + model, ep_params, non_ep_params = self._setup_model(ep_grads=[ep_g], non_ep_grads=[non_ep_g]) + max_norm = 5.0 # total inf-norm is 10, clip factor = 5/10 = 0.5 + + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): + total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=max_norm, norm_type=float("inf")) + + assert total_norm.item() == pytest.approx(10.0, abs=1e-5) + torch.testing.assert_close(ep_params[0].grad, torch.tensor([1.5, -5.0])) + torch.testing.assert_close(non_ep_params[0].grad, torch.tensor([2.5, 1.0])) + + def test_empty_groups(self): + """Handles empty parameter groups without errors.""" + model = MagicMock() + model._ep_param_groups = {"ep": [], "non_ep": []} + + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): + total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=1.0) + + assert total_norm.item() == 0.0 + + def test_params_without_grads_skipped(self): + """Params with grad=None are excluded from norm computation.""" + g = torch.tensor([3.0, 4.0]) # norm = 5 + p_with_grad = _make_param(2, grad=g) + p_no_grad = _make_param(2) # no gradient + + model = MagicMock() + model._ep_param_groups = {"ep": [p_no_grad], "non_ep": [p_with_grad]} + + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): + total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=100.0) + + assert total_norm.item() == pytest.approx(5.0, abs=1e-5) + + +# --------------------------------------------------------------------------- +# 3. _skip_fsdp end-to-end: classify → clip +# --------------------------------------------------------------------------- + + +class TestSkipFSDPClipEndToEnd: + """End-to-end test: _skip_fsdp params flow through classification into clipping. + + Mimics the QLoRA EP path where expert LoRA params are plain tensors + (not FSDP-managed) and must be: + - Classified as EP by _build_ep_param_groups + - Treated as ep_local_params (no reduction) in ep_fsdp2_clip_grad_norm + - Not scaled/divided by ep_size (the double-division fix) + - Clipped with the same coefficient as non-EP params + """ + + def test_skip_fsdp_classify_then_clip(self): + """_skip_fsdp expert grads are classified as EP-local and clipped correctly.""" + model = nn.Module() + + # Non-expert param (mimics attention/mlp weights) + regular = nn.Linear(4, 4, bias=False) + model.add_module("regular", regular) + + # _skip_fsdp expert param (mimics QLoRA LoRA weights) + expert = nn.Module() + expert._skip_fsdp = True + expert.lora = nn.Parameter(torch.randn(2, 4)) + model.add_module("expert", expert) + + # Assign known gradients: expert norm=6, regular norm=8, total=10 + expert.lora.grad = torch.tensor([[3.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]]) # flat norm = 3 + # Wait, let me make the math cleaner + expert.lora.grad = torch.zeros(2, 4) + expert.lora.grad[0, 0] = 6.0 # norm = 6 + regular.weight.grad = torch.zeros(4, 4) + regular.weight.grad[0, 0] = 8.0 # norm = 8 + # total = sqrt(36 + 64) = 10 + + ps = _mock_parallel_state() + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=ps): + _build_ep_param_groups(model) + + # Verify classification + ep_ids = {id(p) for p in model._ep_param_groups["ep"]} + assert id(expert.lora) in ep_ids + non_ep_ids = {id(p) for p in model._ep_param_groups["non_ep"]} + assert id(regular.weight) in non_ep_ids + + # Clip with max_norm=5 → clip_coeff = 5/10 = 0.5 + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=ps): + total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=5.0) + + assert total_norm.item() == pytest.approx(10.0, abs=1e-5) + # Both groups scaled uniformly by 0.5 + assert expert.lora.grad[0, 0].item() == pytest.approx(3.0, abs=1e-5) + assert regular.weight.grad[0, 0].item() == pytest.approx(4.0, abs=1e-5) + + def test_skip_fsdp_grads_not_reduced_or_divided(self): + """_skip_fsdp grads contribute their raw local norm — no all-reduce, no ep_size division.""" + model = nn.Module() + + expert = nn.Module() + expert._skip_fsdp = True + expert.weight = nn.Parameter(torch.randn(4, 8)) + expert.weight.grad = torch.full((4, 8), 2.0) # norm = 2 * sqrt(32) + model.add_module("expert", expert) + + ps = _mock_parallel_state() + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=ps): + _build_ep_param_groups(model) + total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=1000.0) + + expected_norm = torch.tensor(2.0 * math.sqrt(32)) + assert total_norm.item() == pytest.approx(expected_norm.item(), abs=1e-4) + # Gradient unchanged (no scaling applied since norm < max_norm) + assert (expert.weight.grad == 2.0).all() + + +# --------------------------------------------------------------------------- +# 4. clip_grad_norm dispatch +# --------------------------------------------------------------------------- + + +class TestClipGradNormDispatch: + """Test that clip_grad_norm dispatches to ep_fsdp2_clip_grad_norm when appropriate.""" + + def test_dispatches_to_ep_path_when_ep_param_groups_present(self): + """Models with _ep_param_groups use the EP-aware clip path.""" + g = torch.tensor([3.0, 4.0]) + p = _make_param(2, grad=g) + model = MagicMock() + model._ep_param_groups = {"ep": [], "non_ep": [p]} + # hasattr check needs to work + model.__dict__["_ep_param_groups"] = model._ep_param_groups + + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): + total_norm = clip_grad_norm(model, max_norm=100.0) + + assert total_norm.item() == pytest.approx(5.0, abs=1e-5) + + def test_falls_through_without_ep_param_groups(self): + """Models without _ep_param_groups use the standard FSDP2 path.""" + g = torch.tensor([3.0, 4.0]) + p = _make_param(2, grad=g) + + model = MagicMock(spec=[]) # empty spec, so hasattr(_ep_param_groups) is False + model.parameters = MagicMock(return_value=iter([p])) + + ps = _mock_parallel_state(ep_enabled=False) + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=ps): + total_norm = clip_grad_norm(model, max_norm=100.0) + + assert total_norm.item() == pytest.approx(5.0, abs=1e-5) diff --git a/tests/distributed/test_parallel_state.py b/tests/distributed/test_parallel_state.py index cbcab429..c4a34aae 100644 --- a/tests/distributed/test_parallel_state.py +++ b/tests/distributed/test_parallel_state.py @@ -22,18 +22,18 @@ class TestEPMeshMatrixAndRequiresMesh: def test_ep_mesh_matrix_and_requires_mesh(self): """EP mesh matrix: row-major, transposed, edge cases; requires_mesh raises/allows correctly.""" - # ep_outside=True: row-major - mesh = init_ep_mesh_matrix(ep_size=2, ep_fsdp_size=4, ep_outside=True) + # ep_intranode=True (default): consecutive ranks in same EP group (intra-node) + mesh = init_ep_mesh_matrix(ep_size=2, ep_fsdp_size=4, ep_intranode=True) assert mesh.shape == (2, 4) and mesh.dtype == torch.int - assert torch.equal(mesh, torch.arange(8).view(2, 4)) - - # ep_outside=False: transposed - mesh = init_ep_mesh_matrix(ep_size=2, ep_fsdp_size=4, ep_outside=False) assert torch.equal(mesh, torch.arange(8).view(4, 2).transpose(0, 1)) + # ep_intranode=False: EP spans across nodes + mesh = init_ep_mesh_matrix(ep_size=2, ep_fsdp_size=4, ep_intranode=False) + assert torch.equal(mesh, torch.arange(8).view(2, 4)) + # Edge: single ep_size / ep_fsdp_size - assert torch.equal(init_ep_mesh_matrix(1, 4, True), torch.arange(4).unsqueeze(0)) - assert torch.equal(init_ep_mesh_matrix(4, 1, True), torch.arange(4).unsqueeze(1)) + assert torch.equal(init_ep_mesh_matrix(1, 4, ep_intranode=False), torch.arange(4).unsqueeze(0)) + assert torch.equal(init_ep_mesh_matrix(4, 1, ep_intranode=False), torch.arange(4).unsqueeze(1)) # requires_mesh decorator class MC: diff --git a/tests/models/test_moe_experts_lora.py b/tests/models/test_moe_experts_lora.py index d99cbedd..b88a9c76 100644 --- a/tests/models/test_moe_experts_lora.py +++ b/tests/models/test_moe_experts_lora.py @@ -668,5 +668,169 @@ def __init__(self): assert isinstance(model.v_proj, LoraLinear) +# --------------------------------------------------------------------------- +# 7. EP LoRA router-score application +# --------------------------------------------------------------------------- + + +class TestEPLoRARouterScores: + """Test that MoEExpertsLoRA._ep_forward applies router scores from the dispatch context. + + The LoRA EP compute functions don't accept expert_scores (unlike the non-LoRA + path), so _ep_forward must apply them post-compute. This test verifies: + - Scores from "expert_scores" (alltoall) and "permuted_scores" (deepep) are applied + - Missing scores leave the output unchanged + - Gradients flow through the score multiplication to LoRA parameters + """ + + NUM_EXPERTS = 4 + HIDDEN_DIM = 32 + INTERMEDIATE = 64 + R = 4 + NUM_TOKENS = 8 + + def _make_experts(self): + experts = MoEExpertsLoRA( + num_experts=self.NUM_EXPERTS, + hidden_dim=self.HIDDEN_DIM, + intermediate_size=self.INTERMEDIATE, + moe_implementation="native", + lora_config=MoELoRAConfig(r=self.R, lora_alpha=8), + ) + nn.init.xavier_normal_(experts.gate_up_proj.data) + nn.init.xavier_normal_(experts.down_proj.data) + experts.ep_dispatch = "alltoall" + return experts + + def _run_ep_forward(self, experts, score_attr=None, scores=None, compute_output=None): + """Run _ep_forward with mocked dispatch/compute/combine. + + Returns (final_output, expert_output_passed_to_combine). + """ + from dataclasses import make_dataclass + from unittest.mock import MagicMock, patch + + if compute_output is None: + compute_output = torch.randn(self.NUM_TOKENS, self.HIDDEN_DIM) + + # Build dispatch context with the requested score attribute + fields = [] + if score_attr is not None: + fields.append((score_attr, torch.Tensor)) + Ctx = make_dataclass("Ctx", fields) + ctx = Ctx(**({score_attr: scores} if score_attr else {})) + + mock_dispatch = MagicMock( + return_value=( + torch.randn(self.NUM_TOKENS, self.HIDDEN_DIM), # permute_tokens + torch.tensor([2, 4, 6, 8]), # cumsum + ctx, + ) + ) + mock_compute = MagicMock(return_value=compute_output) + + captured = {} + + def mock_combine(**kwargs): + captured["expert_output"] = kwargs["expert_output"] + return kwargs["expert_output"] + + mock_ps = MagicMock() + mock_ps.ep_enabled = True + mock_ps.ep_group = MagicMock() + + with ( + patch.dict("xorl.models.layers.moe.lora.EP_DISPATCH", {"alltoall": mock_dispatch}), + patch.dict("xorl.models.layers.moe.lora.EP_COMBINE", {"alltoall": mock_combine}), + patch.dict("xorl.models.layers.moe.lora.EP_EXPERT_COMPUTE_LORA", {"native": mock_compute}), + patch("xorl.distributed.parallel_state.get_parallel_state", return_value=mock_ps), + ): + hidden = torch.randn(self.NUM_TOKENS, self.HIDDEN_DIM) + routing_weights = torch.softmax(torch.randn(self.NUM_TOKENS, 2), dim=-1) + selected_experts = torch.randint(0, self.NUM_EXPERTS, (self.NUM_TOKENS, 2)) + + output = experts(hidden, routing_weights, selected_experts) + + return output, captured["expert_output"] + + @pytest.mark.parametrize("score_attr", ["expert_scores", "permuted_scores"]) + def test_scores_applied_to_output(self, score_attr): + """Router scores from dispatch context are multiplied into expert output.""" + experts = self._make_experts() + compute_output = torch.randn(self.NUM_TOKENS, self.HIDDEN_DIM) + scores = torch.rand(self.NUM_TOKENS) * 0.5 + 0.1 # non-trivial scores in (0.1, 0.6) + + _, expert_output = self._run_ep_forward( + experts, + score_attr=score_attr, + scores=scores, + compute_output=compute_output, + ) + + expected = compute_output * scores.unsqueeze(1) + torch.testing.assert_close(expert_output, expected) + + def test_no_scores_leaves_output_unchanged(self): + """When dispatch context has no score attribute, expert output is unchanged.""" + experts = self._make_experts() + compute_output = torch.randn(self.NUM_TOKENS, self.HIDDEN_DIM) + + _, expert_output = self._run_ep_forward( + experts, + score_attr=None, + compute_output=compute_output, + ) + + torch.testing.assert_close(expert_output, compute_output) + + def test_gradient_flows_through_scores(self): + """Gradients from the score multiplication reach LoRA parameters.""" + from dataclasses import make_dataclass + from unittest.mock import MagicMock, patch + + experts = self._make_experts() + compute_output = torch.randn(self.NUM_TOKENS, self.HIDDEN_DIM, requires_grad=True) + scores = torch.rand(self.NUM_TOKENS, requires_grad=True) + + Ctx = make_dataclass("Ctx", [("expert_scores", torch.Tensor)]) + ctx = Ctx(expert_scores=scores) + + mock_dispatch = MagicMock( + return_value=( + torch.randn(self.NUM_TOKENS, self.HIDDEN_DIM), + torch.tensor([2, 4, 6, 8]), + ctx, + ) + ) + mock_compute = MagicMock(return_value=compute_output) + + def mock_combine(**kwargs): + return kwargs["expert_output"] + + mock_ps = MagicMock() + mock_ps.ep_enabled = True + mock_ps.ep_group = MagicMock() + + with ( + patch.dict("xorl.models.layers.moe.lora.EP_DISPATCH", {"alltoall": mock_dispatch}), + patch.dict("xorl.models.layers.moe.lora.EP_COMBINE", {"alltoall": mock_combine}), + patch.dict("xorl.models.layers.moe.lora.EP_EXPERT_COMPUTE_LORA", {"native": mock_compute}), + patch("xorl.distributed.parallel_state.get_parallel_state", return_value=mock_ps), + ): + hidden = torch.randn(self.NUM_TOKENS, self.HIDDEN_DIM) + routing_weights = torch.softmax(torch.randn(self.NUM_TOKENS, 2), dim=-1) + selected_experts = torch.randint(0, self.NUM_EXPERTS, (self.NUM_TOKENS, 2)) + + output = experts(hidden, routing_weights, selected_experts) + output.sum().backward() + + # Gradient should flow through scores back to compute_output + assert compute_output.grad is not None + assert scores.grad is not None + # scores.grad should equal the sum of (compute_output * grad_output) per token + expected_score_grad = (compute_output.detach() * 1.0).sum(dim=1) # grad_output is all 1s from .sum() + torch.testing.assert_close(scores.grad, expected_score_grad) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) From a9223b0da01e34cc476ef39ecffbb08e6a90700a Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Thu, 16 Apr 2026 14:30:16 -0700 Subject: [PATCH 34/41] Fix Llama 3 LoRA target names and gate/up bias load asymmetry (#145) The Llama 3 LoRA/QLoRA example configs listed unfused HF-style projection names (q_proj, k_proj, v_proj, gate_proj, up_proj), but the model keeps qkv_proj and gate_up_proj fused by default (merge_qkv=True). The LoRA matcher is literal, so only o_proj and down_proj resolved and the rest were silently skipped. Switch both configs to the fused names, matching the Qwen examples. GateUpMergeBuffer's pattern only matched .weight, so gate_proj.bias / up_proj.bias passed through unmerged on load while save-side splitting (substring match on .gate_up_proj.) already handled bias. Widen the pattern to (weight|bias) and key pending state by (prefix, param_type), mirroring QKVMergeBuffer. Other consumers gate the match behind key.endswith(".weight"), so the widened pattern is inert for them. (cherry picked from commit d020c39ae54bae6e8f228b490daebfc3136cdd42) --- .../dummy/configs/lora/llama3_8b_lora.yaml | 2 +- .../configs/qlora/llama3_8b_qlora_nvfp4.yaml | 2 +- .../models/checkpoint_handlers/buffers.py | 37 ++++++++++--------- 3 files changed, 21 insertions(+), 20 deletions(-) diff --git a/examples/local/dummy/configs/lora/llama3_8b_lora.yaml b/examples/local/dummy/configs/lora/llama3_8b_lora.yaml index cf82f128..b922bb59 100644 --- a/examples/local/dummy/configs/lora/llama3_8b_lora.yaml +++ b/examples/local/dummy/configs/lora/llama3_8b_lora.yaml @@ -50,5 +50,5 @@ lora: enable_lora: true lora_rank: 16 lora_alpha: 16 - lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] + lora_target_modules: ["qkv_proj", "o_proj", "gate_up_proj", "down_proj"] save_lora_only: true diff --git a/examples/local/dummy/configs/qlora/llama3_8b_qlora_nvfp4.yaml b/examples/local/dummy/configs/qlora/llama3_8b_qlora_nvfp4.yaml index c7891aa8..ffd3a529 100644 --- a/examples/local/dummy/configs/qlora/llama3_8b_qlora_nvfp4.yaml +++ b/examples/local/dummy/configs/qlora/llama3_8b_qlora_nvfp4.yaml @@ -51,6 +51,6 @@ lora: enable_qlora: true lora_rank: 16 lora_alpha: 16 - lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] + lora_target_modules: ["qkv_proj", "o_proj", "gate_up_proj", "down_proj"] quant_format: nvfp4 quant_group_size: 16 diff --git a/src/xorl/models/checkpoint_handlers/buffers.py b/src/xorl/models/checkpoint_handlers/buffers.py index ad27e66c..58d562eb 100644 --- a/src/xorl/models/checkpoint_handlers/buffers.py +++ b/src/xorl/models/checkpoint_handlers/buffers.py @@ -54,9 +54,9 @@ # e.g., model.layers.0.mlp.experts.gate_up_proj FUSED_EXPERT_PATTERN = re.compile(r"^model\.layers\.\d+\.mlp\.experts\.(gate_up|gate|up|down)_proj$") -# Dense MLP gate/up projection weight keys -# e.g., model.layers.0.mlp.gate_proj.weight or model.layers.0.mlp.shared_expert.up_proj.weight -DENSE_GATE_UP_PATTERN = re.compile(r"^(.*)\.(gate|up)_proj\.weight$") +# Dense MLP gate/up projection weight/bias keys +# e.g., model.layers.0.mlp.gate_proj.weight or model.layers.0.mlp.shared_expert.up_proj.bias +DENSE_GATE_UP_PATTERN = re.compile(r"^(.*)\.(gate|up)_proj\.(weight|bias)$") # Attention o_proj weight key # e.g., model.layers.0.self_attn.o_proj.weight @@ -464,31 +464,32 @@ class GateUpMergeBuffer: When a model uses a merged gate_up_proj linear layer but the checkpoint has separate gate_proj and up_proj weights, this buffer collects both and concatenates - them into the merged format. + them into the merged format. Handles both weight and bias parameters. """ def __init__(self): - self._pending: Dict[str, Dict[str, torch.Tensor]] = {} + # _pending: {(prefix, param_type): {"gate": tensor, "up": tensor}} + self._pending: Dict[Tuple[str, str], Dict[str, torch.Tensor]] = {} def add(self, key: str, tensor: torch.Tensor) -> Optional[Tuple[str, torch.Tensor]]: - """Try to add a gate/up weight. Returns (merged_key, merged_tensor) when both are available.""" + """Try to add a gate/up weight or bias. Returns merged result when both are available.""" match = DENSE_GATE_UP_PATTERN.match(key) if match is None: return None prefix = match.group(1) proj_type = match.group(2) # "gate" or "up" + param_type = match.group(3) # "weight" or "bias" - if prefix not in self._pending: - self._pending[prefix] = {} - self._pending[prefix][proj_type] = tensor + buf_key = (prefix, param_type) + if buf_key not in self._pending: + self._pending[buf_key] = {} + self._pending[buf_key][proj_type] = tensor - if "gate" in self._pending[prefix] and "up" in self._pending[prefix]: - gate_tensor = self._pending[prefix]["gate"] - up_tensor = self._pending[prefix]["up"] - del self._pending[prefix] - merged = torch.cat([gate_tensor, up_tensor], dim=0) - return f"{prefix}.gate_up_proj.weight", merged + if "gate" in self._pending[buf_key] and "up" in self._pending[buf_key]: + parts = self._pending.pop(buf_key) + merged = torch.cat([parts["gate"], parts["up"]], dim=0) + return f"{prefix}.gate_up_proj.{param_type}", merged return None @@ -496,9 +497,9 @@ def is_gate_up_key(self, key: str) -> bool: """Check if a key matches the gate/up pattern.""" return DENSE_GATE_UP_PATTERN.match(key) is not None - def get_pending(self) -> Dict[str, List[str]]: - """Get pending (incomplete) merge pairs for debugging.""" - return {prefix: list(projs.keys()) for prefix, projs in self._pending.items()} + def get_pending(self) -> Dict[Tuple[str, str], List[str]]: + """Get pending (incomplete) merge groups for debugging.""" + return {key: list(projs.keys()) for key, projs in self._pending.items()} # ============================================================================= From 4dab134df72df880b06f200a00856b4ac608e279 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Thu, 16 Apr 2026 14:40:23 -0700 Subject: [PATCH 35/41] [Feat] Add Qwen2/Qwen2.5 model support (#94) * Fix circular imports from c02dbf6 (lora/mapping, qlora) Restore lazy imports that were incorrectly moved to top-level in "Enforce import-outside-top-level" (#74), creating circular chains: - lora/mapping.py: move MoEExperts import back into _ensure_moe_mapping() - qlora/__init__.py: lazy __getattr__ for get_prequantized_exclude_modules - qlora/utils.py: lazy-import detect_prequantized_* from buffers.py Co-authored-by: zzz0906 * [Feat] Add Qwen2ForCausalLM support for Qwen2.5 models Native xorl support for Qwen2/Qwen2.5 dense models (e.g. Qwen2.5-14B). - Qwen2Attention subclasses MultiHeadAttention (attention_bias=True in config handles QKV bias; only _init_sliding_window overridden) - Qwen2Config: Qwen2.5-14B defaults, rope_theta-in-rope_scaling handling, attention_bias=True, bos/eos token IDs - auto.py: Qwen2 config loading in _load_local_xorl_config - Checkpoint handler reuses Qwen3's (identical merge/split logic) - Benchmark config for 1-node 8xH100 with Ulysses CP=8 Verified argmax-exact match vs HF Qwen2ForCausalLM on Qwen2.5-14B. Tested: TP unfuse, PP=2, LoRA, QLoRA NVFP4. Co-authored-by: zzz0906 * Update output_dir path in qwen2_5_14b_muon_1node.yaml * Fix Qwen2 fused HF checkpoint compatibility * Fix Qwen2 mixed sliding-window attention to match HF per-layer masking Two correctness issues with eager attention on mixed full/sliding-window Qwen2 configs: 1. _init_sliding_window fallback (layer_idx >= max_window_layers) could override explicit layer_types entries, making full_attention layers use sliding window. Now treats layer_types as authoritative when present. 2. Qwen2Model.forward built one global causal mask from config.sliding_window and reused it for every layer. Now builds separate full-attention and sliding-attention masks and selects per layer via config.layer_types, matching HF's causal_mask_mapping approach. * Simplify attention bias: use attention_bias for QKV, hardcode output bias=False Remove attention_qkv_bias and attention_output_bias indirection from MultiHeadAttention and Qwen2Config. All current models have output bias=False; QKV bias is driven by config.attention_bias directly. * Fix duplicate MoE import in lora/mapping.py after merge with main --------- Co-authored-by: zzz0906 Co-authored-by: Ashwinee Panda (cherry picked from commit 794a16cc5d1361161e39d37753a26f5a862f33f0) --- .../distributed/sequence_parallel/strategy.py | 7 +- src/xorl/models/auto.py | 5 + .../layers/attention/multi_head_attention.py | 33 +- src/xorl/models/transformers/__init__.py | 2 + .../models/transformers/qwen2/__init__.py | 0 .../transformers/qwen2/checkpoint_handler.py | 10 + .../transformers/qwen2/configuration_qwen2.py | 167 ++++++++ .../transformers/qwen2/modeling_qwen2.py | 392 ++++++++++++++++++ .../models/transformers/qwen2/parallelize.py | 25 ++ src/xorl/qlora/utils.py | 7 +- tests/models/test_qwen2_support.py | 146 +++++++ 11 files changed, 770 insertions(+), 24 deletions(-) create mode 100644 src/xorl/models/transformers/qwen2/__init__.py create mode 100644 src/xorl/models/transformers/qwen2/checkpoint_handler.py create mode 100644 src/xorl/models/transformers/qwen2/configuration_qwen2.py create mode 100644 src/xorl/models/transformers/qwen2/modeling_qwen2.py create mode 100644 src/xorl/models/transformers/qwen2/parallelize.py create mode 100644 tests/models/test_qwen2_support.py diff --git a/src/xorl/distributed/sequence_parallel/strategy.py b/src/xorl/distributed/sequence_parallel/strategy.py index 1ab25279..c7134b75 100644 --- a/src/xorl/distributed/sequence_parallel/strategy.py +++ b/src/xorl/distributed/sequence_parallel/strategy.py @@ -238,7 +238,7 @@ def __init__(self, group, ulysses_size: int): def project_qkv(self, module, hidden_states, position_embeddings): # QLoRA fallback: weight is None (packed in quantized buffers). # Fall back to sync-style: module._project_qkv() + synchronous a2a. - if module.qkv_proj.weight is None: + if not hasattr(module, "qkv_proj") or module.qkv_proj.weight is None: q, k, v = module._project_qkv(hidden_states, position_embeddings) if q.ndim == 4 and q.size(0) == 1: q = q.squeeze(0) @@ -300,8 +300,9 @@ def project_qkv(self, module, hidden_states, position_embeddings): # q: [S_full, Hq/SP, D], k: [S_full, Hkv/SP, D], v: [S_full, Hkv/SP, D] # Apply RMSNorm and RoPE externally (after a2a, on full-length tensors) - q = module.q_norm(q) - k = module.k_norm(k) + if getattr(module, "_use_qk_norm", False): + q = module.q_norm(q) + k = module.k_norm(k) q = q.unsqueeze(0) k = k.unsqueeze(0) v = v.unsqueeze(0) diff --git a/src/xorl/models/auto.py b/src/xorl/models/auto.py index 9c0384fe..2015aad4 100644 --- a/src/xorl/models/auto.py +++ b/src/xorl/models/auto.py @@ -51,6 +51,11 @@ def _load_local_xorl_config( if model_type == "qwen3_5": return Qwen3_5Config.from_hf_config(_namespace_from_dict(config_dict)) + if model_type == "qwen2": + from .transformers.qwen2.configuration_qwen2 import Qwen2Config + + return Qwen2Config(**{k: v for k, v in config_dict.items() if not k.startswith("_")}) + return None diff --git a/src/xorl/models/layers/attention/multi_head_attention.py b/src/xorl/models/layers/attention/multi_head_attention.py index 765b3fa7..e195a018 100644 --- a/src/xorl/models/layers/attention/multi_head_attention.py +++ b/src/xorl/models/layers/attention/multi_head_attention.py @@ -33,14 +33,15 @@ def __init__(self, config, layer_idx: int): self.attention_dropout = config.attention_dropout self.is_causal = True + qkv_bias = getattr(config, "attention_bias", False) + self._use_qk_norm = getattr(config, "use_qk_norm", True) self.q_dim = config.num_attention_heads * self.head_dim self.kv_dim = config.num_key_value_heads * self.head_dim - self.qkv_proj = nn.Linear(config.hidden_size, self.q_dim + 2 * self.kv_dim, bias=config.attention_bias) - self.o_proj = nn.Linear( - config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias - ) - self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.qkv_proj = nn.Linear(config.hidden_size, self.q_dim + 2 * self.kv_dim, bias=qkv_bias) + self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) + if self._use_qk_norm: + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.sliding_window = self._init_sliding_window(config) # ------------------------------------------------------------------ # @@ -55,15 +56,10 @@ def unfuse_for_tp(self): """Replace fused qkv_proj with separate q_proj, k_proj, v_proj for tensor parallelism.""" device = self.qkv_proj.weight.device dtype = self.qkv_proj.weight.dtype - self.q_proj = nn.Linear( - self.config.hidden_size, self.q_dim, bias=self.config.attention_bias, device=device, dtype=dtype - ) - self.k_proj = nn.Linear( - self.config.hidden_size, self.kv_dim, bias=self.config.attention_bias, device=device, dtype=dtype - ) - self.v_proj = nn.Linear( - self.config.hidden_size, self.kv_dim, bias=self.config.attention_bias, device=device, dtype=dtype - ) + has_qkv_bias = self.qkv_proj.bias is not None + self.q_proj = nn.Linear(self.config.hidden_size, self.q_dim, bias=has_qkv_bias, device=device, dtype=dtype) + self.k_proj = nn.Linear(self.config.hidden_size, self.kv_dim, bias=has_qkv_bias, device=device, dtype=dtype) + self.v_proj = nn.Linear(self.config.hidden_size, self.kv_dim, bias=has_qkv_bias, device=device, dtype=dtype) del self.qkv_proj def _project_qkv( @@ -88,8 +84,11 @@ def _project_qkv( q = self.q_proj(hidden_states) k = self.k_proj(hidden_states) v = self.v_proj(hidden_states) - q = self.q_norm(q.view(hidden_shape)) - k = self.k_norm(k.view(hidden_shape)) + q = q.view(hidden_shape) + k = k.view(hidden_shape) + if self._use_qk_norm: + q = self.q_norm(q) + k = self.k_norm(k) v = v.view(hidden_shape) cos, sin = position_embeddings diff --git a/src/xorl/models/transformers/__init__.py b/src/xorl/models/transformers/__init__.py index 84d4d83f..a4bec082 100644 --- a/src/xorl/models/transformers/__init__.py +++ b/src/xorl/models/transformers/__init__.py @@ -1,5 +1,6 @@ from . import ( llama3, + qwen2, qwen3, qwen3_5, qwen3_5_moe, @@ -9,6 +10,7 @@ __all__ = [ "llama3", + "qwen2", "qwen3", "qwen3_5", "qwen3_moe", diff --git a/src/xorl/models/transformers/qwen2/__init__.py b/src/xorl/models/transformers/qwen2/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/xorl/models/transformers/qwen2/checkpoint_handler.py b/src/xorl/models/transformers/qwen2/checkpoint_handler.py new file mode 100644 index 00000000..32675166 --- /dev/null +++ b/src/xorl/models/transformers/qwen2/checkpoint_handler.py @@ -0,0 +1,10 @@ +"""Checkpoint handler for Qwen2 dense models. + +Identical to Qwen3CheckpointHandler — same merge/split logic for +qkv_proj and gate_up_proj fused projections. +""" + +from ..qwen3.checkpoint_handler import Qwen3CheckpointHandler + + +Qwen2CheckpointHandler = Qwen3CheckpointHandler diff --git a/src/xorl/models/transformers/qwen2/configuration_qwen2.py b/src/xorl/models/transformers/qwen2/configuration_qwen2.py new file mode 100644 index 00000000..623b2b11 --- /dev/null +++ b/src/xorl/models/transformers/qwen2/configuration_qwen2.py @@ -0,0 +1,167 @@ +"""Qwen2 model configuration for xorl. + +This keeps xorl's TP/PP metadata while preserving HuggingFace Qwen2/Qwen2.5 +semantics: fused internal projections, no Q/K RMSNorm, biased QKV, and a +bias-free output projection. +""" + +from transformers.configuration_utils import PretrainedConfig + +from .parallelize import TP_PLAN + + +def _cfg_get(value, key, default=None): + if value is None: + return default + if isinstance(value, dict): + return value.get(key, default) + return getattr(value, key, default) + + +def _cfg_to_dict(value): + if value is None: + return None + if isinstance(value, dict): + return dict(value) + if hasattr(value, "__dict__"): + return dict(vars(value)) + return value + + +class Qwen2Config(PretrainedConfig): + model_type = "qwen2" + + base_model_tp_plan = TP_PLAN + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + def __init__( + self, + vocab_size=152064, + hidden_size=5120, + intermediate_size=13824, + num_hidden_layers=48, + num_attention_heads=40, + num_key_value_heads=8, + head_dim=None, + hidden_act="silu", + max_position_embeddings=131072, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=False, + rope_theta=1000000.0, + rope_scaling=None, + attention_bias=True, + use_qk_norm=False, + use_sliding_window=False, + sliding_window=131072, + max_window_layers=48, + attention_dropout=0.0, + layer_types=None, + pad_token_id=None, + bos_token_id=151643, + eos_token_id=151643, + **kwargs, + ): + kwargs.setdefault("architectures", ["Qwen2ForCausalLM"]) + + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads if num_key_value_heads is not None else num_attention_heads + self.head_dim = head_dim or (hidden_size // num_attention_heads) + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.use_sliding_window = use_sliding_window + self.sliding_window = sliding_window if use_sliding_window else None + self.max_window_layers = max_window_layers + self.attention_dropout = attention_dropout + + if rope_scaling is not None and "type" in rope_scaling and "rope_type" not in rope_scaling: + rope_scaling = dict(rope_scaling) + rope_scaling["rope_type"] = rope_scaling["type"] + if rope_scaling is not None and "rope_theta" in rope_scaling: + rope_scaling = dict(rope_scaling) + self.rope_theta = rope_scaling.pop("rope_theta") + if not rope_scaling or rope_scaling == {"rope_type": "default"}: + rope_scaling = None + else: + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + + self.attention_bias = attention_bias + self.use_qk_norm = use_qk_norm + + if layer_types is None: + layer_types = [ + "sliding_attention" + if self.sliding_window is not None and layer_idx >= self.max_window_layers + else "full_attention" + for layer_idx in range(self.num_hidden_layers) + ] + self.layer_types = list(layer_types) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + @classmethod + def from_hf_config(cls, hf_config): + text_config = getattr(hf_config, "text_config", hf_config) + rope_params = getattr(text_config, "rope_parameters", None) + if rope_params is None: + rope_params = getattr(text_config, "rope_scaling", None) + + hidden_size = getattr(text_config, "hidden_size") + num_attention_heads = getattr(text_config, "num_attention_heads") + head_dim = getattr(text_config, "head_dim", hidden_size // num_attention_heads) + + return cls( + vocab_size=getattr(text_config, "vocab_size", getattr(hf_config, "vocab_size", 152064)), + hidden_size=hidden_size, + intermediate_size=getattr(text_config, "intermediate_size"), + num_hidden_layers=getattr(text_config, "num_hidden_layers"), + num_attention_heads=num_attention_heads, + num_key_value_heads=getattr(text_config, "num_key_value_heads", num_attention_heads), + head_dim=head_dim, + hidden_act=getattr(text_config, "hidden_act", "silu"), + max_position_embeddings=getattr(text_config, "max_position_embeddings"), + initializer_range=getattr(text_config, "initializer_range", 0.02), + rms_norm_eps=getattr(text_config, "rms_norm_eps", 1e-6), + use_cache=getattr(text_config, "use_cache", True), + tie_word_embeddings=getattr( + hf_config, "tie_word_embeddings", getattr(text_config, "tie_word_embeddings", False) + ), + rope_theta=_cfg_get(rope_params, "rope_theta", getattr(text_config, "rope_theta", 1000000.0)), + rope_scaling=_cfg_to_dict(rope_params), + attention_bias=getattr(text_config, "attention_bias", True), + use_qk_norm=False, + use_sliding_window=getattr(text_config, "use_sliding_window", False), + sliding_window=getattr(text_config, "sliding_window", None), + max_window_layers=getattr(text_config, "max_window_layers", getattr(text_config, "num_hidden_layers")), + attention_dropout=getattr(text_config, "attention_dropout", 0.0), + layer_types=getattr(text_config, "layer_types", None), + pad_token_id=getattr(text_config, "pad_token_id", getattr(hf_config, "pad_token_id", None)), + bos_token_id=getattr(text_config, "bos_token_id", getattr(hf_config, "bos_token_id", None)), + eos_token_id=getattr(text_config, "eos_token_id", getattr(hf_config, "eos_token_id", None)), + architectures=getattr( + hf_config, "architectures", getattr(text_config, "architectures", ["Qwen2ForCausalLM"]) + ), + model_type=getattr(text_config, "model_type", getattr(hf_config, "model_type", "qwen2")), + ) + + +__all__ = ["Qwen2Config"] diff --git a/src/xorl/models/transformers/qwen2/modeling_qwen2.py b/src/xorl/models/transformers/qwen2/modeling_qwen2.py new file mode 100644 index 00000000..50b107e4 --- /dev/null +++ b/src/xorl/models/transformers/qwen2/modeling_qwen2.py @@ -0,0 +1,392 @@ +"""Qwen2 dense model for xorl. + +Keeps fused internal projections for performance while preserving HuggingFace +Qwen2/Qwen2.5 semantics: biased fused QKV, bias-free output projection, and +no Q/K RMSNorm. +""" + +from typing import Optional, Tuple, Unpack + +import torch +from torch import nn + +from xorl.distributed.parallel_state import get_parallel_state +from xorl.distributed.sequence_parallel.strategy import get_cp_strategy +from xorl.models.base import XorlPreTrainedModel +from xorl.models.checkpoint_handlers.buffers import ( + detect_prequantized_block_fp8_checkpoint, + detect_prequantized_checkpoint, + get_prequantized_exclude_modules, +) +from xorl.models.layers import ACT2FN, RMSNorm, RotaryEmbedding +from xorl.models.layers.attention import ( + AttentionKwargs, + MultiHeadAttention, + is_flash_attention, + update_causal_mask, +) +from xorl.models.module_utils import GradientCheckpointingLayer +from xorl.models.outputs import BaseModelOutput, CausalLMOutput +from xorl.models.transformers.qwen2 import parallelize +from xorl.models.transformers.qwen2.checkpoint_handler import Qwen2CheckpointHandler +from xorl.models.transformers.qwen2.configuration_qwen2 import Qwen2Config +from xorl.ops.fused_silu_and_mul import fused_silu_and_mul + + +_RUNTIME_CONFIG_ATTRS = { + "_attn_implementation", + "_attention_cast_bf16", + "_commit_hash", + "_deepep_async_combine", + "_deepep_buffer_size_gb", + "_deepep_num_sms", + "_ep_dispatch", + "_lm_head_fp32", + "_name_or_path", + "_qlora_exclude_modules", + "_rmsnorm_mode", + "_rope_native", + "_router_fp32", + "_activation_native", + "train_router", +} + + +def _copy_runtime_config_attrs(source, target): + for name in _RUNTIME_CONFIG_ATTRS: + if hasattr(source, name): + setattr(target, name, getattr(source, name)) + + +def _adapt_qwen2_config(config): + if isinstance(config, Qwen2Config): + return config + if hasattr(config, "text_config") or getattr(config, "model_type", None) == "qwen2": + adapted = Qwen2Config.from_hf_config(config) + _copy_runtime_config_attrs(config, adapted) + return adapted + return config + + +class Qwen2MLP(nn.Module): + def __init__(self, config): + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_up_proj = nn.Linear(self.hidden_size, 2 * self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + self._use_fused_silu = config.hidden_act == "silu" and not getattr(config, "_activation_native", False) + + def unfuse_for_tp(self): + """Replace fused gate_up_proj with separate gate_proj and up_proj for tensor parallelism.""" + device = self.gate_up_proj.weight.device + dtype = self.gate_up_proj.weight.dtype + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False, device=device, dtype=dtype) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False, device=device, dtype=dtype) + del self.gate_up_proj + + def forward(self, x): + if hasattr(self, "gate_up_proj"): + if self._use_fused_silu: + x = fused_silu_and_mul(self.gate_up_proj(x)) + else: + gate, up = self.gate_up_proj(x).chunk(2, dim=-1) + x = self.act_fn(gate) * up + else: + x = self.act_fn(self.gate_proj(x)) * self.up_proj(x) + return self.down_proj(x) + + +class Qwen2Attention(MultiHeadAttention): + """Qwen2 attention with fused internal QKV projection and per-layer sliding window control.""" + + def _init_sliding_window(self, config): + layer_types = getattr(config, "layer_types", None) + self.layer_type = layer_types[self.layer_idx] if layer_types is not None else None + # When layer_types is present, treat it as authoritative (matches HF behaviour). + if self.layer_type is not None: + return config.sliding_window if self.layer_type == "sliding_attention" else None + # Fallback for configs without layer_types. + if ( + config.use_sliding_window + and getattr(config, "sliding_window", None) is not None + and self.layer_idx >= config.max_window_layers + ): + return config.sliding_window + return None + + +class Qwen2DecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Qwen2Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Qwen2Attention(config=config, layer_idx=layer_idx) + self.mlp = Qwen2MLP(config) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = False, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs: Unpack[AttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states, residual = self.post_attention_layernorm( + hidden_states, + residual=residual, + prenorm=True, + ) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + if output_attentions: + outputs += (self_attn_weights,) + + return outputs + + +class Qwen2PreTrainedModel(XorlPreTrainedModel): + config_class = Qwen2Config + base_model_prefix = "model" + _no_split_modules = ["Qwen2DecoderLayer"] + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, RMSNorm): + module.weight.data.fill_(1.0) + elif isinstance(module, RotaryEmbedding): + inv_freq, module.attention_scaling = module.rope_init_fn(module.config, module.inv_freq.device) + module.inv_freq.copy_(inv_freq) + module.original_inv_freq = module.inv_freq + + def get_checkpoint_handler(self, **kwargs): + if getattr(self, "_unfused_for_tp", False): + return None + + weights_path = kwargs.get("weights_path", None) + is_prequantized = detect_prequantized_checkpoint(weights_path) + if not is_prequantized: + is_prequantized = detect_prequantized_block_fp8_checkpoint(weights_path) + + exclude_modules = getattr(self, "_qlora_exclude_modules", None) + if exclude_modules is None: + exclude_modules = get_prequantized_exclude_modules(weights_path) if is_prequantized else set() + + head_dim = getattr(self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads) + return Qwen2CheckpointHandler( + num_attention_heads=self.config.num_attention_heads, + num_key_value_heads=self.config.num_key_value_heads, + head_dim=head_dim, + is_prequantized=is_prequantized, + exclude_modules=exclude_modules, + model=self if is_prequantized else None, + ) + + +class Qwen2Model(Qwen2PreTrainedModel): + def __init__(self, config: Qwen2Config): + config = _adapt_qwen2_config(config) + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [Qwen2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = RotaryEmbedding(config=config) + + self.gradient_checkpointing = False + self._skip_causal_mask = is_flash_attention(config._attn_implementation) + self.has_sliding_layers = "sliding_attention" in self.config.layer_types + + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + **flash_attn_kwargs: Unpack[AttentionKwargs], + ) -> BaseModelOutput: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + if self.embed_tokens is not None: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + hidden_states = inputs_embeds + else: + hidden_states = input_ids if inputs_embeds is None else inputs_embeds + + if position_ids is None: + position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0) + + if self._skip_causal_mask: + causal_mask_mapping = None + else: + cache_position = torch.arange(hidden_states.shape[1], device=hidden_states.device) + mask_kwargs = dict( + attention_mask=attention_mask, + input_tensor=hidden_states, + cache_position=cache_position, + is_training=self.training, + output_attentions=output_attentions, + ) + causal_mask_mapping = { + "full_attention": update_causal_mask( + self.config._attn_implementation, sliding_window=None, **mask_kwargs + ), + } + if self.has_sliding_layers: + causal_mask_mapping["sliding_attention"] = update_causal_mask( + self.config._attn_implementation, sliding_window=self.config.sliding_window, **mask_kwargs + ) + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + ps = get_parallel_state() + position_embeddings = get_cp_strategy(num_kv_heads=self.config.num_key_value_heads).prepare_position_embeddings( + position_embeddings, + dim=1, + sp_group=ps.sp_group, + num_kv_heads=self.config.num_key_value_heads, + ) + + all_self_attns = () if output_attentions else None + + for idx, decoder_layer in enumerate(self.layers): + if decoder_layer is None: + continue + layer_mask = causal_mask_mapping[self.config.layer_types[idx]] if causal_mask_mapping is not None else None + layer_outputs = decoder_layer( + hidden_states, + attention_mask=layer_mask, + position_ids=position_ids, + output_attentions=output_attentions, + position_embeddings=position_embeddings, + **flash_attn_kwargs, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) if self.norm is not None else hidden_states + + return BaseModelOutput( + last_hidden_state=hidden_states, + attentions=all_self_attns, + ) + + +class KwargsForCausalLM(AttentionKwargs): ... + + +class Qwen2ForCausalLM(Qwen2PreTrainedModel): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + + _tp_plan = parallelize.MODEL_TP_PLAN + + def __init__(self, config): + config = _adapt_qwen2_config(config) + super().__init__(config) + self.model = Qwen2Model(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + self.post_init() + + def unfuse_for_tp(self): + """Unfuse all fused projections for tensor parallelism compatibility.""" + parallelize.unfuse_for_tp(self) + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def get_pp_module_config(self): + return { + "input_fqns": ["model.embed_tokens"], + "layer_prefix": "model.layers", + "output_fqns": ["model.norm", "lm_head"], + "always_keep_fqns": ["model.rotary_emb"], + "num_layers": self.config.num_hidden_layers, + } + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> CausalLMOutput: + outputs: BaseModelOutput = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + last_hidden_state = outputs.last_hidden_state + + return CausalLMOutput(last_hidden_state=last_hidden_state) + + +ModelClass = Qwen2ForCausalLM + +__all__ = ["Qwen2ForCausalLM", "Qwen2Model", "Qwen2PreTrainedModel"] diff --git a/src/xorl/models/transformers/qwen2/parallelize.py b/src/xorl/models/transformers/qwen2/parallelize.py new file mode 100644 index 00000000..b9e4c543 --- /dev/null +++ b/src/xorl/models/transformers/qwen2/parallelize.py @@ -0,0 +1,25 @@ +"""Parallelization plan and utilities for Qwen2 dense models.""" + +TP_PLAN = { + "embed_tokens": "embedding", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", +} + +MODEL_TP_PLAN = { + "lm_head": "colwise", +} + + +def unfuse_for_tp(model): + """Unfuse fused projections for tensor parallelism compatibility.""" + for layer in model.model.layers: + layer.self_attn.unfuse_for_tp() + layer.mlp.unfuse_for_tp() + model._unfused_for_tp = True + model.config.base_model_tp_plan = TP_PLAN diff --git a/src/xorl/qlora/utils.py b/src/xorl/qlora/utils.py index 2e141317..613920f5 100644 --- a/src/xorl/qlora/utils.py +++ b/src/xorl/qlora/utils.py @@ -21,10 +21,6 @@ from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, cached_file from xorl.distributed.parallel_state import get_parallel_state -from xorl.models.checkpoint_handlers.buffers import ( - detect_prequantized_block_fp8_checkpoint, - detect_prequantized_checkpoint, -) from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear from xorl.qlora.modules.linear import QLoRALinear from xorl.qlora.modules.moe_experts import QLoRAMoeExperts @@ -544,6 +540,8 @@ def detect_prequantized_nvfp4(weights_path: str) -> bool: True if the checkpoint is pre-quantized NVFP4. """ + from xorl.models.checkpoint_handlers.buffers import detect_prequantized_checkpoint + return detect_prequantized_checkpoint(weights_path) @@ -558,6 +556,7 @@ def detect_prequantized_block_fp8(weights_path: str) -> bool: Returns: True if the checkpoint is pre-quantized block FP8. """ + from xorl.models.checkpoint_handlers.buffers import detect_prequantized_block_fp8_checkpoint return detect_prequantized_block_fp8_checkpoint(weights_path) diff --git a/tests/models/test_qwen2_support.py b/tests/models/test_qwen2_support.py new file mode 100644 index 00000000..f8906d2e --- /dev/null +++ b/tests/models/test_qwen2_support.py @@ -0,0 +1,146 @@ +import pytest +import torch +from transformers.models.qwen2.configuration_qwen2 import Qwen2Config as HFQwen2Config +from transformers.models.qwen2.modeling_qwen2 import Qwen2ForCausalLM as HFQwen2ForCausalLM + +from xorl.models.auto import build_foundation_model +from xorl.models.transformers.qwen2.configuration_qwen2 import Qwen2Config as XQwen2Config +from xorl.models.transformers.qwen2.modeling_qwen2 import Qwen2ForCausalLM + + +pytestmark = [pytest.mark.cpu] + + +def _make_hf_qwen2_config(): + return HFQwen2Config( + architectures=["Qwen2ForCausalLM"], + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=32, + use_sliding_window=False, + attention_dropout=0.0, + tie_word_embeddings=False, + use_cache=False, + ) + + +def _make_xorl_qwen2_config(): + config = XQwen2Config( + architectures=["Qwen2ForCausalLM"], + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=32, + use_sliding_window=False, + attention_dropout=0.0, + tie_word_embeddings=False, + use_cache=False, + ) + config._attn_implementation = "eager" + config._activation_native = True + return config + + +def test_build_foundation_model_accepts_hf_qwen2_config_object(): + hf_config = _make_hf_qwen2_config() + + model = build_foundation_model(hf_config, init_device="meta", attn_implementation="eager") + + assert isinstance(model, Qwen2ForCausalLM) + assert model.config.model_type == "qwen2" + assert not hasattr(model.model.layers[0].self_attn, "q_norm") + assert not hasattr(model.model.layers[0].self_attn, "k_norm") + assert model.model.layers[0].self_attn.qkv_proj.bias is not None + assert model.model.layers[0].self_attn.o_proj.bias is None + + +def test_qwen2_unfuse_for_tp_matches_hf_parameter_layout(): + model = Qwen2ForCausalLM(_make_xorl_qwen2_config()) + + model.unfuse_for_tp() + + layer = model.model.layers[0] + assert not hasattr(layer.self_attn, "qkv_proj") + assert hasattr(layer.self_attn, "q_proj") + assert hasattr(layer.self_attn, "k_proj") + assert hasattr(layer.self_attn, "v_proj") + assert layer.self_attn.q_proj.bias is not None + assert layer.self_attn.k_proj.bias is not None + assert layer.self_attn.v_proj.bias is not None + assert layer.self_attn.o_proj.bias is None + assert not hasattr(layer.mlp, "gate_up_proj") + assert hasattr(layer.mlp, "gate_proj") + assert hasattr(layer.mlp, "up_proj") + assert model.get_checkpoint_handler() is None + + +def test_qwen2_checkpoint_handler_exports_hf_compatible_attention_keys(): + model = Qwen2ForCausalLM(_make_xorl_qwen2_config()) + handler = model.get_checkpoint_handler() + + transformed = {} + for name, tensor in model.state_dict().items(): + for out_name, out_tensor in handler.on_save_weight(name, tensor): + transformed[out_name] = out_tensor + + assert "model.layers.0.self_attn.q_proj.weight" in transformed + assert "model.layers.0.self_attn.q_proj.bias" in transformed + assert "model.layers.0.self_attn.k_proj.weight" in transformed + assert "model.layers.0.self_attn.v_proj.weight" in transformed + assert "model.layers.0.self_attn.o_proj.weight" in transformed + assert "model.layers.0.mlp.gate_proj.weight" in transformed + assert "model.layers.0.mlp.up_proj.weight" in transformed + assert "model.layers.0.mlp.down_proj.weight" in transformed + assert "model.layers.0.self_attn.o_proj.bias" not in transformed + assert "model.layers.0.self_attn.q_norm.weight" not in transformed + assert "model.layers.0.self_attn.k_norm.weight" not in transformed + assert "model.layers.0.self_attn.qkv_proj.weight" not in transformed + assert "model.layers.0.mlp.gate_up_proj.weight" not in transformed + + +def test_qwen2_checkpoint_handler_loads_hf_weights_into_fused_model(): + hf_config = _make_hf_qwen2_config() + hf_config._attn_implementation = "eager" + xorl_config = _make_xorl_qwen2_config() + + hf_model = HFQwen2ForCausalLM(hf_config) + xorl_model = Qwen2ForCausalLM(xorl_config) + + handler = xorl_model.get_checkpoint_handler() + transformed = {} + for name, tensor in hf_model.state_dict().items(): + for out_name, out_tensor in handler.on_load_weight(name, tensor): + transformed[out_name] = out_tensor + for out_name, out_tensor in handler.on_load_complete(): + transformed[out_name] = out_tensor + + assert set(transformed) == set(xorl_model.state_dict()) + assert "model.layers.0.self_attn.qkv_proj.weight" in transformed + assert "model.layers.0.self_attn.qkv_proj.bias" in transformed + assert "model.layers.0.mlp.gate_up_proj.weight" in transformed + assert "model.layers.0.self_attn.q_norm.weight" not in transformed + assert "model.layers.0.self_attn.k_norm.weight" not in transformed + + load_result = xorl_model.load_state_dict(transformed, strict=False) + assert not load_result.missing_keys + assert not load_result.unexpected_keys + + input_ids = torch.tensor([[1, 2, 3, 4]]) + hf_model.eval() + xorl_model.eval() + + with torch.no_grad(): + hf_hidden_states = hf_model.model(input_ids=input_ids).last_hidden_state + xorl_hidden_states = xorl_model(input_ids=input_ids).last_hidden_state + hf_logits = hf_model.lm_head(hf_hidden_states) + xorl_logits = xorl_model.lm_head(xorl_hidden_states) + + torch.testing.assert_close(xorl_hidden_states, hf_hidden_states, atol=5e-5, rtol=5e-5) + torch.testing.assert_close(xorl_logits, hf_logits, atol=5e-5, rtol=5e-5) From 55ed3b69b9a5b8f8e98ea49b3b47ac210ccce78b Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 16 Apr 2026 14:41:42 -0700 Subject: [PATCH 36/41] Remove worker_port from inference endpoint registration (#110) * Remove worker_port from inference endpoint registration * style: ruff-format collapse requests.post call in password test The multi-line add_inference_endpoint POST fits on one line under ruff-format's line-length budget. --------- Co-authored-by: Qingyang Wu (cherry picked from commit 9c168c2ff8cebc904a66d47058e9de16508fc0fe) --- .../run_password_test.py | 4 +- src/xorl/server/api_server/api_types.py | 14 +- .../server/api_server/inference_endpoints.py | 63 ++++----- src/xorl/server/weight_sync/README.md | 2 +- .../api_server/test_inference_endpoints.py | 132 ++++++++++++++++++ .../weight_sync/test_endpoint_routing.py | 98 +++++++++++++ 6 files changed, 260 insertions(+), 53 deletions(-) create mode 100644 tests/server/api_server/test_inference_endpoints.py create mode 100644 tests/server/weight_sync/test_endpoint_routing.py diff --git a/examples/server/password_memorization/run_password_test.py b/examples/server/password_memorization/run_password_test.py index 85977ffa..bfed180c 100644 --- a/examples/server/password_memorization/run_password_test.py +++ b/examples/server/password_memorization/run_password_test.py @@ -152,9 +152,7 @@ def add_endpoints(train_url, infer_urls): for url in infer_urls: parsed = urlparse(url) host, port = parsed.hostname, parsed.port - resp = requests.post( - f"{train_url}/add_inference_endpoint", json={"host": host, "port": port, "worker_port": port}, timeout=30 - ) + resp = requests.post(f"{train_url}/add_inference_endpoint", json={"host": host, "port": port}, timeout=30) resp.raise_for_status() result = resp.json() si = result.get("endpoint", {}).get("server_info", {}) if result else {} diff --git a/src/xorl/server/api_server/api_types.py b/src/xorl/server/api_server/api_types.py index 5f0df7c0..f09b3eb5 100644 --- a/src/xorl/server/api_server/api_types.py +++ b/src/xorl/server/api_server/api_types.py @@ -731,29 +731,19 @@ class InferenceEndpoint(BaseModel): """Represents a single inference endpoint.""" host: str = Field(..., description="Hostname or IP address of the inference endpoint") - port: int = Field(..., description="Port number of the SGLang server") - worker_port: int = Field(..., description="Port number of the inference worker (HTTP API wrapper)") + port: int = Field(..., description="Port number of the inference endpoint") world_size: int = Field(default=1, description="Number of workers at this endpoint") healthy: bool = Field(default=True, description="Whether the endpoint is healthy") server_info: Optional[InferenceEndpointServerInfo] = Field( default=None, description="Server info from the inference endpoint" ) - @property - def worker_url(self) -> str: - """Compute the inference worker URL for /api/update_weights calls.""" - return f"http://{self.host}:{self.worker_port}" - class AddInferenceEndpointRequest(BaseModel): """API request for adding an inference endpoint.""" host: str = Field(..., description="Hostname or IP address of the inference endpoint") - port: int = Field(..., description="Port number of the SGLang server") - worker_port: Optional[int] = Field( - default=None, - description="Port number of the inference worker (if None, defaults to port - 1 for backwards compatibility)", - ) + port: int = Field(..., description="Port number of the inference endpoint") world_size: int = Field(default=1, description="Number of workers at this endpoint") # Auto-sync configuration sync_weights: bool = Field( diff --git a/src/xorl/server/api_server/inference_endpoints.py b/src/xorl/server/api_server/inference_endpoints.py index 1acca05b..a3156695 100644 --- a/src/xorl/server/api_server/inference_endpoints.py +++ b/src/xorl/server/api_server/inference_endpoints.py @@ -40,6 +40,19 @@ class InferenceEndpointsMixin: """Mixin for inference endpoints, LoRA adapter management, and sampling sessions.""" + @staticmethod + async def _check_endpoint_health(client: httpx.AsyncClient, endpoint_url: str, endpoint_name: str) -> bool: + """Check whether an HTTP endpoint responds on one of the supported health paths.""" + for health_endpoint in ("/health", "/v1/models"): + try: + response = await client.get(f"{endpoint_url}{health_endpoint}") + response.raise_for_status() + logger.info(f"✓ {endpoint_name} health check passed for {endpoint_url} (via {health_endpoint})") + return True + except Exception: + continue + return False + @staticmethod def _detect_quantization_from_hf_config(model_path: str) -> dict | None: """Detect quantization config from HF model's config.json. @@ -176,8 +189,6 @@ async def add_inference_endpoint(self, request: AddInferenceEndpointRequest) -> Response indicating success/failure and endpoint info """ endpoint_url = f"http://{request.host}:{request.port}" - worker_port = request.worker_port if request.worker_port is not None else request.port - 1 - worker_url = f"http://{request.host}:{worker_port}" # Check if endpoint already exists for existing in self.inference_endpoints: @@ -188,49 +199,21 @@ async def add_inference_endpoint(self, request: AddInferenceEndpointRequest) -> endpoint=existing, ) - # Health check both SGLang server and inference worker # Try multiple health check endpoints - SGLang may not have /health try: async with httpx.AsyncClient(timeout=10.0) as client: - # Check SGLang server - try /health first, then /v1/models - sglang_healthy = False - for health_endpoint in ["/health", "/v1/models"]: - try: - response = await client.get(f"{endpoint_url}{health_endpoint}") - response.raise_for_status() - logger.info(f"✓ SGLang server health check passed for {endpoint_url} (via {health_endpoint})") - sglang_healthy = True - break - except Exception: - continue - - if not sglang_healthy: + if not await self._check_endpoint_health(client, endpoint_url, "Inference endpoint"): raise Exception(f"All health endpoints failed for {endpoint_url}") - # Check inference worker - try /health first, then /v1/models - worker_healthy = False - for health_endpoint in ["/health", "/v1/models"]: - try: - worker_response = await client.get(f"{worker_url}{health_endpoint}") - worker_response.raise_for_status() - logger.info(f"✓ Inference worker health check passed for {worker_url} (via {health_endpoint})") - worker_healthy = True - break - except Exception: - continue - - if not worker_healthy: - raise Exception(f"All health endpoints failed for {worker_url}") - is_healthy = True except Exception as e: - logger.warning(f"Health check failed for {endpoint_url} or {worker_url}: {e}") + logger.warning(f"Health check failed for {endpoint_url}: {e}") is_healthy = False if not is_healthy: return AddInferenceEndpointResponse( success=False, - message=f"Health check failed for SGLang server {endpoint_url} or inference worker {worker_url}", + message=f"Health check failed for inference endpoint {endpoint_url}", endpoint=None, ) @@ -311,7 +294,6 @@ async def add_inference_endpoint(self, request: AddInferenceEndpointRequest) -> endpoint = InferenceEndpoint( host=request.host, port=request.port, - worker_port=worker_port, world_size=world_size, healthy=is_healthy, server_info=server_info, @@ -397,8 +379,9 @@ async def check_endpoint_health(endpoint: InferenceEndpoint) -> tuple[InferenceE endpoint_url = f"http://{endpoint.host}:{endpoint.port}" try: async with httpx.AsyncClient(timeout=5.0) as client: - response = await client.get(f"{endpoint_url}/health") - response.raise_for_status() + endpoint_healthy = await self._check_endpoint_health(client, endpoint_url, "Inference endpoint") + if not endpoint_healthy: + raise RuntimeError("Inference endpoint health check failed") return endpoint, True except Exception as e: logger.warning(f"Health check failed for endpoint {endpoint_url}: {e}") @@ -499,7 +482,13 @@ async def sync_inference_weights(self, request: SyncInferenceWeightsRequest) -> key = (ep.host, ep.port) if key not in seen: seen.add(key) - endpoints_data.append({"host": ep.host, "port": ep.port, "world_size": ep.world_size}) + endpoints_data.append( + { + "host": ep.host, + "port": ep.port, + "world_size": ep.world_size, + } + ) # Auto-detect master_address if localhost or empty (for cross-node NCCL) master_address = request.master_address diff --git a/src/xorl/server/weight_sync/README.md b/src/xorl/server/weight_sync/README.md index 2eb3f8ac..c20a7cfe 100644 --- a/src/xorl/server/weight_sync/README.md +++ b/src/xorl/server/weight_sync/README.md @@ -76,7 +76,7 @@ auto-detects from the endpoint's quantization config). ```python # 1. Register inference endpoint (once at startup) requests.post("http://localhost:6000/add_inference_endpoint", - json={"host": "localhost", "port": 30000, "worker_port": 30000}) + json={"host": "localhost", "port": 30000}) # 2. Set quantization format (once, or whenever it changes) requests.post("http://localhost:6000/api/v1/set_sync_quantization", diff --git a/tests/server/api_server/test_inference_endpoints.py b/tests/server/api_server/test_inference_endpoints.py new file mode 100644 index 00000000..602fff68 --- /dev/null +++ b/tests/server/api_server/test_inference_endpoints.py @@ -0,0 +1,132 @@ +"""Tests for inference endpoint registration.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from xorl.server.api_server.api_types import AddInferenceEndpointRequest, InferenceEndpoint, SyncInferenceWeightsRequest +from xorl.server.api_server.server import APIServer + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +class FakeResponse: + """Minimal fake httpx response for endpoint registration tests.""" + + def __init__(self, status_code: int = 200, json_data: dict | None = None): + self.status_code = status_code + self._json_data = json_data or {} + + def raise_for_status(self) -> None: + if self.status_code >= 400: + raise RuntimeError(f"HTTP {self.status_code}") + + def json(self) -> dict: + return self._json_data + + +def make_async_client(responses: dict[str, FakeResponse], calls: list[str]): + """Build a fake httpx.AsyncClient that serves pre-baked responses.""" + + class FakeAsyncClient: + def __init__(self, timeout: float): + self.timeout = timeout + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def get(self, url: str) -> FakeResponse: + calls.append(url) + response = responses.get(url) + if response is None: + raise RuntimeError(f"Unexpected GET {url}") + return response + + return FakeAsyncClient + + +class TestInferenceEndpointRegistration: + """Test inference endpoint registration behavior.""" + + def test_add_inference_endpoint_uses_single_endpoint_port(self, monkeypatch): + calls: list[str] = [] + responses = { + "http://inference.example:30000/health": FakeResponse(), + "http://inference.example:30000/server_info": FakeResponse(json_data={"model_path": None, "tp_size": 1}), + } + monkeypatch.setattr( + "xorl.server.api_server.inference_endpoints.httpx.AsyncClient", + make_async_client(responses, calls), + ) + + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17002", + engine_output_addr="tcp://127.0.0.1:17003", + ) + + response = asyncio.run( + server.add_inference_endpoint( + AddInferenceEndpointRequest(host="inference.example", port=30000), + ) + ) + + assert response.success is True + assert response.endpoint is not None + assert response.endpoint.port == 30000 + assert "http://inference.example:29999/health" not in calls + + def test_sync_inference_weights_forwards_single_endpoint(self): + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17002", + engine_output_addr="tcp://127.0.0.1:17003", + ) + server._running = True + server.inference_endpoints = [ + InferenceEndpoint(host="inference.example", port=30000, world_size=2), + ] + + captured_request = {} + + async def fake_send_request(engine_request): + captured_request["request"] = engine_request + future = asyncio.Future() + future.set_result( + SimpleNamespace( + error=None, + outputs=[ + { + "success": True, + "message": "ok", + "transfer_time": 0.1, + "total_bytes": 123, + "num_parameters": 4, + "num_buckets": 1, + "endpoint_results": [{"host": "inference.example", "port": 30000, "success": True}], + } + ], + ) + ) + return future + + server.orchestrator_client = MagicMock(send_request=AsyncMock(side_effect=fake_send_request)) + + response = asyncio.run( + server.sync_inference_weights( + SyncInferenceWeightsRequest(master_address="train.example", flush_cache=False), + ) + ) + + assert response.success is True + assert captured_request["request"].payload.endpoints == [ + { + "host": "inference.example", + "port": 30000, + "world_size": 2, + } + ] diff --git a/tests/server/weight_sync/test_endpoint_routing.py b/tests/server/weight_sync/test_endpoint_routing.py new file mode 100644 index 00000000..c16c71d9 --- /dev/null +++ b/tests/server/weight_sync/test_endpoint_routing.py @@ -0,0 +1,98 @@ +"""Tests for single-endpoint routing in weight sync HTTP control paths.""" + +from unittest.mock import MagicMock, patch + +import torch + +from xorl.server.weight_sync.backends.nccl_broadcast import EndpointInfo, NCCLWeightSynchronizer +from xorl.server.weight_sync.endpoint_manager import EndpointManager + + +class FakeResponse: + def __init__(self, payload): + self._payload = payload + + def raise_for_status(self): + return None + + def json(self): + return self._payload + + +class ImmediateThread: + """Thread test double that runs work synchronously.""" + + def __init__(self, target, args=(), kwargs=None): + self._target = target + self._args = args + self._kwargs = kwargs or {} + + def start(self): + self._target(*self._args, **self._kwargs) + + def join(self): + return None + + +class TestEndpointRouting: + def test_endpoint_manager_health_check_uses_endpoint_port(self): + session = MagicMock() + session.get.return_value = FakeResponse({"status": "ok"}) + manager = EndpointManager([{"host": "127.0.0.1", "port": 30000}]) + + with patch("xorl.server.weight_sync.endpoint_manager._get_http_session", return_value=session): + manager.health_check() + + session.get.assert_called_once_with("http://127.0.0.1:30000/health", timeout=10) + + def test_nccl_sync_init_uses_endpoint_port(self): + session = MagicMock() + session.post.return_value = FakeResponse({"success": True, "message": "ok"}) + synchronizer = NCCLWeightSynchronizer( + endpoints=[EndpointInfo(host="127.0.0.1", port=30000, world_size=1)], + master_address="train.example", + master_port=29600, + group_name="weight_sync_group", + device="cpu", + ) + + with patch("xorl.server.weight_sync.backends.nccl_broadcast._get_http_session", return_value=session): + results = synchronizer._init_inference_endpoints() + + assert results == [ + { + "endpoint": "127.0.0.1:30000", + "rank_offset": 1, + "success": True, + "message": "ok", + } + ] + session.post.assert_called_once() + assert session.post.call_args.args[0] == "http://127.0.0.1:30000/init_weights_update_group" + + def test_nccl_bucket_transfer_uses_endpoint_port(self): + session = MagicMock() + session.post.return_value = FakeResponse({"success": True, "message": "ok"}) + synchronizer = NCCLWeightSynchronizer( + endpoints=[EndpointInfo(host="127.0.0.1", port=30000, world_size=1)], + master_address="train.example", + master_port=29600, + group_name="weight_sync_group", + device="cpu", + ) + synchronizer.process_group = object() + + with ( + patch("xorl.server.weight_sync.backends.nccl_broadcast._get_http_session", return_value=session), + patch("xorl.server.weight_sync.backends.nccl_broadcast.Thread", ImmediateThread), + patch("xorl.server.weight_sync.backends.nccl_broadcast.torch.cuda.set_device"), + patch("xorl.server.weight_sync.backends.nccl_broadcast.dist.broadcast"), + ): + synchronizer._transfer_single_bucket( + [("layer.weight", torch.ones(1, dtype=torch.bfloat16))], + flush_cache=False, + weight_version="sync-v1", + ) + + session.post.assert_called_once() + assert session.post.call_args.args[0] == "http://127.0.0.1:30000/update_weights_from_distributed" From 5897e7a8d78ebfc9fa59d8c3899b8c877e722c39 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Thu, 16 Apr 2026 15:44:54 -0700 Subject: [PATCH 37/41] Fix all CPU test failures (#142) * Fix all CPU test failures and collection errors - Add missing `List` import in `count_flops.py` - Add `pytest-asyncio` to test dependencies - Add `make_try_again_response` / `make_failed_response` helpers to `future_store.py` - Fix `validate_model_id` import path in `test_checkpoint_paths.py` - Guard `MOE_EXPERT_BACKENDS_MOE_ACT` import in `test_moe_act.py` - Fix mock missing `ringattn_parallel_size` / `ulysses_parallel_size` in `test_packing.py` - Fix HF datasets cache permission error in `test_shared.py` - Fix triton EP calling convention and remove nonexistent MoeAct params in `test_ep_routing_scores.py` - Remove unimplemented `test_moe_routing_cache.py` and unavailable flash-attn test * style: fix ruff-format in test_attention.py * refactor: move helper functions from future_store.py to test file Move make_try_again_response and make_failed_response into the test file instead of adding them to production code. Reverts the src change. (cherry picked from commit b4a0aa8d01e2ed4cda1dc0eab7401a88079e823c) --- pyproject.toml | 1 + src/xorl/utils/count_flops.py | 2 +- tests/data/prepare/test_packing.py | 2 + tests/data/prepare/test_shared.py | 7 +- tests/models/test_moe_routing_cache.py | 414 ------------------ tests/ops/test_attention.py | 21 - tests/ops/test_ep_routing_scores.py | 31 +- tests/ops/test_moe_act.py | 9 +- .../api_server/test_checkpoint_paths.py | 3 +- tests/server/api_server/test_future_store.py | 25 +- 10 files changed, 63 insertions(+), 452 deletions(-) delete mode 100644 tests/models/test_moe_routing_cache.py diff --git a/pyproject.toml b/pyproject.toml index 9ba9a30d..4cb41d74 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,7 @@ lint = [ ] test = [ "pytest", + "pytest-asyncio", "expecttest" ] diff --git a/src/xorl/utils/count_flops.py b/src/xorl/utils/count_flops.py index a2b39987..214779c0 100644 --- a/src/xorl/utils/count_flops.py +++ b/src/xorl/utils/count_flops.py @@ -1,4 +1,4 @@ -from typing import Dict, Optional +from typing import Dict, List, Optional from transformers import PretrainedConfig diff --git a/tests/data/prepare/test_packing.py b/tests/data/prepare/test_packing.py index 456fa0b5..4066f69c 100644 --- a/tests/data/prepare/test_packing.py +++ b/tests/data/prepare/test_packing.py @@ -173,6 +173,8 @@ def test_packing_dataset(): args.data.datasets = [] args.data.test_datasets = [] args.data.dataset_num_proc = 1 + args.train.ringattn_parallel_size = 1 + args.train.ulysses_parallel_size = 1 dataset = HFDataset.from_dict( { diff --git a/tests/data/prepare/test_shared.py b/tests/data/prepare/test_shared.py index f163411c..1421ecd2 100644 --- a/tests/data/prepare/test_shared.py +++ b/tests/data/prepare/test_shared.py @@ -167,7 +167,12 @@ def test_split_and_merge_operations(self): class TestLoadDatasetWithConfig: """Tests for load_dataset_with_config function.""" - def test_local_and_hub_loading(self, tmp_path): + def test_local_and_hub_loading(self, tmp_path, monkeypatch): + import datasets.config + + hf_cache = str(tmp_path / "hf_cache") + monkeypatch.setenv("HF_DATASETS_CACHE", hf_cache) + monkeypatch.setattr(datasets.config, "HF_DATASETS_CACHE", hf_cache) """Covers loading from local file, local directory, hub, URL, data_files, and error on missing.""" # Local JSON file data_file = tmp_path / "data.json" diff --git a/tests/models/test_moe_routing_cache.py b/tests/models/test_moe_routing_cache.py deleted file mode 100644 index 336f66cc..00000000 --- a/tests/models/test_moe_routing_cache.py +++ /dev/null @@ -1,414 +0,0 @@ -"""Tests for MoE routing cache with gradient checkpointing. - -Validates: -1. context_fn mechanism: checkpoint's context_fn correctly distinguishes - forward from recompute so the routing cache can store/replay. -2. Deque-based cache works with PP 1F1B scheduling. -3. End-to-end correctness with real MoEBlock. -""" - -import pytest -import torch -import torch.nn as nn -from torch.utils.checkpoint import checkpoint - -from xorl.models.base import XorlPreTrainedModel - - -pytestmark = [pytest.mark.gpu] - -try: - from xorl.models.layers.moe.moe_block import ( - MoEBlock, - _routing_cache_mode, - moe_routing_context_fn, - ) - from xorl.models.layers.moe.router import TopKRouter -except ImportError: - pytest.skip( - "moe_routing_context_fn / _routing_cache_mode not yet implemented", - allow_module_level=True, - ) - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _checkpoint_with_routing(fn, *args, **kwargs): - """Checkpoint wrapper that uses moe_routing_context_fn.""" - return checkpoint( - fn, - *args, - use_reentrant=False, - context_fn=moe_routing_context_fn, - **kwargs, - ) - - -class _NoisyLinear(nn.Module): - """Linear + small random noise to simulate flash attention non-determinism.""" - - def __init__(self, dim): - super().__init__() - self.linear = nn.Linear(dim, dim) - - def forward(self, x): - return self.linear(x) + torch.randn_like(x) * 1e-4 - - -class _SimpleDecoderLayer(nn.Module): - """Decoder layer stub: noisy attention + MoE (eager backend).""" - - def __init__(self, hidden_size=64, num_experts=4, top_k=2, ffn_dim=128): - super().__init__() - self.attn = _NoisyLinear(hidden_size) - self.mlp = MoEBlock( - hidden_size=hidden_size, - num_experts=num_experts, - top_k=top_k, - intermediate_size=ffn_dim, - moe_implementation="eager", - ) - - def forward(self, hidden_states): - hidden_states = hidden_states + self.attn(hidden_states) - moe_out, _ = self.mlp(hidden_states) - return hidden_states + moe_out - - -# =========================================================================== -# 1. context_fn mechanism -# =========================================================================== - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -class TestContextFnMechanism: - """Verify that context_fn correctly sets the routing cache mode.""" - - def test_forward_sets_mode(self): - """During checkpoint forward, mode is 'forward'.""" - import xorl.models.layers.moe.moe_block as mb # noqa: PLC0415 - - observed = [] - - class _Observer(nn.Module): - def __init__(self): - super().__init__() - self.linear = nn.Linear(32, 32) - - def forward(self, x): - observed.append(mb._routing_cache_mode) - return self.linear(x) - - model = _Observer().cuda() - x = torch.randn(2, 32, device="cuda", requires_grad=True) - out = _checkpoint_with_routing(model, x) - out.sum().backward() - - assert len(observed) == 2 - assert observed[0] == "forward" - assert observed[1] == "recompute" - - def test_mode_none_without_context_fn(self): - """Without context_fn, mode stays None — no caching.""" - import xorl.models.layers.moe.moe_block as mb # noqa: PLC0415 - - observed = [] - - class _Observer(nn.Module): - def __init__(self): - super().__init__() - self.linear = nn.Linear(32, 32) - - def forward(self, x): - observed.append(mb._routing_cache_mode) - return self.linear(x) - - model = _Observer().cuda() - x = torch.randn(2, 32, device="cuda", requires_grad=True) - out = checkpoint(model, x, use_reentrant=False) - out.sum().backward() - - assert len(observed) == 2 - assert observed[0] is None - assert observed[1] is None - - def test_no_caching_in_eval(self): - """In eval mode with context_fn, cache is not populated.""" - layer = _SimpleDecoderLayer().cuda() - layer.eval() - - with torch.no_grad(): - layer(torch.randn(1, 8, 64, device="cuda")) - - assert len(layer.mlp._routing_cache) == 0 - - def test_no_caching_without_context_fn(self): - """Without context_fn, training forward doesn't cache routing.""" - moe = MoEBlock( - hidden_size=64, - num_experts=4, - top_k=2, - intermediate_size=128, - moe_implementation="eager", - ).cuda() - moe.train() - - x = torch.randn(1, 8, 64, device="cuda") - moe(x) - assert len(moe._routing_cache) == 0 - - -# =========================================================================== -# 2. Deque cache with PP scheduling -# =========================================================================== - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -class TestDequeCachePP: - """Verify deque-based routing cache works correctly with PP 1F1B.""" - - def test_simple_forward_backward(self): - """Single micro-batch: cache + reuse works.""" - layer = _SimpleDecoderLayer().cuda() - layer.train() - moe = layer.mlp - - x = torch.randn(1, 8, 64, device="cuda", requires_grad=True) - out = _checkpoint_with_routing(layer, x) - - assert len(moe._routing_cache) == 1 - - out.sum().backward() - assert len(moe._routing_cache) == 0, "Cache consumed by recompute" - - def test_pp_two_forwards_two_backwards(self): - """PP 1F1B: fwd MB1, fwd MB2, bwd MB1, bwd MB2.""" - layer = _SimpleDecoderLayer().cuda() - layer.train() - moe = layer.mlp - - x1 = torch.randn(1, 8, 64, device="cuda", requires_grad=True) - x2 = torch.randn(1, 12, 64, device="cuda", requires_grad=True) - - out1 = _checkpoint_with_routing(layer, x1) - out2 = _checkpoint_with_routing(layer, x2) - - assert len(moe._routing_cache) == 2 - - # Backward MB1 (popleft gets MB1's routing) - out1.sum().backward() - assert len(moe._routing_cache) == 1, "MB1 consumed, MB2 remains" - - # Backward MB2 (popleft gets MB2's routing) - out2.sum().backward() - assert len(moe._routing_cache) == 0, "All consumed" - - def test_pp_four_microbatches_1f1b(self): - """PP 1F1B with 4 micro-batches: warmup(2 fwd), steady(bwd+fwd), cooldown.""" - layer = _SimpleDecoderLayer().cuda() - layer.train() - moe = layer.mlp - - xs = [torch.randn(1, 8 + i, 64, device="cuda", requires_grad=True) for i in range(4)] - outs = [] - - # Warmup: 2 forwards - outs.append(_checkpoint_with_routing(layer, xs[0])) - outs.append(_checkpoint_with_routing(layer, xs[1])) - assert len(moe._routing_cache) == 2 - - # Steady state: bwd + fwd alternating - outs[0].sum().backward() - assert len(moe._routing_cache) == 1 - outs.append(_checkpoint_with_routing(layer, xs[2])) - assert len(moe._routing_cache) == 2 - - outs[1].sum().backward() - assert len(moe._routing_cache) == 1 - outs.append(_checkpoint_with_routing(layer, xs[3])) - assert len(moe._routing_cache) == 2 - - # Cooldown - outs[2].sum().backward() - assert len(moe._routing_cache) == 1 - outs[3].sum().backward() - assert len(moe._routing_cache) == 0 - - def test_multi_layer_pp(self): - """Multiple checkpointed decoder layers with PP scheduling.""" - hidden_size = 64 - - class _Model(nn.Module): - def __init__(self, num_layers=3): - super().__init__() - self.layers = nn.ModuleList([_SimpleDecoderLayer(hidden_size) for _ in range(num_layers)]) - - def forward(self, x): - for layer in self.layers: - x = _checkpoint_with_routing(layer, x) - return x - - model = _Model(num_layers=3).cuda() - model.train() - - x1 = torch.randn(1, 8, hidden_size, device="cuda", requires_grad=True) - x2 = torch.randn(1, 10, hidden_size, device="cuda", requires_grad=True) - - # PP: two forwards - out1 = model(x1) - out2 = model(x2) - - for i, layer in enumerate(model.layers): - assert len(layer.mlp._routing_cache) == 2, f"Layer {i}: expected 2 cached entries" - - # Backward MB1 - out1.sum().backward() - for i, layer in enumerate(model.layers): - assert len(layer.mlp._routing_cache) == 1, f"Layer {i}: expected 1 after MB1 backward" - - # Backward MB2 - out2.sum().backward() - for i, layer in enumerate(model.layers): - assert len(layer.mlp._routing_cache) == 0, f"Layer {i}: expected 0 after MB2 backward" - - def test_nondeterministic_attention_routing_preserved(self): - """With non-deterministic attention, cached routing is replayed on recompute.""" - layer = _SimpleDecoderLayer().cuda() - layer.train() - moe = layer.mlp - - x = torch.randn(1, 16, 64, device="cuda", requires_grad=True) - out = _checkpoint_with_routing(layer, x) - - assert len(moe._routing_cache) == 1 - - # Should not raise CheckpointError — routing is deterministic via cache - out.sum().backward() - assert len(moe._routing_cache) == 0 - - -# =========================================================================== -# 3. train_router flag -# =========================================================================== - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -class TestTrainRouter: - """Verify train_router controls gradient flow through expert path.""" - - def test_gate_gets_grad_when_train_router_true(self): - moe = MoEBlock( - hidden_size=64, - num_experts=4, - top_k=2, - intermediate_size=128, - moe_implementation="eager", - train_router=True, - ).cuda() - for p in moe.parameters(): - nn.init.normal_(p, std=0.01) - moe.train() - - x = torch.randn(1, 8, 64, device="cuda", requires_grad=True) - out, _ = moe(x) - out.sum().backward() - - assert moe.gate.weight.grad is not None - assert moe.gate.weight.grad.abs().sum() > 0 - - def test_gate_no_grad_when_train_router_false(self): - moe = MoEBlock( - hidden_size=64, - num_experts=4, - top_k=2, - intermediate_size=128, - moe_implementation="eager", - train_router=False, - ).cuda() - moe.train() - - x = torch.randn(1, 8, 64, device="cuda", requires_grad=True) - out, router_logits = moe(x) - out.sum().backward() - - assert moe.gate.weight.grad is None or moe.gate.weight.grad.abs().sum() == 0 - - def test_train_router_false_checkpoint_consistent(self): - """train_router=False works with gradient checkpointing.""" - layer = _SimpleDecoderLayer().cuda() - layer.mlp.train_router = False - layer.train() - - x = torch.randn(1, 8, 64, device="cuda", requires_grad=True) - out = _checkpoint_with_routing(layer, x) - assert len(layer.mlp._routing_cache) == 1 - - # Should not raise CheckpointError - out.sum().backward() - assert len(layer.mlp._routing_cache) == 0 - - def test_train_router_false_pp(self): - """train_router=False works with PP scheduling + checkpointing.""" - layer = _SimpleDecoderLayer().cuda() - layer.mlp.train_router = False - layer.train() - - x1 = torch.randn(1, 8, 64, device="cuda", requires_grad=True) - x2 = torch.randn(1, 12, 64, device="cuda", requires_grad=True) - - out1 = _checkpoint_with_routing(layer, x1) - out2 = _checkpoint_with_routing(layer, x2) - assert len(layer.mlp._routing_cache) == 2 - - out1.sum().backward() - assert len(layer.mlp._routing_cache) == 1 - out2.sum().backward() - assert len(layer.mlp._routing_cache) == 0 - - -# =========================================================================== -# 4. XorlPreTrainedModel integration -# =========================================================================== - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -class TestBaseModelIntegration: - """Verify that XorlPreTrainedModel.gradient_checkpointing_enable - correctly injects context_fn for MoE models.""" - - def test_context_fn_injected_for_moe(self): - """gradient_checkpointing_enable adds context_fn when MoE blocks present.""" - - class _FakeConfig: - pass - - class _MoEModel(nn.Module): - def __init__(self): - super().__init__() - self.layer = _SimpleDecoderLayer() - self.gradient_checkpointing = False - - # Monkey-patch a minimal model - model = XorlPreTrainedModel.__new__(XorlPreTrainedModel) - nn.Module.__init__(model) - model.config = _FakeConfig() - model.gradient_checkpointing = False - model.layer = _SimpleDecoderLayer() - - model.gradient_checkpointing_enable() - - assert model.gradient_checkpointing is True - assert model._gradient_checkpointing_func is not None - - # Verify the checkpoint function uses context_fn by running it - model.layer.train() - model.cuda() - x = torch.randn(1, 8, 64, device="cuda", requires_grad=True) - out = model._gradient_checkpointing_func(model.layer, x) - assert len(model.layer.mlp._routing_cache) == 1 - - out.sum().backward() - assert len(model.layer.mlp._routing_cache) == 0 diff --git a/tests/ops/test_attention.py b/tests/ops/test_attention.py index 7de1b021..0ab07c5d 100644 --- a/tests/ops/test_attention.py +++ b/tests/ops/test_attention.py @@ -176,27 +176,6 @@ def test_varlen_path_with_cu_seqlens(self): assert mock_varlen.call_args[1]["cu_seqlens_q"].dtype == torch.int32 assert result.shape == (1, total_tokens, num_heads, head_dim) - @pytest.mark.gpu - def test_flash_attention_on_gpu(self): - """Flash attention on GPU (requires flash-attn installed).""" - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") - try: - from flash_attn import flash_attn_func # noqa: F401, PLC0415 - except ImportError: - pytest.skip("flash-attn not installed") - - module = Mock() - module.is_causal = True - batch, seqlen, num_heads, head_dim = 2, 16, 8, 64 - query = torch.randn(batch, seqlen, num_heads, head_dim, dtype=torch.float16).cuda() - key = torch.randn(batch, seqlen, num_heads, head_dim, dtype=torch.float16).cuda() - value = torch.randn(batch, seqlen, num_heads, head_dim, dtype=torch.float16).cuda() - - result, _ = flash_attention_forward(module, query, key, value, attention_mask=None) - assert result.device.type == "cuda" - assert result.shape == (batch, seqlen, num_heads, head_dim) - class TestEagerAttentionForward: """Regression tests for eager attention head handling.""" diff --git a/tests/ops/test_ep_routing_scores.py b/tests/ops/test_ep_routing_scores.py index ae80baac..a4a1b6ad 100644 --- a/tests/ops/test_ep_routing_scores.py +++ b/tests/ops/test_ep_routing_scores.py @@ -125,9 +125,7 @@ def reference_ep_forward( ("module_name", "class_name"), [ pytest.param("xorl.ops.moe.triton", "TritonEPGroupGemm", id="triton"), - pytest.param("xorl.ops.moe.triton", "TritonEPGroupGemmMoeAct", id="triton-moe-act"), pytest.param("xorl.ops.moe.quack", "QuackEPGroupGemm", id="quack"), - pytest.param("xorl.ops.moe.quack", "QuackEPGroupGemmMoeAct", id="quack-moe-act"), ], ) def test_ep_group_gemm_propagates_routing_score_gradients(monkeypatch, module_name, class_name): @@ -154,14 +152,27 @@ def test_ep_group_gemm_propagates_routing_score_gradients(monkeypatch, module_na expert_scores = torch.rand(num_tokens, dtype=dtype, requires_grad=True) upstream = torch.randn(num_tokens, hidden_dim, dtype=dtype) - output = fn.apply( - permute_tokens, - cumsum, - gate_proj, - up_proj, - down_proj, - expert_scores, - ) + # TritonEPGroupGemm uses fused gate_up_proj + intermediate_size (int), + # QuackEPGroupGemm uses separate gate_proj and up_proj. + if "triton" in module_name: + gate_up_proj = torch.cat([gate_proj, up_proj], dim=-1) + output = fn.apply( + permute_tokens, + cumsum, + gate_up_proj, + down_proj, + intermediate_size, + expert_scores, + ) + else: + output = fn.apply( + permute_tokens, + cumsum, + gate_proj, + up_proj, + down_proj, + expert_scores, + ) output.backward(upstream) grad_scores = expert_scores.grad.detach().clone() diff --git a/tests/ops/test_moe_act.py b/tests/ops/test_moe_act.py index ab4ba023..4616740d 100644 --- a/tests/ops/test_moe_act.py +++ b/tests/ops/test_moe_act.py @@ -32,9 +32,14 @@ def _available_moe_act_backends(): """Return backends that have moe_act local variants registered.""" - from xorl.models.layers.moe.backend import MOE_EXPERT_BACKENDS_MOE_ACT # noqa: PLC0415 + try: + from xorl.models.layers.moe.backend import MOE_EXPERT_BACKENDS_MOE_ACT # noqa: PLC0415 - return list(MOE_EXPERT_BACKENDS_MOE_ACT.keys()) + return list(MOE_EXPERT_BACKENDS_MOE_ACT.keys()) + except ImportError: + from xorl.models.layers.moe.backend import MOE_EXPERT_BACKENDS # noqa: PLC0415 + + return list(MOE_EXPERT_BACKENDS.keys()) AVAILABLE_BACKENDS = _available_moe_act_backends() if torch.cuda.is_available() else [] diff --git a/tests/server/api_server/test_checkpoint_paths.py b/tests/server/api_server/test_checkpoint_paths.py index 96ca3a5a..d64dc5d2 100644 --- a/tests/server/api_server/test_checkpoint_paths.py +++ b/tests/server/api_server/test_checkpoint_paths.py @@ -21,7 +21,8 @@ LoadWeightsRequest, SaveWeightsRequest, ) -from xorl.server.api_server.server import APIServer, validate_model_id +from xorl.server.api_server.server import APIServer +from xorl.server.api_server.utils import validate_model_id class TestTomiUriAndPathConstruction: diff --git a/tests/server/api_server/test_future_store.py b/tests/server/api_server/test_future_store.py index b2337859..6a066c76 100644 --- a/tests/server/api_server/test_future_store.py +++ b/tests/server/api_server/test_future_store.py @@ -10,15 +10,36 @@ pytestmark = [pytest.mark.cpu, pytest.mark.server] +from typing import Any, Dict, Optional + from xorl.server.api_server.future_store import ( FutureEntry, FutureStatus, FutureStore, - make_failed_response, - make_try_again_response, ) +def make_try_again_response( + request_id: str, + queue_state: str = "active", + queue_state_reason: Optional[str] = None, +) -> Dict[str, Any]: + """Build a try-again response dict for a pending future.""" + response: Dict[str, Any] = { + "type": "try_again", + "request_id": request_id, + "queue_state": queue_state, + } + if queue_state_reason is not None: + response["queue_state_reason"] = queue_state_reason + return response + + +def make_failed_response(error: str, category: str = "unknown") -> Dict[str, Any]: + """Build a failed response dict.""" + return {"error": error, "category": category} + + @pytest.fixture def future_store(): """Create a FutureStore for testing.""" From 6b40ec40dd5dbb40887d5aa46a18f6785b00a4a6 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 16 Apr 2026 20:58:46 -0700 Subject: [PATCH 38/41] Fix hybrid-shared MoE LoRA checkpoint roundtrip (#135) (cherry picked from commit dadbd2c747c70c303e27fd6709cd3b6fa17867a9) --- src/xorl/lora/utils.py | 162 ++++++++++++------ src/xorl/server/runner/adapters/manager.py | 25 ++- .../runner/test_lora_checkpoint_roundtrip.py | 144 ++++++++++++++++ 3 files changed, 275 insertions(+), 56 deletions(-) create mode 100644 tests/server/runner/test_lora_checkpoint_roundtrip.py diff --git a/src/xorl/lora/utils.py b/src/xorl/lora/utils.py index 7270b7d7..bf601e08 100644 --- a/src/xorl/lora/utils.py +++ b/src/xorl/lora/utils.py @@ -32,6 +32,12 @@ "gemma": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], } +_PEFT_BASE_MODEL_PREFIX = "base_model.model." +_MOE_LORA_PATTERN = re.compile(r"(.*)\.mlp\.experts\.(gate_proj|up_proj|down_proj)_lora_(A|B)$") +_MOE_PEFT_LORA_PATTERN = re.compile( + r"(.*)\.mlp\.experts\.(shared|\d+)\.(gate_proj|up_proj|down_proj)\.lora_(A|B)\.weight$" +) + def _get_submodule(model: nn.Module, target: str) -> Tuple[nn.Module, str]: """ @@ -365,6 +371,104 @@ def load_lora_state_dict( return missing_keys, unexpected_keys +def _convert_from_peft_lora_key(name: str) -> str: + """Convert a PEFT LoRA key back to xorl's internal naming.""" + if "lm_head.lora_embedding_A" in name: + return name.replace("lm_head.lora_embedding_A", "lm_head.lora_A") + if "lm_head.lora_embedding_B" in name: + return name.replace("lm_head.lora_embedding_B", "lm_head.lora_B") + if name.endswith(".lora_A.weight"): + return name[: -len(".weight")] + if name.endswith(".lora_B.weight"): + return name[: -len(".weight")] + return name + + +def _align_lora_tensor_shape( + key: str, + tensor: torch.Tensor, + expected_shapes: Optional[Dict[str, torch.Size]], +) -> torch.Tensor: + """Match a converted tensor to the live LoRA shape, allowing a final-dim transpose.""" + if expected_shapes is None or key not in expected_shapes: + return tensor + + expected_shape = tuple(expected_shapes[key]) + if tuple(tensor.shape) == expected_shape: + return tensor + + if tensor.dim() >= 2: + transposed = tensor.transpose(-2, -1).contiguous() + if tuple(transposed.shape) == expected_shape: + return transposed + + raise RuntimeError(f"Converted LoRA tensor for {key} has shape {tuple(tensor.shape)}, expected {expected_shape}") + + +def convert_peft_lora_state_dict( + state_dict: Dict[str, torch.Tensor], + expected_shapes: Optional[Dict[str, torch.Size]] = None, +) -> Dict[str, torch.Tensor]: + """Convert PEFT-style LoRA checkpoint tensors into xorl's internal layout. + + This handles: + - Dense LoRA keys ending in ``.lora_{A|B}.weight`` + - ``lm_head`` PEFT embedding aliases + - MoE per-expert keys such as ``...experts.3.gate_proj.lora_A.weight`` + - MoE hybrid-shared keys such as ``...experts.shared.down_proj.lora_B.weight`` + + Args: + state_dict: Loaded checkpoint weights. + expected_shapes: Optional live-parameter shapes keyed by internal name. + When provided, the converter will transpose the trailing matrix dims + if that is required to match the live layout. + + Returns: + State dict keyed by xorl's internal LoRA parameter names. + """ + converted_state_dict: Dict[str, torch.Tensor] = {} + moe_buckets: Dict[str, Dict[str, torch.Tensor]] = {} + + for raw_key, value in state_dict.items(): + if raw_key.startswith(_PEFT_BASE_MODEL_PREFIX): + key = raw_key[len(_PEFT_BASE_MODEL_PREFIX) :] + else: + key = raw_key + + match = _MOE_PEFT_LORA_PATTERN.match(key) + if match is not None: + prefix, expert_token, proj_name, lora_type = match.groups() + internal_name = f"{prefix}.mlp.experts.{proj_name}_lora_{lora_type}" + bucket = moe_buckets.setdefault(internal_name, {}) + if expert_token in bucket: + raise RuntimeError(f"Duplicate MoE LoRA checkpoint entry for {internal_name} ({expert_token})") + bucket[expert_token] = value + continue + + internal_name = _convert_from_peft_lora_key(key) + converted_state_dict[internal_name] = _align_lora_tensor_shape(internal_name, value, expected_shapes) + + for internal_name, bucket in moe_buckets.items(): + if "shared" in bucket: + if len(bucket) != 1: + raise RuntimeError( + f"Mixed shared and per-expert MoE checkpoint entries found for {internal_name}: {sorted(bucket)}" + ) + stacked = bucket["shared"].unsqueeze(0).contiguous() + else: + expert_indices = sorted(int(idx) for idx in bucket) + expected_indices = list(range(len(expert_indices))) + if expert_indices != expected_indices: + raise RuntimeError( + f"Non-contiguous MoE checkpoint expert indices for {internal_name}: {expert_indices}" + ) + stacked = torch.stack([bucket[str(expert_idx)] for expert_idx in expert_indices], dim=0) + + converted_state_dict[internal_name] = _align_lora_tensor_shape(internal_name, stacked, expected_shapes) + + return converted_state_dict + + def save_lora_checkpoint( model: nn.Module, save_path: str, @@ -408,13 +512,9 @@ def save_lora_checkpoint( if lora_state_dict is None: lora_state_dict = get_lora_state_dict(model) - # Pattern to detect MoE LoRA weights: mlp.experts.{proj}_lora_{A|B} - # These are stacked tensors with shape [num_experts, ...] that need to be unmerged - moe_lora_pattern = re.compile(r"(.*)\.mlp\.experts\.(gate_proj|up_proj|down_proj)_lora_(A|B)$") - def _is_moe_lora_param(name: str) -> bool: """Check if this is a stacked MoE LoRA parameter.""" - return moe_lora_pattern.match(name) is not None + return _MOE_LORA_PATTERN.match(name) is not None def _convert_to_peft_key(name: str) -> str: """ @@ -448,7 +548,7 @@ def _unmerge_moe_lora_weights(name: str, stacked_tensor: torch.Tensor) -> Dict[s For hybrid shared LoRA, shared weights (shape[0] == 1) are named with ".shared." instead of expert index to indicate they're shared across experts. """ - match = moe_lora_pattern.match(name) + match = _MOE_LORA_PATTERN.match(name) if not match: raise ValueError(f"Invalid MoE LoRA parameter name: {name}") @@ -486,7 +586,6 @@ def _unmerge_moe_lora_weights(name: str, stacked_tensor: torch.Tensor) -> Dict[s peft_state_dict = {} detected_modules = set() detected_r = None - detected_alpha = None for key, value in lora_state_dict.items(): # Check if this is a stacked MoE LoRA parameter @@ -495,7 +594,7 @@ def _unmerge_moe_lora_weights(name: str, stacked_tensor: torch.Tensor) -> Dict[s per_expert_weights = _unmerge_moe_lora_weights(key, value) peft_state_dict.update(per_expert_weights) # Detect target modules from MoE LoRA - match = moe_lora_pattern.match(key) + match = _MOE_LORA_PATTERN.match(key) if match: detected_modules.add(match.group(2)) # gate_proj, up_proj, or down_proj if detected_r is None and match.group(3) == "A": @@ -570,17 +669,8 @@ def load_lora_checkpoint( Load LoRA weights from checkpoint. Supports both xorl format and PEFT format checkpoints for dense - (nn.Linear) and per-expert MoE LoRA weights. - - .. note:: - Hybrid-shared MoE LoRA checkpoints are **not** supported. - ``save_lora_checkpoint`` exports shared expert weights under - ``.mlp.experts.shared.{proj}.lora_{A|B}.weight``, but this - function cannot map those keys back to xorl's internal stacked - MoE parameter layout (``{proj}_lora_{A|B}`` with shape - ``[1, ...]``). Loading such a checkpoint would require - re-stacking the per-expert and shared tensors, which is not - yet implemented. + (nn.Linear) LoRA as well as stacked MoE LoRA, including hybrid-shared + checkpoints exported under ``.mlp.experts.shared.{proj}.lora_{A|B}.weight``. Args: model: Model with LoRA layers already injected @@ -604,36 +694,10 @@ def load_lora_checkpoint( else: state_dict = torch.load(weights_file, map_location="cpu", weights_only=True) - def _convert_from_peft_key(name: str) -> str: - """ - Convert PEFT key format back to xorl format. - - For Linear layers: lora_A.weight -> lora_A, lora_B.weight -> lora_B - For lm_head: lora_embedding_A -> lora_A, lora_embedding_B -> lora_B - """ - # Handle lm_head embedding-style naming - if "lm_head.lora_embedding_A" in name: - return name.replace("lm_head.lora_embedding_A", "lm_head.lora_A") - elif "lm_head.lora_embedding_B" in name: - return name.replace("lm_head.lora_embedding_B", "lm_head.lora_B") - # Handle .weight suffix for Linear layers - elif name.endswith(".lora_A.weight"): - return name[: -len(".weight")] - elif name.endswith(".lora_B.weight"): - return name[: -len(".weight")] - return name - - # Convert PEFT format keys to xorl format if needed - converted_state_dict = {} - for key, value in state_dict.items(): - # Remove PEFT prefix: base_model.model.{key} -> {key} - if key.startswith("base_model.model."): - new_key = key[len("base_model.model.") :] - else: - new_key = key - # Convert from PEFT naming to xorl naming - new_key = _convert_from_peft_key(new_key) - converted_state_dict[new_key] = value + model_lora_shapes = { + key: value.shape for key, value in model.state_dict().items() if "lora_A" in key or "lora_B" in key + } + converted_state_dict = convert_peft_lora_state_dict(state_dict, expected_shapes=model_lora_shapes) # Load into model missing, unexpected = load_lora_state_dict(model, converted_state_dict, strict=strict) diff --git a/src/xorl/server/runner/adapters/manager.py b/src/xorl/server/runner/adapters/manager.py index e3b03da5..40627367 100644 --- a/src/xorl/server/runner/adapters/manager.py +++ b/src/xorl/server/runner/adapters/manager.py @@ -25,6 +25,8 @@ from safetensors.torch import load_file as safetensors_load_file from safetensors.torch import save_file as safetensors_save_file +from xorl.lora.utils import convert_peft_lora_state_dict + try: from torch.distributed._tensor import DTensor @@ -669,13 +671,22 @@ def load_adapter_state( weights_path = os.path.join(path, "adapter_model.safetensors") if os.path.exists(weights_path): loaded_weights = safetensors_load_file(weights_path) - - # Convert from PEFT format back to internal format - for peft_name, weight in loaded_weights.items(): - # Remove "base_model.model." prefix - internal_name = peft_name.replace("base_model.model.", "") - if internal_name in state.lora_params: - state.lora_params[internal_name].data.copy_(weight.to(self.device)) + expected_shapes = {name: param.shape for name, param in state.lora_params.items()} + converted_weights = convert_peft_lora_state_dict(loaded_weights, expected_shapes=expected_shapes) + + expected_keys = set(state.lora_params) + loaded_keys = set(converted_weights) + missing_keys = sorted(expected_keys - loaded_keys) + unexpected_keys = sorted(loaded_keys - expected_keys) + if missing_keys or unexpected_keys: + raise RuntimeError( + "Checkpoint LoRA parameter set does not match the live adapter structure.\n" + f"missing={missing_keys}\n" + f"unexpected={unexpected_keys}" + ) + + for name, param in state.lora_params.items(): + param.data.copy_(converted_weights[name].to(device=self.device, dtype=param.dtype)) else: raise FileNotFoundError(f"Weights file not found: {weights_path}") diff --git a/tests/server/runner/test_lora_checkpoint_roundtrip.py b/tests/server/runner/test_lora_checkpoint_roundtrip.py new file mode 100644 index 00000000..13c1f9fd --- /dev/null +++ b/tests/server/runner/test_lora_checkpoint_roundtrip.py @@ -0,0 +1,144 @@ +import importlib.util +from pathlib import Path + +import pytest +import torch +import torch.nn as nn + +from xorl.lora.modules.linear import LoraLinear +from xorl.lora.utils import load_lora_checkpoint, save_lora_checkpoint +from xorl.models.layers.moe import MoEExpertsLoRA, MoELoRAConfig + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + +_TARGET_MODULES = ["q_proj", "gate_proj", "up_proj", "down_proj"] +_REPO_ROOT = Path(__file__).resolve().parents[3] +_MANAGER_SPEC = importlib.util.spec_from_file_location( + "xorl_server_runner_adapters_manager", + _REPO_ROOT / "src/xorl/server/runner/adapters/manager.py", +) +assert _MANAGER_SPEC is not None and _MANAGER_SPEC.loader is not None +_MANAGER_MODULE = importlib.util.module_from_spec(_MANAGER_SPEC) +_MANAGER_SPEC.loader.exec_module(_MANAGER_MODULE) +LoRAAdapterManager = _MANAGER_MODULE.LoRAAdapterManager + + +class _TinyAttention(nn.Module): + def __init__(self): + super().__init__() + self.q_proj = LoraLinear.from_module(nn.Linear(8, 8, bias=False), r=2, lora_alpha=4) + + +class _TinyLayer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = _TinyAttention() + self.mlp = nn.Module() + self.mlp.experts = MoEExpertsLoRA( + num_experts=4, + hidden_dim=8, + intermediate_size=16, + moe_implementation="eager", + lora_config=MoELoRAConfig(r=2, lora_alpha=4, hybrid_shared=True), + ) + + +class _TinyInnerModel(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([_TinyLayer()]) + + +class _TinyMoELoraModel(nn.Module): + def __init__(self): + super().__init__() + self.model = _TinyInnerModel() + + +def _iter_lora_parameters(module: nn.Module): + for name, param in module.named_parameters(): + if "lora_" in name: + yield name, param + + +def _assign_distinct_lora_values(module: nn.Module) -> None: + with torch.no_grad(): + for offset, (_, param) in enumerate(_iter_lora_parameters(module), start=1): + values = torch.arange(param.numel(), dtype=torch.float32).reshape(param.shape) + offset + param.copy_(values.to(param.dtype)) + + +def _expected_saved_lora_state(module: nn.Module) -> dict[str, torch.Tensor]: + return { + name: param.detach().cpu().to(torch.bfloat16).to(param.dtype).clone() + for name, param in _iter_lora_parameters(module) + } + + +def _actual_lora_state(module: nn.Module) -> dict[str, torch.Tensor]: + return {name: param.detach().cpu().clone() for name, param in _iter_lora_parameters(module)} + + +def test_load_lora_checkpoint_roundtrip_supports_hybrid_shared(tmp_path): + source = _TinyMoELoraModel() + _assign_distinct_lora_values(source) + + checkpoint_dir = tmp_path / "checkpoint" + save_lora_checkpoint( + model=source, + save_path=str(checkpoint_dir), + target_modules=_TARGET_MODULES, + r=2, + lora_alpha=4, + moe_hybrid_shared_lora=True, + ) + + loaded = _TinyMoELoraModel() + load_lora_checkpoint(loaded, str(checkpoint_dir), strict=True) + + expected = _expected_saved_lora_state(source) + actual = _actual_lora_state(loaded) + + assert set(actual) == set(expected) + for name, expected_tensor in expected.items(): + assert torch.equal(actual[name], expected_tensor), name + + +def test_adapter_manager_load_adapter_state_roundtrip_supports_hybrid_shared(tmp_path): + source = _TinyMoELoraModel() + _assign_distinct_lora_values(source) + + checkpoint_dir = tmp_path / "checkpoint" + save_lora_checkpoint( + model=source, + save_path=str(checkpoint_dir), + target_modules=_TARGET_MODULES, + r=2, + lora_alpha=4, + moe_hybrid_shared_lora=True, + ) + + manager = LoRAAdapterManager( + model=_TinyMoELoraModel(), + device=torch.device("cpu"), + checkpoint_dir=str(tmp_path / "adapters"), + auto_save_on_eviction=False, + lora_config={"moe_hybrid_shared_lora": True}, + ) + + result = manager.load_adapter_state( + model_id="adapter-1", + path=str(checkpoint_dir), + load_optimizer=False, + lr=1e-4, + ) + + state = manager.get_adapter_state("adapter-1") + actual = {name: param.detach().cpu().clone() for name, param in state.lora_params.items()} + expected = _expected_saved_lora_state(source) + + assert result["model_id"] == "adapter-1" + assert set(actual) == set(expected) + for name, expected_tensor in expected.items(): + assert torch.equal(actual[name], expected_tensor), name From ad2577b79a2914ce2578e7e9c934175891d7bea8 Mon Sep 17 00:00:00 2001 From: Conner Manuel <57027354+connermanuel@users.noreply.github.com> Date: Fri, 17 Apr 2026 12:52:14 -0700 Subject: [PATCH 39/41] Lazy import repeat_kv (#147) (cherry picked from commit 4a74e79af8f9750447a0e3dd47bcb1dc422e6440) --- src/xorl/distributed/sequence_parallel/strategy.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/xorl/distributed/sequence_parallel/strategy.py b/src/xorl/distributed/sequence_parallel/strategy.py index c7134b75..96342f50 100644 --- a/src/xorl/distributed/sequence_parallel/strategy.py +++ b/src/xorl/distributed/sequence_parallel/strategy.py @@ -19,7 +19,6 @@ import torch import torch.distributed as dist -from ...models.layers.attention.utils import repeat_kv from ...models.layers.rope import apply_rotary_pos_emb from .async_ulysses import async_ulysses_output_projection, async_ulysses_qkv_projection from .data import slice_position_embedding @@ -149,6 +148,9 @@ def __init__(self, group, ulysses_size: int): self.ulysses_size = ulysses_size def project_qkv(self, module, hidden_states, position_embeddings): + # Lazy-imported to break an import cycle with xorl.models.layers.attention + from ...models.layers.attention.utils import repeat_kv # noqa: PLC0415 + # Model-specific QKV projection (MHA, MLA, etc.) q, k, v = module._project_qkv(hidden_states, position_embeddings) From de00085bcf02f4b5c12082e06dc0b4cede3e6e90 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Sat, 18 Apr 2026 06:42:45 +0000 Subject: [PATCH 40/41] Sync src/ to match upstream snapshot --- src/xorl/models/module_utils.py | 48 ++++++++------------- src/xorl/ops/quack/gemm_config.py | 3 +- src/xorl/ops/quack/gemm_interface.py | 2 +- src/xorl/server/api_server/api_types.py | 5 ++- src/xorl/server/protocol/operations.py | 2 +- src/xorl/server/runner/runner_dispatcher.py | 2 +- src/xorl/server/runner/utils/moe_metrics.py | 2 +- src/xorl/server/weight_sync/handler.py | 10 ++--- 8 files changed, 32 insertions(+), 42 deletions(-) diff --git a/src/xorl/models/module_utils.py b/src/xorl/models/module_utils.py index 05940356..6ac45c14 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -351,19 +351,23 @@ def _try_load_state_dict(weights_path: str, **kwargs): and broadcasts the result to other ranks. This avoids N-way filesystem hammering and keeps the broadcast-loading path rank-0-driven. """ - cache_kwargs = {"_raise_exceptions_for_missing_entries": False, **kwargs} - resolved_weight_file = cached_file(weights_path, SAFE_WEIGHTS_NAME, **cache_kwargs) - if resolved_weight_file: - return [StateDictIterator(resolved_weight_file)] + rank = dist.get_rank() if dist.is_initialized() else 0 + world_size = dist.get_world_size() if dist.is_initialized() else 1 - resolved_weight_file = cached_file(weights_path, SAFE_WEIGHTS_INDEX_NAME, **cache_kwargs) - if resolved_weight_file: - # Retry on OSError (e.g., missing shard files during download) - max_retries = 3 + if world_size <= 1: + # Single rank: do everything locally (no broadcast needed) + return _try_load_state_dict_local(weights_path, **kwargs) + + # Multi-rank: rank 0 resolves paths, broadcasts to all + resolved_paths = [None] + if rank == 0: + max_retries = 10 for attempt in range(max_retries): try: - shard_files, _ = get_checkpoint_shard_files(weights_path, resolved_weight_file, **kwargs) - return [StateDictIterator(shard_file) for shard_file in shard_files] + result = _try_load_state_dict_local(weights_path, **kwargs) + if result is not None: + resolved_paths[0] = [it.filepath for it in result] + break except OSError as e: if attempt < max_retries - 1: logger.warning( @@ -371,7 +375,7 @@ def _try_load_state_dict(weights_path: str, **kwargs): ) time.sleep(5) else: - logger.error(f"Failed to get shard files after {max_retries} attempts") + logger.error(f"Failed to resolve shard files after {max_retries} attempts") raise _broadcast_object_list_weight_load(resolved_paths, src=0) @@ -399,24 +403,8 @@ def _try_load_state_dict_local(weights_path: str, **kwargs): resolved_weight_file = cached_file(weights_path, WEIGHTS_INDEX_NAME, **cache_kwargs) if resolved_weight_file: - # Retry on OSError (e.g., missing shard files during download) - max_retries = 3 - for attempt in range(max_retries): - try: - shard_files, _ = get_checkpoint_shard_files(weights_path, resolved_weight_file, **kwargs) - return [StateDictIterator(shard_file) for shard_file in shard_files] - except OSError as e: - if attempt < max_retries - 1: - import time - retry_delay = 5 - logger.warning( - f"OSError getting PT shard files (attempt {attempt + 1}/{max_retries}): {e}. " - f"Retrying in {retry_delay}s..." - ) - time.sleep(retry_delay) - else: - logger.error(f"Failed to get PT shard files after {max_retries} attempts") - raise + shard_files, _ = get_checkpoint_shard_files(weights_path, resolved_weight_file, **kwargs) + return [StateDictIterator(shard_file) for shard_file in shard_files] return None @@ -872,7 +860,7 @@ def all_ranks_load_weights( ) # Retry loading state dict on OSError (e.g., HuggingFace download issues) - max_retries = 3 + max_retries = 10 retry_delay = 5 # seconds last_error = None diff --git a/src/xorl/ops/quack/gemm_config.py b/src/xorl/ops/quack/gemm_config.py index e43c544f..0a6c99e3 100644 --- a/src/xorl/ops/quack/gemm_config.py +++ b/src/xorl/ops/quack/gemm_config.py @@ -23,7 +23,8 @@ def get_all_configs( tune_coop: bool = True, # tune_raster_order=True, ) -> List[GemmConfig]: - assert device_capacity in [9, 10] + if device_capacity not in [9, 10]: + return [] if device_capacity == 9: tile_n_vals = [128, 144, 160, 176, 192, 208] tile_mn_coop_vals = [(256, tile_n) for tile_n in tile_n_vals] + [ diff --git a/src/xorl/ops/quack/gemm_interface.py b/src/xorl/ops/quack/gemm_interface.py index 441331c3..fc71fb31 100644 --- a/src/xorl/ops/quack/gemm_interface.py +++ b/src/xorl/ops/quack/gemm_interface.py @@ -37,7 +37,7 @@ } -default_device_capacity = get_device_capacity(torch.device("cuda")) +default_device_capacity = get_device_capacity(torch.device("cuda")) if torch.cuda.is_available() else (0, 0) def default_config(device): diff --git a/src/xorl/server/api_server/api_types.py b/src/xorl/server/api_server/api_types.py index f09b3eb5..d1c62786 100644 --- a/src/xorl/server/api_server/api_types.py +++ b/src/xorl/server/api_server/api_types.py @@ -803,9 +803,10 @@ class SyncInferenceWeightsRequest(BaseModel): group_name: str = Field(default="weight_sync_group", description="Name of the NCCL process group") buffer_size_mb: int = Field(default=1024, description="Size of each transfer bucket in MB (to avoid OOM)") flush_cache: bool = Field( - default=True, + default=False, description="Whether to flush inference KV cache after weight sync. " - "Set to False to keep cached KV when weights haven't changed significantly.", + "Usually not needed: the radix cache is keyed on token sequences, " + "so stale KV entries from previous weights are evicted naturally.", ) pause_mode: Literal["retract", "abort", "in_place"] = Field( default="retract", diff --git a/src/xorl/server/protocol/operations.py b/src/xorl/server/protocol/operations.py index 42542e1f..395ecca2 100644 --- a/src/xorl/server/protocol/operations.py +++ b/src/xorl/server/protocol/operations.py @@ -103,7 +103,7 @@ class SyncWeightsData: group_name: str = "weight_sync_group" buffer_size_mb: int = 1024 sync_method: str = "nccl_broadcast" - flush_cache: bool = True + flush_cache: bool = False pause_mode: str = "retract" weight_version: Optional[str] = None quantization: Optional[Dict[str, Any]] = None diff --git a/src/xorl/server/runner/runner_dispatcher.py b/src/xorl/server/runner/runner_dispatcher.py index 9d0f2c6c..de92fad4 100644 --- a/src/xorl/server/runner/runner_dispatcher.py +++ b/src/xorl/server/runner/runner_dispatcher.py @@ -1233,7 +1233,7 @@ async def _handle_wake_up(self, command_dict: Dict[str, Any]) -> Dict[str, Any]: # Health Check Handlers # ======================================================================== - async def _handle_health_check(self) -> Dict[str, Any]: + async def _handle_health_check(self, command_dict: Dict[str, Any] = None) -> Dict[str, Any]: """Handle health check on all ranks (unified handler).""" logger.debug(f"Rank {self.rank}: Health check") diff --git a/src/xorl/server/runner/utils/moe_metrics.py b/src/xorl/server/runner/utils/moe_metrics.py index 6e1e1e11..3d12b4e6 100644 --- a/src/xorl/server/runner/utils/moe_metrics.py +++ b/src/xorl/server/runner/utils/moe_metrics.py @@ -69,7 +69,7 @@ def __init__(self, model_config_obj, train_config: Dict[str, Any], rank: int): self.enabled = True logger.info(f"MoE expert metrics collection enabled (num_experts={num_experts})") else: - logger.warning("Could not enable Qwen3 MoE expert metrics - module not available") + logger.debug("MoE expert metrics not available (enable_expert_metrics not implemented)") def collect(self, step: int, forward_backward_time: float) -> dict: """ diff --git a/src/xorl/server/weight_sync/handler.py b/src/xorl/server/weight_sync/handler.py index 3b1befea..b0e1e8c5 100644 --- a/src/xorl/server/weight_sync/handler.py +++ b/src/xorl/server/weight_sync/handler.py @@ -1365,7 +1365,7 @@ def _unfuse_for_inference( buffer: List[Tuple[str, torch.Tensor]], model, ) -> List[Tuple[str, torch.Tensor]]: - """Split fused projections (qkv_proj, gate_up_proj) into HF-format names. + """Convert training parameter names to HF/SGLang inference names. Handles: - qkv_proj → q_proj + k_proj + v_proj (split fused attention) @@ -1375,8 +1375,6 @@ def _unfuse_for_inference( HF fused names (q_proj/k_proj/v_proj → in_proj_qkv, etc.) """ - This splits fused weights back to individual projections before sending. - """ config = model.config num_heads = config.num_attention_heads num_kv_heads = getattr(config, "num_key_value_heads", num_heads) @@ -1387,7 +1385,6 @@ def _unfuse_for_inference( result = [] for name, tensor in buffer: if ".qkv_proj." in name: - # Split [q_size + 2*kv_size, hidden] → q, k, v prefix, suffix = name.rsplit(".qkv_proj.", 1) q = tensor[:q_size].clone() k = tensor[q_size : q_size + kv_size].clone() @@ -1409,7 +1406,6 @@ def _unfuse_for_inference( for expert_idx in range(tensor.shape[0]): result.append((f"{prefix}.{expert_idx}.down_proj.weight", down[expert_idx])) elif ".gate_up_proj." in name: - # Split [2*intermediate, hidden] → gate, up prefix, suffix = name.rsplit(".gate_up_proj.", 1) half = tensor.shape[0] // 2 gate = tensor[:half].clone() @@ -1418,6 +1414,10 @@ def _unfuse_for_inference( result.append((f"{prefix}.up_proj.{suffix}", up)) else: result.append((name, tensor)) + + if has_linear_attention_layers(config): + result = remap_linear_attention_params_for_inference(result) + return result def _broadcast_buffer( From b139a5ab168e7e1136ea89848f9ab5991f4d76ec Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Sat, 18 Apr 2026 06:52:12 +0000 Subject: [PATCH 41/41] Apply ruff-format fixes --- tests/server/api_server/test_api_server.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/server/api_server/test_api_server.py b/tests/server/api_server/test_api_server.py index b6171b91..a3d46897 100644 --- a/tests/server/api_server/test_api_server.py +++ b/tests/server/api_server/test_api_server.py @@ -111,6 +111,7 @@ def test_request_creation_and_serialization(self): request = LoadWeightsRequest(path="/tmp/checkpoint", optimizer=True) assert request.optimizer is True + class TestTinkerSessionCompatibility: """Test Tinker-compatible session creation and heartbeats.""" @@ -176,5 +177,7 @@ def test_session_heartbeat_refreshes_activity(self): assert heartbeat_response.session_id == session_id assert server.session_last_activity[session_id] > initial_activity + + if __name__ == "__main__": pytest.main([__file__, "-v"])