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/.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/.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/README.md b/README.md index 64604525..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 @@ -78,6 +133,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. @@ -85,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/config-reference/local.md b/docs/src/content/docs/config-reference/local.md index 65f8f309..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` @@ -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. | @@ -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/config-reference/server.md b/docs/src/content/docs/config-reference/server.md index e2faf5d8..21eccae1 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. | @@ -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. | @@ -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/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/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/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/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/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/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/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/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..190200fe 100644 --- a/docs/src/content/docs/server-training/overview.md +++ b/docs/src/content/docs/server-training/overview.md @@ -1,493 +1,174 @@ --- -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. - -### Creating a training client - -**Full-weight training** (no LoRA — server auto-registers `model_id="default"` on startup): - -```python -training_client = xorl_client.TrainingClient( - holder=service_client.holder, - model_id="default", - base_model="Qwen/Qwen3-8B", -) -``` - -**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", -) -``` +service = xorl_client.ServiceClient(base_url="http://localhost:6000") +client = service.create_training_client(base_model="Qwen/Qwen3-8B") -### Preparing data +# 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, +}) -Training data is passed as a list of `Datum` objects: +# 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 -datum = xorl_client.Datum( - model_input=xorl_client.ModelInput.from_ints(input_ids), - loss_fn_inputs={"labels": labels}, -) + if step % sync_interval == 0: + client.sync_inference_weights( + master_address="localhost", + master_port=29600, + ).result() ``` -`ModelInput.from_ints` wraps a flat token list into the expected `{"input_ids": [...]}` format. +For multi-node setup, server configuration, and launcher CLI options, see [Launching & Configuration](/server-training/training-server/launching/). -### Training loop +### Available configs -`forward_backward` and `optim_step` return futures — call `.result()` to block and retrieve the output: - -```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}") -``` +| 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 | -### Available loss functions +### Matching inference models -| `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 -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. - -### Tinker-compatible endpoints +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_model` | Create/register a model session | +| `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 | -### Field mappings - -xorl automatically maps Tinker's field names to its own format, so Tinker clients work without changes: - -| 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] -} -``` +## Recommended Reading -`LossFnOutput` responses always use `TensorData` values: +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: -```json -{ - "loss": {"data": [2.345], "dtype": "float32", "shape": [1]}, - "logprobs": {"data": [...], "dtype": "float32", "shape": [512]}, - "elementwise_loss": {"data": [...], "dtype": "float32", "shape": [512]} -} -``` +### Foundational RL for LLMs -`logprobs` and `elementwise_loss` are only present when `loss_fn: "per_token_ce"` is used. +| Paper | Description | +|---|---| +| [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. | ---- +### System Design: Pipelining and Async RL -## 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 9f5697b8..81993257 100644 --- a/docs/src/content/docs/server-training/sglang.mdx +++ b/docs/src/content/docs/server-training/sglang.mdx @@ -16,37 +16,40 @@ 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.). --- ## 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: @@ -78,52 +90,112 @@ 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 ``` -```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..9dddf172 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` | 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 @@ -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/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/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/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/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_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..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 @@ -34,7 +35,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 +53,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 +61,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/local/dummy/configs/lora/llama3_8b_lora.yaml b/examples/local/dummy/configs/lora/llama3_8b_lora.yaml new file mode 100644 index 00000000..b922bb59 --- /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: ["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 new file mode 100644 index 00000000..ffd3a529 --- /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: ["qkv_proj", "o_proj", "gate_up_proj", "down_proj"] + quant_format: nvfp4 + quant_group_size: 16 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..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,10 +8,13 @@ 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_native: true +rmsnorm_mode: native activation_native: true rope_native: true attention_cast_bf16: true 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 new file mode 100644 index 00000000..ef6bb58a --- /dev/null +++ b/examples/server/configs/full/qwen3_5_35b_a3b_full.yaml @@ -0,0 +1,45 @@ +# 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 +load_weights_mode: all_ranks + +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/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/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/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/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/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 7f2e2a07..bfed180c 100644 --- a/examples/server/password_memorization/run_password_test.py +++ b/examples/server/password_memorization/run_password_test.py @@ -31,11 +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", @@ -43,8 +47,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 +55,7 @@ # Training data # --------------------------------------------------------------------------- + def build_training_data(tokenizer): data = [] for project, code in CODES.items(): @@ -61,14 +65,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 +88,20 @@ def build_training_data(tokenizer): # Server helpers # --------------------------------------------------------------------------- + +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 = 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) @@ -91,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 @@ -103,77 +124,146 @@ 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): 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 {} - 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", - json={"model_id": "default", "forward_backward_input": {"data": data, "loss_fn": "causallm_loss"}}, - timeout=30) - fb_result = wait_for_future(train_url, fb.json()["request_id"]) + fb = requests.post( + f"{train_url}/api/v1/forward_backward", + json={"model_id": MODEL_ID, "forward_backward_input": {"data": data, "loss_fn": "causallm_loss"}}, + timeout=30, + ) + 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}, - timeout=30) - opt_result = wait_for_future(train_url, opt.json()["request_id"]) + opt = requests.post( + f"{train_url}/api/v1/optim_step", + json={"model_id": MODEL_ID, "adam_params": {"learning_rate": lr}, "gradient_clip": 1.0}, + timeout=30, + ) + 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 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) + 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) + resp = requests.post( + 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')}") + print( + f" Sync: success={result.get('success')}, {time.time() - t0:.1f}s, " + 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", - "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)" - - -def test_inference(infer_urls, label=""): + resp = requests.post( + f"{infer_url}/v1/chat/completions", + json={ + "model": MODEL_ID, + "messages": [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt}], + "max_tokens": 32, + "temperature": 0, + "chat_template_kwargs": {"enable_thinking": False}, + }, + timeout=30, + ) + 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="", 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 @@ -182,6 +272,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,21 +299,28 @@ 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 + 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): - print(f" FAILED: {url} not running"); return None + 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() @@ -232,46 +330,68 @@ 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) + 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 - - correct = test_inference(infer_urls, "after sync") + 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, f"after sync ({sync_version})", expected_weight_version=sync_version) print(f" Score: {correct}/{total_codes}") return correct 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") + 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) @@ -285,14 +405,14 @@ def main(): correct = run_test(args, training_data) except Exception as e: print(f" ERROR: {e}") - import traceback; traceback.print_exc() + 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.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 a93d5a79..4cb41d74 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", @@ -49,11 +50,12 @@ dev = [ {include-group = "test"}, ] lint = [ - "pre-commit", + "pre-commit>=4.5,<5", "ruff", ] test = [ "pytest", + "pytest-asyncio", "expecttest" ] @@ -80,7 +82,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 +92,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/__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 3704c9a6..27c8a10d 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -2,12 +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 @@ -24,16 +22,34 @@ get_type_hints, ) +import torch import yaml +from .utils import logging +from .utils.checkpoint_utils import get_checkpoint_path -from .utils import helper, logging T = TypeVar("T") 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) @@ -43,43 +59,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, @@ -89,9 +96,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, @@ -102,37 +107,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): @@ -183,17 +178,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""" @@ -240,33 +231,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, @@ -306,15 +287,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, @@ -336,9 +313,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, @@ -348,33 +323,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, @@ -384,9 +349,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, @@ -395,7 +358,6 @@ class DataArguments: }, ) - def _convert_datasets_to_config( self, datasets: Optional[List[Dict[str, Any]]], dataset_name: str ) -> List[DatasetConfig]: @@ -410,25 +372,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.") @@ -456,21 +412,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, @@ -480,28 +430,39 @@ 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", 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')."}, ) 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, @@ -509,21 +470,28 @@ 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)." + }, + ) + 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: - 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 @@ -538,10 +506,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: @@ -555,9 +520,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, @@ -580,9 +543,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", @@ -590,7 +556,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, @@ -612,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]: @@ -623,19 +640,33 @@ 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( 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, @@ -671,20 +702,21 @@ def optimizer_kwargs(self) -> Dict[str, Any]: default=False, metadata={"help": "Use reentrant gradient checkpointing."}, ) - recompute_modules: Optional[List[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( + gradient_checkpointing_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" + ) }, ) @@ -695,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, @@ -745,9 +775,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", @@ -765,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, @@ -833,13 +861,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, @@ -851,7 +879,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, @@ -869,6 +899,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."}, @@ -909,16 +943,32 @@ 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): + # 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")) 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( @@ -932,36 +982,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}) * " @@ -969,9 +1008,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. " @@ -979,19 +1016,14 @@ 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 - self.load_checkpoint_path = get_checkpoint_path( output_dir=self.output_dir, is_local_rank0=self.local_rank == 0, @@ -1017,16 +1049,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 @@ -1039,7 +1068,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", @@ -1094,23 +1125,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, @@ -1129,9 +1166,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, @@ -1223,7 +1258,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) @@ -1234,9 +1269,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() @@ -1247,9 +1280,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: @@ -1259,17 +1290,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) @@ -1292,9 +1317,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: @@ -1306,9 +1329,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] @@ -1320,24 +1341,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 @@ -1346,12 +1361,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 @@ -1372,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: @@ -1450,22 +1463,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 @@ -1494,6 +1498,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..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 @@ -100,8 +102,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"} @@ -225,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: @@ -234,9 +234,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 +243,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 +290,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 +303,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 +329,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 +475,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) @@ -543,7 +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 bd7ffb05..e6bf5606 100644 --- a/src/xorl/cli/direct_train.py +++ b/src/xorl/cli/direct_train.py @@ -1,19 +1,44 @@ +# 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 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 + # 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 import time -from dataclasses import asdict, dataclass, field -from functools import partial -from typing import Any, Dict, List +from dataclasses import asdict +from typing import Any, Dict, List, Optional -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,9 +49,14 @@ 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.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, @@ -34,20 +64,27 @@ 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 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(): @@ -55,18 +92,67 @@ 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() if args.train.global_rank == 0: save_args(args, args.train.output_dir) + if args.train.use_wandb: + import wandb # noqa: PLC0415 + + 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, 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} + 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 # noqa: PLC0415 + + 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) - + init_parallel_state( dp_size=args.train.data_parallel_size, dp_replicate_size=args.train.data_parallel_replicate_size, @@ -85,7 +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") @@ -105,7 +190,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,18 +203,24 @@ 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_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, + 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, + rmsnorm_mode=args.model.rmsnorm_mode, init_device=args.train.init_device, ) model_config = model.config @@ -152,7 +243,6 @@ 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 if detect_prequantized_nvfp4(args.model.model_path): is_prequantized = True checkpoint_quant_format = "nvfp4" @@ -163,16 +253,12 @@ 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( @@ -199,15 +285,37 @@ 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, - 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( + 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, @@ -219,13 +327,15 @@ 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, 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 @@ -248,14 +358,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_quantize_qlora, maybe_load_and_quantize_moe_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 +374,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: @@ -303,18 +411,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) @@ -337,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, ) @@ -363,18 +458,19 @@ 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): """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 +537,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() @@ -450,14 +547,13 @@ 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 set_replay_stage, RoutingReplay + use_routing_replay = ps.ep_size > 1 and args.train.moe_recomputed if pp_enabled: @@ -471,18 +567,14 @@ 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 +649,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 +696,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, ) @@ -637,21 +734,26 @@ 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: 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, - }) + import wandb # noqa: PLC0415 + + 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 +761,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 +782,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 +805,8 @@ 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 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 @@ -715,18 +817,22 @@ 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, 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, 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: 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/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..1e0327f6 100644 --- a/src/xorl/cli/train.py +++ b/src/xorl/cli/train.py @@ -1,8 +1,27 @@ +# 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() + +# 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/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..7c54041c 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"]: @@ -73,7 +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 @@ -85,7 +84,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 +107,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 +121,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 +146,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 +162,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 +184,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 +220,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..ba2e3ff6 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 @@ -12,6 +12,7 @@ PackingConcatCollator, ShiftTokensCollator, TextSequenceShardCollator, + ToTensorCollator, ) @@ -25,17 +26,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 +46,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 +66,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 +94,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 +164,40 @@ 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 e61cd5ed..ea894e9b 100644 --- a/src/xorl/data/prepare/packing.py +++ b/src/xorl/data/prepare/packing.py @@ -3,29 +3,25 @@ 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 shutil 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 +93,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 +119,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 +169,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 +179,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 +299,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 +316,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 +351,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 +395,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 +421,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,35 +434,39 @@ 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) - + 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 - }) - + 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.""" - # 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")) - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - world_size = int(os.environ.get("WORLD_SIZE", "1")) # Try to load from cache first cached_bins = self._load_cached_bins() @@ -490,52 +475,36 @@ 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) 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...") - # 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 = 3 - 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: - 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") @@ -549,27 +518,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 +542,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 +561,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..85c50c93 100644 --- a/src/xorl/data/prepare/prepare_datasets.py +++ b/src/xorl/data/prepare/prepare_datasets.py @@ -1,36 +1,40 @@ """Data handling specific to SFT.""" -import functools -import os -import tempfile from typing import Literal +import numpy as np +import pyarrow as pa +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__) @@ -45,58 +49,39 @@ 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()) - 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 +109,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 +170,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 +249,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 +266,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 +283,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 +317,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 +331,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 +360,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..41e3f978 100644 --- a/src/xorl/data/prepare/shared.py +++ b/src/xorl/data/prepare/shared.py @@ -3,17 +3,16 @@ from __future__ import annotations import functools -import hashlib import os +from dataclasses import replace 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, @@ -26,10 +25,31 @@ 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.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 +61,8 @@ pass from ...utils import logging + + LOG = logging.get_logger(__name__) EXTENSIONS_TO_DATASET_TYPES = { @@ -51,8 +73,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: @@ -81,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): @@ -148,9 +167,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,49 +204,31 @@ 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: - 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, {} @@ -260,9 +259,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 +284,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 +307,7 @@ def _load_from_hub( **load_dataset_kwargs, ) + def _load_from_cloud( dataset_config: DatasetConfig, remote_fs: S3FileSystem | GCSFileSystem | AzureBlobFileSystem | OCIFileSystem, @@ -336,9 +330,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 +374,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 +401,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 +499,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 +516,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 +544,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/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/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..6159359d 100644 --- a/src/xorl/distributed/moe/__init__.py +++ b/src/xorl/distributed/moe/__init__.py @@ -9,13 +9,13 @@ def __getattr__(name): "alltoall_pre_dispatch", "alltoall_post_combine", ): - from .alltoall import ( + from .alltoall import ( # noqa: PLC0415 + 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"): - from .deepep import ( - DeepEPBuffer, + if name in ( + "DeepEPBuffer", + "DEEPEP_AVAILABLE", + "token_pre_dispatch", + "tokens_post_combine", + "get_default_buffer", + "destroy_default_buffer", + ): + from .deepep import ( # noqa: PLC0415 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..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 @@ -144,14 +163,14 @@ 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 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 @@ -180,9 +199,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, @@ -201,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, @@ -208,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, ) @@ -237,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 fcb73333..100d834e 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 @@ -327,12 +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, device=grad_expert_input.device, + ctx.num_recv_tokens, + ctx.hidden_dim, + 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()) @@ -352,16 +353,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,18 +377,30 @@ 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, 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) + # 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=expert_output.dtype, device=device, + dispatch_ctx.num_recv_tokens, + hidden_dim, + dtype=torch.float32, + device=device, ) idx_2d = dispatch_ctx.permuted_indices.unsqueeze(1).expand(-1, hidden_dim) - gather_output.scatter_add_(0, idx_2d, weighted) + _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].float()) + gather_output = gather_output.to(dtype) # Step 2: Combine previous_event = EventOverlap(EventHandle()) @@ -405,9 +418,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 @@ -418,7 +429,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 @@ -437,11 +448,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 @@ -501,9 +509,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 +532,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 +573,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 +585,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 +618,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/moe/utils.py b/src/xorl/distributed/moe/utils.py index 8bbe3ad8..4f044e39 100644 --- a/src/xorl/distributed/moe/utils.py +++ b/src/xorl/distributed/moe/utils.py @@ -28,37 +28,34 @@ 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) - - # 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) + # 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) + - 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: diff --git a/src/xorl/distributed/parallel_plan.py b/src/xorl/distributed/parallel_plan.py index b72bc3e9..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 @@ -112,7 +113,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()} @@ -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 733141fd..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: @@ -83,7 +84,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.") @@ -94,10 +97,10 @@ def __post_init__(self): ) if self.cp_enabled: - from ..distributed.sequence_parallel import ( + from ..distributed.sequence_parallel import ( # noqa: PLC0415 init_sequence_parallel, - set_ringattn_group, set_data_parallel_group, + set_ringattn_group, set_ulysses_sequence_parallel_group, set_unified_sequence_parallel_group, ) @@ -145,7 +148,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() @@ -157,7 +160,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() @@ -402,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"]: @@ -410,7 +454,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() @@ -422,7 +466,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() @@ -442,7 +486,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() @@ -454,7 +498,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() @@ -470,7 +514,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() @@ -482,7 +526,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() @@ -531,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. @@ -641,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, @@ -652,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/pipeline_parallel.py b/src/xorl/distributed/pipeline_parallel.py index 518d6585..baf0a0b3 100644 --- a/src/xorl/distributed/pipeline_parallel.py +++ b/src/xorl/distributed/pipeline_parallel.py @@ -33,7 +33,11 @@ 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__) @@ -103,8 +107,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 +124,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 +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 .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..e81d5b35 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,26 +29,26 @@ 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 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/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..96342f50 100644 --- a/src/xorl/distributed/sequence_parallel/strategy.py +++ b/src/xorl/distributed/sequence_parallel/strategy.py @@ -14,11 +14,12 @@ """ from abc import ABC, abstractmethod -from typing import Optional, Tuple +from typing import Optional import torch import torch.distributed as dist +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 @@ -62,6 +63,7 @@ def _scale_cu_seqlens_for_ringattn(kwargs, ringattn_group): # Base class # ------------------------------------------------------------------ # + class CPStrategy(ABC): """Base class for sequence parallel strategies. @@ -110,6 +112,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 +121,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 +135,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. @@ -146,7 +148,8 @@ 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 + # 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) @@ -155,8 +158,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 +192,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 +202,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 +224,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. @@ -245,11 +238,9 @@ 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: + 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) @@ -311,8 +302,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) @@ -323,9 +315,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 +324,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 +353,7 @@ def prepare_position_embeddings(self, position_embeddings, **kwargs): # Ring attention # ------------------------------------------------------------------ # + class RingAttentionStrategy(CPStrategy): """Ring attention without Ulysses. @@ -382,11 +369,11 @@ 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(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 +401,7 @@ def prepare_position_embeddings(self, position_embeddings, dim, sp_group, **kwar # Hybrid Ulysses + Ring # ------------------------------------------------------------------ # + class HybridUlyssesRingStrategy(CPStrategy): """Hybrid Ulysses + Ring attention. @@ -438,11 +426,11 @@ 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(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( @@ -489,7 +477,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/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..e92dbca1 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 @@ -218,20 +220,25 @@ 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: 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) @@ -242,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) @@ -295,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. - from torch.distributed._composable.fsdp.fully_shard import FSDPModule + # 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) @@ -323,8 +349,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": @@ -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, @@ -368,8 +455,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 ---- @@ -409,8 +500,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): @@ -421,17 +511,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: @@ -459,15 +547,17 @@ 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") 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) @@ -545,23 +635,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 @@ -575,6 +662,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) @@ -603,12 +691,12 @@ 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": - 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..1a698376 100644 --- a/src/xorl/lora/__init__.py +++ b/src/xorl/lora/__init__.py @@ -1,85 +1,75 @@ """ 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 -from xorl.lora.modules import LoraModule, LoraLinear - -# Mapping -from xorl.lora.mapping import ( - LORA_MAPPING, - can_apply_lora, - get_lora_class_for_module, - register_lora_mapping, -) - -# Utility functions -from xorl.lora.utils import ( - 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_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.lora.modules import LoraLinear, LoraModule 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, ) + +_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", @@ -88,24 +78,62 @@ ] -# 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 ( - MoELoRAConfig, + from xorl.models.layers.moe.lora import ( # noqa: PLC0415 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 g["MoEExpertsLoRA"] = MoEExpertsLoRA @@ -113,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/lora/mapping.py b/src/xorl/lora/mapping.py index 422e8c52..25222143 100644 --- a/src/xorl/lora/mapping.py +++ b/src/xorl/lora/mapping.py @@ -27,8 +27,10 @@ def _ensure_moe_mapping(): global _moe_registered if not _moe_registered: + from xorl.models.layers.moe import MoEExperts, MoEExpertsLoRA # noqa: PLC0415 + _moe_registered = True - from xorl.models.layers.moe import MoEExperts, MoEExpertsLoRA + LORA_MAPPING[MoEExperts] = MoEExpertsLoRA @@ -55,7 +57,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 06e5c31e..bf601e08 100644 --- a/src/xorl/lora/utils.py +++ b/src/xorl/lora/utils.py @@ -8,15 +8,18 @@ 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.distributed as dist import torch.nn as nn -from safetensors.torch import save_file, load_file +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 from xorl.lora.modules.base import LoraModule -from xorl.lora.mapping import can_apply_lora, get_lora_class_for_module + logger = logging.getLogger(__name__) @@ -29,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]: """ @@ -244,8 +253,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. @@ -257,8 +264,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 @@ -293,9 +298,8 @@ 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) + fqn2spec_info = getattr(model, "_fqn2spec_info", None) lora_state_dict = {} for name, param in model.named_parameters(): @@ -354,9 +358,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 @@ -369,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, @@ -378,6 +478,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,11 +498,13 @@ 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 """ - import re os.makedirs(save_path, exist_ok=True) @@ -409,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: """ @@ -439,15 +538,17 @@ 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. """ - match = moe_lora_pattern.match(name) + match = _MOE_LORA_PATTERN.match(name) if not match: raise ValueError(f"Invalid MoE LoRA parameter name: {name}") @@ -459,17 +560,21 @@ 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 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" @@ -481,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 @@ -490,12 +594,14 @@ 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": - # 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(".") @@ -562,7 +668,9 @@ 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) 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 @@ -586,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) @@ -655,7 +737,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: @@ -675,7 +756,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 @@ -701,12 +781,11 @@ 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, lora_alpha=lora_alpha, - shared_lora=shared_lora, target_modules=target_modules, hybrid_shared=hybrid_shared, ) @@ -716,12 +795,11 @@ 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( - "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 @@ -732,7 +810,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: """ @@ -754,7 +831,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: @@ -796,7 +872,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, ) @@ -818,10 +893,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/auto.py b/src/xorl/models/auto.py index cf3affdd..2015aad4 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, @@ -13,7 +14,10 @@ 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.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, @@ -42,15 +46,16 @@ 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)) + 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 @@ -80,7 +85,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 @@ -105,12 +109,13 @@ 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, 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, @@ -138,13 +143,21 @@ 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 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/base.py b/src/xorl/models/base.py index d03e11a2..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 @@ -46,7 +45,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. @@ -94,28 +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/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..58d562eb 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, +) # ============================================================================= @@ -20,9 +36,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 +48,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$" -) +# 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 -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 +83,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. @@ -135,7 +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) @@ -169,10 +169,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 +227,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 +283,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 +296,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") @@ -329,7 +318,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: @@ -344,6 +333,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 +351,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 +391,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) @@ -441,6 +430,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. @@ -464,36 +458,38 @@ 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. 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 @@ -501,15 +497,16 @@ 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()} # ============================================================================= # QKVMergeBuffer # ============================================================================= + class QKVMergeBuffer: """Buffer for merging separate q_proj, k_proj, v_proj checkpoint weights into qkv_proj. @@ -529,9 +526,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 +563,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 +573,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__( @@ -614,8 +619,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] = {} @@ -632,11 +635,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 +654,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 +665,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 +692,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 +726,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 ----- @@ -746,11 +747,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 ( - nvfp4_quantize, nvfp4_dequantize, - block_fp8_quantize_gkn, block_fp8_dequantize_gkn, - ) - from xorl.qlora.modules.linear import QLoRALinear module = accum.module M, K = module.out_features, module.in_features @@ -791,7 +787,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,11 +815,7 @@ class NvFP4WeightBuffer(_QLoRAWeightBufferBase): _suffixes = _QLORA_SUFFIXES_NVFP4 _source_format = "nvfp4" - def _pack( - self, fqn: str, accum: _ModuleAccum - ) -> List[Tuple[str, "torch.Tensor"]]: - from xorl.qlora.modules.linear import QLoRALinear - + def _pack(self, fqn: str, accum: _ModuleAccum) -> List[Tuple[str, "torch.Tensor"]]: module = accum.module merge_sources = accum.merge_sources @@ -872,11 +866,7 @@ class BlockFP8WeightBuffer(_QLoRAWeightBufferBase): _suffixes = _QLORA_SUFFIXES_FP8 _source_format = "block_fp8" - def _pack( - self, fqn: str, accum: _ModuleAccum - ) -> List[Tuple[str, "torch.Tensor"]]: - from xorl.qlora.modules.linear import QLoRALinear - + def _pack(self, fqn: str, accum: _ModuleAccum) -> List[Tuple[str, "torch.Tensor"]]: module = accum.module merge_sources = accum.merge_sources @@ -908,7 +898,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: @@ -945,8 +934,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 @@ -990,9 +977,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. @@ -1078,8 +1063,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: @@ -1093,12 +1076,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 +1092,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 @@ -1131,11 +1110,8 @@ 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 + packed = pieces["weight"] # [N, K//2] uint8 block_scales = pieces["weight_scale"] # [N, K//bs] fp8 global_scale = pieces["weight_scale_2"] # [1] fp32 @@ -1154,8 +1130,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") @@ -1168,18 +1142,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): @@ -1188,11 +1158,9 @@ 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 + 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 +1174,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") @@ -1228,17 +1194,10 @@ 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 ( - QLoRAMoeExperts, BlockFP8QLoRAMoeExperts, - ) 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..202714cd 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, @@ -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, @@ -27,6 +27,7 @@ rotate_half, ) + __all__ = [ "ACT2FN", "ATTENTION_FUNCTIONS", @@ -40,6 +41,7 @@ "MoELoRAConfig", "TopKRouter", "RMSNorm", + "get_rmsnorm_mode", "ROPE_INIT_FUNCTIONS", "RotaryEmbedding", "apply_rotary_pos_emb", @@ -49,5 +51,6 @@ "repeat_kv", "rope_config_validation", "rotate_half", + "set_rmsnorm_mode", "update_causal_mask", ] diff --git a/src/xorl/models/layers/activations.py b/src/xorl/models/layers/activations.py index 462a7da1..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__) @@ -225,9 +226,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 @@ -236,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/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..e195a018 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. @@ -32,16 +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) # ------------------------------------------------------------------ # @@ -56,9 +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( @@ -83,15 +84,18 @@ 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 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 +136,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/__init__.py b/src/xorl/models/layers/moe/__init__.py index 8ed8ed71..715812d8 100644 --- a/src/xorl/models/layers/moe/__init__.py +++ b/src/xorl/models/layers/moe/__init__.py @@ -10,18 +10,12 @@ - :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, - copy_weights_to_lora_experts, - mark_only_lora_as_trainable, - lora_state_dict, -) -from .backend import MOE_EXPERT_BACKENDS +from .router import TopKRouter + __all__ = [ "MoEExperts", @@ -34,4 +28,37 @@ "mark_only_lora_as_trainable", "lora_state_dict", "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/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/layers/moe/backend/__init__.py b/src/xorl/models/layers/moe/backend/__init__.py index d51c4c6b..5b8f0796 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) @@ -41,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 @@ -54,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 + + 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, expert_scores) - EP_EXPERT_COMPUTE["quack"] = QuackEPGroupGemm.apply + 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 - EP_EXPERT_COMPUTE["native"] = native_ep_compute + 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, expert_scores) + + EP_EXPERT_COMPUTE["native"] = _native_ep_fused except ImportError: pass @@ -83,7 +95,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 +106,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, ) @@ -172,65 +186,6 @@ 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 -try: - from xorl.ops.moe.quack import QuackEPGroupGemmMoeAct - - EP_EXPERT_COMPUTE_MOE_ACT["quack"] = QuackEPGroupGemmMoeAct.apply -except ImportError: - pass - -# Native EP moe_act compute -try: - from .native import native_ep_compute_moe_act - - EP_EXPERT_COMPUTE_MOE_ACT["native"] = native_ep_compute_moe_act -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", @@ -238,6 +193,4 @@ "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/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..118fb395 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,12 +120,14 @@ def _unpad( # Core grouped-mm expert forward (compiled like torchtitan) # --------------------------------------------------------------------------- + def _run_experts_grouped_mm( 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 using ``torch._grouped_mm``. @@ -156,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( @@ -171,80 +177,16 @@ 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, -) -> 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, - ) - - out = torch._grouped_mm( - h, down_proj.to(torch.bfloat16), offs=offsets, - ).to(x.dtype) - - return out - - # --------------------------------------------------------------------------- # 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 +212,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,12 +246,20 @@ def _native_expert_forward_impl( sorted_hidden_padded[pad_dst] = sorted_hidden # 7. Expert compute (backend-specific) + expert_scores_padded = sorted_hidden_padded.new_zeros(total_padded) + expert_scores_padded[pad_dst] = sorted_weights.to(expert_scores_padded.dtype) + expert_out_padded = compute_fn( - gate_proj, up_proj, down_proj, sorted_hidden_padded, padded_counts + gate_proj, + up_proj, + down_proj, + sorted_hidden_padded, + padded_counts, + expert_scores_padded, ) - # 8. Gather from padded layout (reuse pad_dst) + apply routing weights - expert_out = expert_out_padded[pad_dst] * sorted_weights.unsqueeze(-1) + # 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) @@ -321,6 +272,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 +285,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 +310,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 +342,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 +373,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 +385,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 +400,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 +412,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 +451,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,12 +492,14 @@ 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, 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``. @@ -558,7 +525,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]) @@ -568,6 +545,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,63 +572,18 @@ 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, ) 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, -) -> 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) - out_padded = _run_experts_moe_act(gate_proj, up_proj, down_proj, padded_tokens, padded_counts) - - 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/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 e9a90f6c..86100c44 100644 --- a/src/xorl/models/layers/moe/experts.py +++ b/src/xorl/models/layers/moe/experts.py @@ -1,11 +1,19 @@ """MoE expert weight container with backend dispatch.""" 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, + MOE_EXPERT_BACKENDS, +) +from .common import split_gate_up_proj def _flag_enabled(name: str) -> bool: @@ -25,9 +33,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. @@ -51,29 +61,37 @@ 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.gate_up_proj._fused_gate_up = True self.down_proj = nn.Parameter( torch.empty(num_experts, intermediate_size, hidden_dim), requires_grad=True, ) 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 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, @@ -89,41 +107,37 @@ 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 return fn( hidden_states, expert_idx, - self.gate_proj, - self.up_proj, + self.gate_proj.contiguous(), + self.up_proj.contiguous(), self.down_proj, self.act_fn, ) # 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() 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: - fn = MOE_EXPERT_BACKENDS_MOE_ACT[self.moe_implementation] - else: - fn = MOE_EXPERT_BACKENDS[self.moe_implementation] + # Local single-GPU path + gate_proj = self.gate_proj.contiguous() + up_proj = self.up_proj.contiguous() + fn = MOE_EXPERT_BACKENDS[self.moe_implementation] return fn( hidden_states, routing_weights, selected_experts, - self.gate_proj, - self.up_proj, + gate_proj, + up_proj, self.down_proj, num_experts=self.num_experts, ) @@ -144,7 +158,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_DISPATCH, EP_COMBINE, EP_EXPERT_COMPUTE, EP_EXPERT_COMPUTE_MOE_ACT if self.moe_implementation not in EP_EXPERT_COMPUTE: raise ValueError( @@ -153,28 +166,24 @@ 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] 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 - ) + 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) @@ -182,16 +191,54 @@ 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, - self.gate_proj, self.up_proj, self.down_proj, + permute_tokens, + cumsum, + self.gate_up_proj, + self.down_proj, + self.intermediate_size, + expert_scores, ) # 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): @@ -201,7 +248,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 @@ -214,16 +260,19 @@ 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, self.up_proj, self.down_proj, + permute_tokens, + cumsum, + self.gate_up_proj, + self.down_proj, + self.intermediate_size, + expert_scores, ) 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() @@ -263,13 +312,14 @@ 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, 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 73fa065f..837ad78b 100644 --- a/src/xorl/models/layers/moe/lora.py +++ b/src/xorl/models/layers/moe/lora.py @@ -15,17 +15,20 @@ 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 +from .backend import EP_COMBINE, EP_DISPATCH, EP_EXPERT_COMPUTE_LORA, MOE_EXPERT_BACKENDS_LORA +from .common import split_gate_up_proj + logger = logging.get_logger(__name__) @@ -45,7 +48,6 @@ def __post_init__(self): self.target_modules = ["gate_proj", "up_proj", "down_proj"] - class MoEExpertsLoRA(LoraModule, nn.Module): """MoE experts with LoRA adapters. @@ -53,7 +55,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,14 +83,11 @@ 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.gate_up_proj._fused_gate_up = True self.down_proj = nn.Parameter( torch.empty(self.num_experts, intermediate_size, hidden_dim), requires_grad=False, @@ -126,6 +127,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 +185,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() @@ -192,25 +207,25 @@ 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() 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 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, @@ -233,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_DISPATCH, EP_COMBINE, EP_EXPERT_COMPUTE_LORA if self.moe_implementation not in EP_EXPERT_COMPUTE_LORA: raise ValueError( @@ -242,34 +256,44 @@ 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] 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 - ) + 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, + gate_proj, + 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, ) + # 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 - ) + 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): @@ -283,13 +307,14 @@ 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, 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): @@ -357,7 +382,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") @@ -377,13 +403,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 @@ -407,7 +437,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") @@ -427,22 +458,23 @@ 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 - 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}") # --------------------------------------------------------------------------- @@ -453,8 +485,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/layers/moe/moe_block.py b/src/xorl/models/layers/moe/moe_block.py index b70d79c3..c33acc90 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 @@ -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``. """ @@ -85,7 +85,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: @@ -98,14 +97,13 @@ 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. hybrid_shared: If True, share ``lora_A`` for gate/up and ``lora_B`` for down across experts. """ - from .lora import inject_lora_into_experts + from .lora import inject_lora_into_experts # noqa: PLC0415 inject_lora_into_experts( self, @@ -130,22 +128,24 @@ 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): + 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) @@ -165,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(): @@ -173,44 +172,89 @@ 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 + routing_weights, selected_experts = self.router(router_logits, hidden_states.dtype) + + 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'." ) - - # When train_router is False, detach routing_weights so the gate is - # only trained via auxiliary losses on router_logits, not through - # expert computation. 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) - final_hidden_states = final_hidden_states.reshape( - batch_size, sequence_length, hidden_dim - ) return final_hidden_states, router_logits def _eager_forward( @@ -223,21 +267,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 @@ -252,5 +291,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/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/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/layers/rope.py b/src/xorl/models/layers/rope.py index 41de2fa3..94eafe71 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 @@ -535,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/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 f3878806..6ac45c14 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -1,30 +1,36 @@ - 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 from functools import partial 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 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 -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 +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 @@ -39,13 +45,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, - parse_expert_key, - checkpoint_has_per_expert_weights, -) + +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]]: @@ -95,7 +143,6 @@ def _get_checkpoint_keys(weights_path: str) -> Optional[Set[str]]: return None - @contextmanager def init_empty_weights(): """ @@ -209,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 @@ -246,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 @@ -284,14 +327,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..." @@ -303,57 +346,65 @@ 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) - 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: - 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) + 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) + 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)] 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 @@ -396,6 +447,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. @@ -429,9 +538,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 +564,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 +610,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 @@ -536,7 +637,7 @@ 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: if isinstance(p, DTShard): @@ -596,8 +697,6 @@ 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 # Check if this is a LoRA parameter and handle specially if "lora_A" in name or "lora_B" in name: @@ -664,7 +763,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. @@ -762,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 @@ -777,7 +875,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") @@ -928,7 +1025,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, ) @@ -984,6 +1084,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). @@ -995,23 +1096,28 @@ 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], 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 # 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( @@ -1030,7 +1136,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: @@ -1045,33 +1157,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): " - f"{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) @@ -1086,7 +1226,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() @@ -1096,25 +1239,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) @@ -1124,7 +1289,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) @@ -1335,11 +1503,17 @@ 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"]): + 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`.") @@ -1419,6 +1593,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) \ No newline at end of file + 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/__init__.py b/src/xorl/models/transformers/__init__.py index 6fde8dfc..a4bec082 100644 --- a/src/xorl/models/transformers/__init__.py +++ b/src/xorl/models/transformers/__init__.py @@ -1,12 +1,16 @@ from . import ( + llama3, + qwen2, qwen3, qwen3_5, - qwen3_moe, qwen3_5_moe, + qwen3_moe, ) __all__ = [ + "llama3", + "qwen2", "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 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/models/transformers/qwen3/checkpoint_handler.py b/src/xorl/models/transformers/qwen3/checkpoint_handler.py index 8da80cb2..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 @@ -7,10 +8,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 +92,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 +111,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 [] @@ -151,7 +159,6 @@ def on_load_weight( 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}") @@ -163,9 +170,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..f977e37f 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}`; " @@ -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 @@ -257,7 +258,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..d487c822 100644 --- a/src/xorl/models/transformers/qwen3_5/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3_5/checkpoint_handler.py @@ -1,19 +1,20 @@ """Checkpoint handler for dense Qwen3_5 models.""" +import warnings from typing import Callable, List, Optional, Set, Tuple import torch 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 +56,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 +111,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. @@ -154,8 +155,6 @@ def on_load_weight( 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}") @@ -165,9 +164,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..938db2d1 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, @@ -14,13 +19,17 @@ ) 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.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 +37,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 @@ -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): @@ -185,7 +209,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 +242,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." @@ -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/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..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,17 +1,25 @@ """Checkpoint handler for Qwen3_5 MoE models.""" import re -from typing import Callable, List, Optional, Set, Tuple +import warnings +from typing import Callable, Dict, List, Optional, Set, Tuple import torch 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_full_key, + parse_expert_key, ) from ..qwen3_5_shared import ( is_excluded_module_key, @@ -58,16 +66,15 @@ 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: - self._expert_buffer = ExpertWeightBuffer( - num_experts, ep_rank=ep_rank, ep_size=ep_size - ) + 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: self._qkv_buffer = QKVMergeBuffer() @@ -79,39 +86,83 @@ 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 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() + 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: + 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: + 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: 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,9 +171,33 @@ 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.transpose(1, 2).contiguous())] + + 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.transpose(1, 2), up.transpose(1, 2)], dim=2).contiguous(), + ) + ] + return [] + 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,21 +213,33 @@ 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: + 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: @@ -176,8 +263,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 +284,21 @@ 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]]: + 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. @@ -229,11 +332,7 @@ def on_load_weight( 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: @@ -266,27 +365,23 @@ 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) 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]]: - import warnings if self._expert_buffer is not None: 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: @@ -297,9 +392,26 @@ 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 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/configuration_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/configuration_qwen3_5_moe.py index ec9d36f7..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 @@ -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), @@ -248,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/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..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 @@ -7,17 +7,27 @@ 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.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 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 +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 +35,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 +68,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 +78,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: @@ -87,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): @@ -183,13 +208,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, @@ -198,48 +225,56 @@ 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') - 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) + 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 = { - '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"), } -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 - 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,30 +289,28 @@ 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) - 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) 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, @@ -287,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, @@ -303,30 +336,36 @@ 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 = 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 + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual=residual, prenorm=True) + 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): 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 +393,28 @@ 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) + 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, @@ -380,6 +427,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, ) @@ -392,7 +440,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 +467,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 @@ -465,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, @@ -492,17 +551,20 @@ 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],) 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 +601,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 +635,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_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_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..0acf7a17 100644 --- a/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3_moe/checkpoint_handler.py @@ -1,19 +1,28 @@ """Checkpoint handler for Qwen3 MoE models.""" -from typing import Callable, List, Optional, Set, Tuple +import re +import warnings +from typing import Callable, Dict, List, Optional, Set, Tuple import torch import torch.nn as nn 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, ) @@ -43,6 +52,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, @@ -54,6 +67,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, @@ -61,9 +75,12 @@ 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, 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: @@ -71,9 +88,16 @@ 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._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 @@ -83,9 +107,93 @@ 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 _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: + 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: + 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: + 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.transpose(1, 2).contiguous())] + + 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.transpose(1, 2), up.transpose(1, 2)], dim=2).contiguous(), + ) + ] + return [] + def get_skip_key_fn(self) -> Optional[Callable[[str], bool]]: """Return predicate to skip keys during loading. @@ -96,23 +204,26 @@ 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: + 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 @@ -121,6 +232,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: @@ -152,8 +272,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 +298,17 @@ 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]]: + 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 [] @@ -208,19 +340,20 @@ def on_load_weight( 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 @@ -234,7 +367,7 @@ def on_load_weight( 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 @@ -248,12 +381,10 @@ def on_load_weight( 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]]: + 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: @@ -264,18 +395,16 @@ def on_skip_weight( 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]]: - import warnings if self._expert_buffer is not None: 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: @@ -291,15 +420,31 @@ 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 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/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/models/transformers/qwen3_moe/modeling_qwen3_moe.py b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py index 3d59a221..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,7 +12,7 @@ # 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, Unpack import torch from torch import nn @@ -20,6 +20,12 @@ 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 +33,16 @@ 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.module_utils import MoEGradientCheckpointingLayer 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 @@ -134,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") @@ -154,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") @@ -174,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") @@ -194,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") @@ -210,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) @@ -230,76 +234,35 @@ 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, - ) - hidden_states = residual + hidden_states - - # Fully Connected - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) - - 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 + 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, + output_router_logits=output_router_logits, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + **kwargs, + ) class Qwen3MoePreTrainedModel(XorlPreTrainedModel): @@ -336,9 +299,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, @@ -353,6 +315,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( @@ -365,6 +335,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"), @@ -457,7 +428,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, ) @@ -465,17 +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, @@ -486,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, @@ -500,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],) @@ -514,89 +499,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/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/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..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, @@ -25,10 +30,17 @@ get_pid_mn, make_blocked, ) -from ..utils.pretuned import algo_key_scaled, pretuned 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}, @@ -47,9 +59,11 @@ 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/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/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/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..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 @@ -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 @@ -220,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, + ) 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/__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/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 new file mode 100644 index 00000000..8628a41f --- /dev/null +++ b/src/xorl/ops/loss/grpo_loss.py @@ -0,0 +1,239 @@ +""" +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/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..99a50ea3 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,71 +342,112 @@ 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 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 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..9237c14e 100644 --- a/src/xorl/ops/moe/quack.py +++ b/src/xorl/ops/moe/quack.py @@ -24,112 +24,14 @@ 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): - max_M = permute_tokens.shape[0] - cu_seqlens = cumsum_to_cu_seqlens(cumsum) - - 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 - 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) - ctx.save_for_backward(permute_tokens, cumsum, gate_proj, up_proj, down_proj) - return down_output - - @staticmethod - def backward(ctx, grad_output): - permute_tokens, cumsum, gate_proj, up_proj, down_proj = 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 - - # dgrad FC2 - grad_gated_output = 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_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 - - # 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 - - class QuackMoeExpertsFunction(torch.autograd.Function): """Memory-optimized: separate gate/up GEMMs, recompute cheap intermediates, 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 +61,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 +110,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 +131,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 +147,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 @@ -248,14 +180,22 @@ 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( - ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, - max_M, cu_seqlens, + ctx, + permute_tokens, + cumsum, + gate_proj, + up_proj, + down_proj, + expert_scores, + max_M, + cu_seqlens, ) gate_output = quack_group_gemm_same_nk( @@ -267,6 +207,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( @@ -274,14 +216,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( @@ -296,6 +243,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() @@ -323,21 +272,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 ) @@ -346,18 +302,28 @@ 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_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, ) - 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 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,20 +339,32 @@ 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 - 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): @@ -425,18 +403,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 +453,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 +474,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 +491,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 @@ -541,151 +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/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..d4324bd2 100644 --- a/src/xorl/ops/moe/triton.py +++ b/src/xorl/ops/moe/triton.py @@ -1,184 +1,82 @@ -import os - 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 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 -class TritonEPGroupGemm(torch.autograd.Function): - """EP expert MLP computation using Triton group GEMM. - - 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``. - """ - - @staticmethod - 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, - ) - 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, - ) - - gate_activation = torch.ops.aten.silu(gate_output) - gated_output = gate_activation * up_output - 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, - ) - del gated_output - - ctx.save_for_backward(permute_tokens, cumsum, gate_proj, up_proj, down_proj, gate_output, up_output) - 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 - max_M = grad_output.shape[0] - - # Recompute activation - gate_activation = torch.ops.aten.silu(gate_output) - gated_output = gate_activation * up_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, - ) - - # 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_output, b=grad_output, c=grad_down_proj, - cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, - ) - del gated_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, - ) - - 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, - ) - 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) - 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, - ) - del grad_gate_output - - return ( - grad_permute_tokens, - None, # cumsum - grad_gate_proj, - grad_up_proj, - grad_down_proj, - ) + from xorl.ops.group_gemm.kernel.moe import ( + expert_histogram, + moe_gather, + moe_index_compute, + moe_scatter, + ) -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). +class TritonEPGroupGemm(torch.autograd.Function): + """EP expert MLP with fused gate+up GEMM. Zero-copy weight references. - Memory savings: 2 * tokens * intermediate_size * dtype_size per MoE layer. + 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): + 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( - a=permute_tokens, b=up_proj, - cumsum_M=cumsum, max_M=max_M, transpose_a=False, transpose_b=False, + 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_output) - gated_output = gate_activation * up_output - del gate_activation + # 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) 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 - # 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_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 = ctx.saved_tensors + permute_tokens, cumsum, gate_up_proj, down_proj, gate_up_output, expert_scores = ctx.saved_tensors max_M = grad_output.shape[0] - # 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, - ) - 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, - ) - - # Rest identical to TritonEPGroupGemm.backward - 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) # 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, + 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 @@ -186,56 +84,56 @@ 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_weighted, + b=grad_output, + c=grad_down_proj, + cumsum_K=cumsum, + max_K=max_M, + transpose_a=True, + transpose_b=False, ) - del gated_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_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_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, - ) + grad_gated_output = grad_gated_weighted * expert_scores.unsqueeze(-1) + del grad_gated_weighted - 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 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_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, - cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, + 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_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, ) @@ -268,12 +166,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 +197,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 +213,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 +230,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 +255,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 +267,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 +287,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 +311,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 @@ -443,190 +389,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/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..639a7b79 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 @@ -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: @@ -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..0a6c99e3 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) @@ -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] + [ @@ -86,9 +87,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..fc71fb31 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 @@ -38,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): @@ -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..3a73cbcf 100644 --- a/src/xorl/ops/quack/gemm_sm100.py +++ b/src/xorl/ops/quack/gemm_sm100.py @@ -2,37 +2,39 @@ # 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 +import time 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 triton.testing import do_bench +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 +341,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 +368,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 +459,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 +478,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 +505,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 +515,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 +530,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 +542,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 +586,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 +608,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 +835,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 +848,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 +875,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 +896,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 +950,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 +975,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 +995,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 +1030,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 +1049,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 +1110,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 +1182,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 +1277,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 +1325,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 +1359,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 +1747,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 +1786,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 +1810,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 +1820,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 +1840,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 +1856,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 +1939,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 +1958,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 +2297,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 +2406,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 +2428,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 +2449,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 +2461,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 +2469,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, ) @@ -2650,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 @@ -2659,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 @@ -2688,9 +2578,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 +2635,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..f3c9dd70 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 @@ -227,36 +208,28 @@ 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: # 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 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: # 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/__init__.py b/src/xorl/optim/__init__.py index b8c4f12a..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 build_optimizer +from .signsgd import SignSGD -__all__ = ["build_lr_scheduler", "build_optimizer", "Muon"] +__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 f58c4f47..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,9 +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): """ @@ -56,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 @@ -78,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. @@ -95,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, @@ -139,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: @@ -153,21 +230,31 @@ 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] - # Reshape to 2D for Newton-Schulz, then restore. + # 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: orig_shape = grad_local.shape - grad_local = grad_local.reshape(-1, grad_local.shape[-1]) + 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[-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 @@ -176,14 +263,9 @@ 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( - 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 @@ -191,27 +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 (uses PyTorch's implementation) - 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) + 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 = 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 86a7642d..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,199 +25,11 @@ _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 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: return False + ps = get_parallel_state() if ps.dp_mode != "fsdp2" or not ps.ep_enabled: return False @@ -287,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, *, @@ -304,16 +121,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": @@ -321,8 +148,14 @@ 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) + 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), @@ -330,16 +163,22 @@ 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 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, signsgd, muon." ) @@ -361,12 +200,19 @@ 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 """ 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) @@ -430,15 +276,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] @@ -475,22 +343,36 @@ 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} + "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. 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 +381,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 +416,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, ) @@ -584,8 +474,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) @@ -598,14 +488,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 +515,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/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 ade69244..2e9c889f 100644 --- a/src/xorl/qlora/__init__.py +++ b/src/xorl/qlora/__init__.py @@ -1,19 +1,5 @@ -from .modules.linear import QLoRALinear, prefetch_aqn_noise -from .modules.nvfp4_linear import NvFP4QLoRALinear -from .modules.block_fp8_linear import BlockFP8QLoRALinear -from .modules.nf4_linear import NF4QLoRALinear -from .modules.moe_experts import QLoRAMoeExperts -from .utils import ( - 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, -) -from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules +import importlib + __all__ = [ "QLoRALinear", @@ -30,5 +16,33 @@ "detect_prequantized_nvfp4", "detect_prequantized_block_fp8", "maybe_load_prequantized_qlora", - "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"), +} + + +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 110723f8..78336c9a 100644 --- a/src/xorl/qlora/modules/__init__.py +++ b/src/xorl/qlora/modules/__init__.py @@ -1,8 +1,5 @@ -from .linear import QLoRALinear, prefetch_aqn_noise -from .nvfp4_linear import NvFP4QLoRALinear -from .block_fp8_linear import BlockFP8QLoRALinear -from .nf4_linear import NF4QLoRALinear -from .moe_experts import QLoRAMoeExperts +import importlib + __all__ = [ "QLoRALinear", @@ -12,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/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..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): @@ -54,9 +59,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 +116,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.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 if fmt == "block_fp8": @@ -146,9 +150,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.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 elif quant_format == "nf4": @@ -215,11 +220,11 @@ 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 +255,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( @@ -266,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) @@ -293,7 +296,7 @@ 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() self._quantize_and_store(w) @@ -323,10 +326,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): @@ -419,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 46a9ff61..2edf0e14 100644 --- a/src/xorl/qlora/modules/moe_experts.py +++ b/src/xorl/qlora/modules/moe_experts.py @@ -25,20 +25,23 @@ import math from typing import Optional +import safetensors.torch import torch import torch.nn as nn -import torch.nn.functional as F 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 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 +115,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 +160,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 +306,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 +376,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(). @@ -363,9 +392,6 @@ def load_and_quantize_weights(self, weights_path: str, weight_map: Optional[dict 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: @@ -393,6 +419,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]""" @@ -424,16 +455,15 @@ 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() 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 + 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] @@ -467,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_DISPATCH, EP_COMBINE, 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( @@ -476,8 +506,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 +514,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): @@ -520,7 +549,8 @@ 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, buffer_size_gb=self.deepep_buffer_size_gb, @@ -551,7 +581,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 +659,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 +719,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 +738,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 +779,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 +797,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 +822,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 +861,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 +905,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 +918,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 +994,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..613920f5 100644 --- a/src/xorl/qlora/utils.py +++ b/src/xorl/qlora/utils.py @@ -6,18 +6,27 @@ import json import logging +import math +import math as _math import os -from typing import Collection, Dict, List, Optional, Set, Tuple +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.qlora.modules.linear import QLoRALinear -from xorl.qlora.modules.nvfp4_linear import NvFP4QLoRALinear +from xorl.distributed.parallel_state import get_parallel_state 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 +105,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 +135,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 +172,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 @@ -212,7 +218,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 @@ -227,17 +232,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) - and hasattr(experts, "down_proj") and isinstance(experts.down_proj, nn.Parameter) - and experts.gate_proj.dim() == 3 + isinstance(fused_gate_up, nn.Parameter) + and hasattr(experts, "down_proj") + and isinstance(experts.down_proj, nn.Parameter) + 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 @@ -255,7 +273,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 @@ -266,7 +284,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 +355,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 @@ -377,13 +392,8 @@ def maybe_load_and_quantize_moe_qlora( Returns: 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 +401,6 @@ def maybe_load_and_quantize_moe_qlora( weight_map = None if weights_path: try: - from transformers.utils import cached_file, SAFE_WEIGHTS_INDEX_NAME index_path = cached_file(weights_path, SAFE_WEIGHTS_INDEX_NAME) if index_path: with open(index_path) as f: @@ -412,10 +421,12 @@ 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 +434,25 @@ 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 +460,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 +469,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 +514,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 +528,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). @@ -523,7 +539,9 @@ 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) @@ -539,6 +557,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 +574,6 @@ 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 rank = dist.get_rank() device = torch.device("cuda") @@ -615,10 +630,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) @@ -645,12 +657,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 cached_file, SAFE_WEIGHTS_INDEX_NAME index_path = cached_file(weights_path, SAFE_WEIGHTS_INDEX_NAME) if index_path: with open(index_path) as f: @@ -669,8 +679,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())) @@ -703,7 +712,7 @@ 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(): is_meta = param.device.type == "meta" @@ -721,11 +730,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: @@ -736,7 +746,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) @@ -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..d1c62786 100644 --- a/src/xorl/server/api_server/api_types.py +++ b/src/xorl/server/api_server/api_types.py @@ -4,11 +4,25 @@ 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 +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,44 +402,104 @@ 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).""" + model_id: str = Field(..., description="Model identifier that was unloaded") 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.""" + 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)" ) @@ -401,17 +511,21 @@ 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." - ) - 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") + @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.""" @@ -422,7 +536,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.""" @@ -440,17 +553,27 @@ 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.""" - 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 +584,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") @@ -471,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. @@ -491,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.""" @@ -518,17 +659,27 @@ 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.""" 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 +714,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", @@ -578,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( @@ -609,7 +752,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)") @@ -656,13 +799,14 @@ 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( - 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/api_server/endpoints.py b/src/xorl/server/api_server/endpoints.py index e69d4d5f..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,15 +51,16 @@ 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 + logger = logging.getLogger(__name__) router = APIRouter() - # ============================================================================ # Training Endpoints (Two-Phase Pattern) # ============================================================================ @@ -441,19 +446,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 +543,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 +568,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] = [] @@ -657,7 +650,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) """ @@ -687,7 +680,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. @@ -794,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/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..a3156695 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, @@ -30,12 +33,26 @@ from xorl.server.protocol.api_orchestrator import OrchestratorRequest from xorl.server.protocol.operations import SyncWeightsData + logger = logging.getLogger(__name__) 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. @@ -46,8 +63,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" @@ -62,7 +77,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) @@ -175,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: @@ -187,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, ) @@ -283,7 +267,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}") @@ -306,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, @@ -392,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}") @@ -494,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 @@ -564,7 +558,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 +598,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 +868,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 +903,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/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/src/xorl/server/api_server/weights.py b/src/xorl/server/api_server/weights.py index bc679578..cb69175f 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: @@ -697,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: @@ -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/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 5f6d1b43..ecbba254 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=0, 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..7f4870b2 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__) @@ -45,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 @@ -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=0, + 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..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,13 +43,12 @@ 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 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 +64,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 +85,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. @@ -97,14 +100,14 @@ 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)) 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 +141,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 +179,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)}") @@ -201,6 +202,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, ) @@ -224,7 +226,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 +235,7 @@ def run_orchestrator( # API Server Process # ============================================================================ + def run_api_server( host: str, port: int, @@ -264,13 +267,10 @@ 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( - 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 +287,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 + 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 @@ -345,6 +344,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 +360,61 @@ 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): - # Nested config - flatten it + if "model" in config and isinstance(config["model"], dict): + # 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')) - - # 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)) - - # 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) - - # 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) + # 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 + + # 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 + + # 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.* keys are prefixed with worker_ in ServerArguments + worker_config = config.get("worker", {}) + _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: @@ -436,7 +424,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 +458,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 +482,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 +499,7 @@ def calculate_world_size_from_config(config_path: str) -> int: # Main Launcher # ============================================================================ + class Launcher: """Main launcher for all server components.""" @@ -599,7 +588,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 +637,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: @@ -726,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 @@ -851,9 +842,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...") @@ -866,9 +855,9 @@ 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)) + s.bind(("", 0)) return s.getsockname()[1] def _launch_workers_with_torchrun(self): @@ -897,7 +886,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 +1000,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 +1109,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 +1208,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 +1313,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 +1322,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 +1354,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 +1393,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..1fa1791c 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) @@ -181,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): @@ -205,7 +202,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 +226,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 +238,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 +259,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 +269,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 +281,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 +305,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 +325,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 +350,7 @@ class SequentialPacker(Packer): >>> batches[0]["num_samples"] 2 """ - + def __init__( self, enable_packing: bool = True, @@ -358,7 +360,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 +369,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 +421,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 +447,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 +480,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 +541,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 +560,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 +623,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 +676,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 +709,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 +737,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 +746,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 +757,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 +773,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 +782,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 +798,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 +817,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 +836,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 +847,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 +866,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 +886,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 +908,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 +920,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 +954,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 +1003,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 +1032,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 +1070,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 +1078,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..395ecca2 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) # ============================================================================ @@ -97,11 +99,11 @@ 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" - 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/__init__.py b/src/xorl/server/runner/__init__.py index d6d8b72c..bb13a89f 100644 --- a/src/xorl/server/runner/__init__.py +++ b/src/xorl/server/runner/__init__.py @@ -1,4 +1,17 @@ -from xorl.server.runner.runner_dispatcher import RunnerDispatcher -from xorl.server.runner.model_runner import ModelRunner +"""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 cc22aa91..c2b14a56 100644 --- a/src/xorl/server/runner/adapters/__init__.py +++ b/src/xorl/server/runner/adapters/__init__.py @@ -1,4 +1,17 @@ -from xorl.server.runner.adapters.manager import LoRAAdapterManager -from xorl.server.runner.adapters.adapter_coordinator import AdapterCoordinator +"""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/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..40627367 100644 --- a/src/xorl/server/runner/adapters/manager.py +++ b/src/xorl/server/runner/adapters/manager.py @@ -25,9 +25,13 @@ 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 from torch.distributed._tensor.placement_types import Shard + _HAS_DTENSOR = True except ImportError: _HAS_DTENSOR = False @@ -47,7 +51,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 @@ -275,23 +279,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) @@ -369,7 +374,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 +386,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() @@ -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/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..f411a31d 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__) @@ -110,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. @@ -126,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: @@ -167,8 +194,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 +265,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 +424,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 +462,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 +550,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 +568,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 +665,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 +720,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 +810,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 +920,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..dec1ef9c 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -17,17 +17,13 @@ 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 +31,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 +38,25 @@ policy_loss_function, ) from xorl.optim import build_optimizer -from xorl.trainers.model_builder import build_training_model +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, + resolve_training_model_dtype, +) 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,8 +156,22 @@ 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"}, - "policy_loss": {"labels", "target_tokens", "logprobs", "advantages", "rollout_logprobs"}, + "importance_sampling": { + "labels", + "target_tokens", + "logprobs", + "advantages", + "_original_position_ids", + "rollout_logprobs", + }, + "policy_loss": { + "labels", + "target_tokens", + "logprobs", + "advantages", + "_original_position_ids", + "rollout_logprobs", + }, } def __init__( @@ -226,6 +239,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() @@ -275,8 +292,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) @@ -446,13 +463,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 @@ -464,6 +479,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), @@ -473,7 +489,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"), @@ -492,7 +507,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), @@ -519,7 +534,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 +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( @@ -603,9 +632,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 +682,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 +788,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 +807,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 +832,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 +841,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 +856,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 +896,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,12 +980,19 @@ 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) # Count valid tokens globally global_valid_tokens = self._count_global_valid_tokens(micro_batches) @@ -956,7 +1018,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 +1054,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 +1074,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 +1156,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 @@ -1203,6 +1271,13 @@ 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. + gc.collect() + torch.cuda.empty_cache() + # Validate single-tenant mode for full-weights training self._validate_single_tenant(model_id) @@ -1241,7 +1316,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: @@ -1266,7 +1349,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) @@ -1289,6 +1372,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) @@ -1314,16 +1402,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 +1491,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 +1573,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 +1595,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 +1618,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 +1669,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..de92fad4 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}, " @@ -797,40 +821,47 @@ 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"] 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 +879,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 +914,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 +1009,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 @@ -1201,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") @@ -1225,7 +1257,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..8a36d191 100644 --- a/src/xorl/server/runner/setup.py +++ b/src/xorl/server/runner/setup.py @@ -18,9 +18,10 @@ 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.compile_cache import configure_rank_local_compile_caches from xorl.utils.device import get_nccl_backend @@ -227,7 +228,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() @@ -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 @@ -293,9 +278,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..3d12b4e6 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 @@ -67,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: """ @@ -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..407d95fb 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 @@ -62,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]: """ @@ -171,17 +184,27 @@ 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 " - f"with {num_moe_layers} layers" + f"with {num_moe_layers} layers (tried topk={candidates})" ) return None @@ -244,9 +267,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 +292,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..c2f452ab 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -7,8 +7,11 @@ parameters like batch size, epochs, and optimizer settings. """ -from dataclasses import dataclass, field, asdict -from typing import Optional, Dict, List, Literal, Any +import sys +from dataclasses import dataclass, field +from typing import Any, Dict, List, Literal, Optional + +import yaml @dataclass @@ -44,96 +47,95 @@ 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)."}, + ) + + 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')."} + 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."} + 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( - 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 +144,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,105 +169,78 @@ 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)" + }, ) # ======================================================================== # 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( @@ -291,7 +250,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 +272,63 @@ 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." + }, + ) + + 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." + }, ) # ======================================================================== @@ -321,57 +337,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 +384,68 @@ 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."} - ) - - moe_shared_lora: bool = field( - default=False, - metadata={"help": "Enable shared LoRA for MoE: share LoRA weights across experts"} + 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_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 +453,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 +483,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): @@ -540,11 +522,13 @@ 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, "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, @@ -554,7 +538,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, @@ -588,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, @@ -597,6 +588,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 @@ -606,7 +600,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, @@ -621,7 +614,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 +623,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 +665,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 +674,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 +696,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 +712,11 @@ 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 # 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 +724,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 +736,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 +772,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..bfb6b266 100644 --- a/src/xorl/server/utils/network.py +++ b/src/xorl/server/utils/network.py @@ -8,8 +8,11 @@ import logging import os import socket +import time +from pathlib import Path from typing import Optional + logger = logging.getLogger(__name__) @@ -148,8 +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' """ - import os - 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/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/README.md b/src/xorl/server/weight_sync/README.md index afdc20f2..c20a7cfe 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" @@ -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", @@ -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/__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..dcb8df5c 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", @@ -25,9 +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}. " - 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..8d534f87 100644 --- a/src/xorl/server/weight_sync/backends/base.py +++ b/src/xorl/server/weight_sync/backends/base.py @@ -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,9 +61,10 @@ 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 + master_port: int = 0 group_name: str = "weight_sync_group" buffer_size_mb: int = 1024 device: str = "cuda:0" @@ -123,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. @@ -134,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 6d90c828..ec4f3f65 100644 --- a/src/xorl/server/weight_sync/backends/nccl_broadcast.py +++ b/src/xorl/server/weight_sync/backends/nccl_broadcast.py @@ -11,10 +11,12 @@ the pluggable backend interface. """ +import inspect import logging 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 @@ -22,6 +24,14 @@ 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__) @@ -48,6 +58,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 +67,7 @@ class EndpointInfo: @dataclass class SyncResult: """Result of weight synchronization.""" + success: bool message: str transfer_time: float = 0.0 @@ -83,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", @@ -94,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 @@ -111,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: " @@ -122,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). @@ -132,26 +195,14 @@ def _init_training_process_group(self) -> dist.ProcessGroup: Returns: The initialized process group """ - from torch.distributed import TCPStore, PrefixStore - 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" - # 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}" ) @@ -164,23 +215,13 @@ 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(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" @@ -191,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)} @@ -229,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 @@ -249,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() @@ -320,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 # ======================================================================== @@ -339,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, @@ -453,10 +498,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 +513,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 +536,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 +568,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 } @@ -552,6 +599,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. @@ -565,6 +613,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 @@ -591,19 +640,20 @@ 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, ) 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 +746,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 +785,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 +822,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 +867,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 +876,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 @@ -859,14 +897,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}" - ) + 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/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..b0e1e8c5 100644 --- a/src/xorl/server/weight_sync/handler.py +++ b/src/xorl/server/weight_sync/handler.py @@ -32,6 +32,25 @@ 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__) @@ -70,7 +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 @@ -81,12 +99,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: @@ -99,6 +118,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, ) @@ -122,6 +142,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]: """ @@ -134,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}" @@ -167,14 +187,12 @@ 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.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 @@ -238,7 +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 @@ -254,7 +271,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 +296,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 +309,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 +329,19 @@ 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 +354,20 @@ 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), + weight_version=weight_version if is_last_overall and not moe_contexts else None, ) total_bytes += b total_params += p @@ -362,10 +377,12 @@ 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), + weight_version=weight_version if is_last_overall else None, quantization=quantization, ps=_ps, ) @@ -375,8 +392,10 @@ 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), + weight_version=weight_version if is_last_overall else None, quantization=quantization, ) total_bytes += b @@ -398,7 +417,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,8 +434,10 @@ 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), + weight_version=weight_version if is_last_overall else None, ) total_bytes += b total_params += p @@ -449,13 +472,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 +485,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 +526,10 @@ 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 + child_fsdp_prefixes = set() for mname, mod in fsdp_mod.named_modules(): if isinstance(mod, FSDPModule) and mname != "": @@ -549,7 +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 @@ -567,13 +585,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 +616,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 +632,9 @@ 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 + child_fsdp_prefixes = set() for mname, mod in fsdp_mod.named_modules(): if isinstance(mod, FSDPModule) and mname != "": @@ -634,16 +651,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]. @@ -661,8 +681,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) @@ -671,12 +689,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 @@ -689,6 +709,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, @@ -715,6 +736,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'}, " @@ -798,10 +833,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 +853,14 @@ 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, + weight_version=weight_version, ) total_bytes += b total_params += p @@ -841,6 +883,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]: @@ -910,10 +953,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 +971,14 @@ 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, + weight_version=weight_version, ) total_bytes += b total_params += p @@ -966,7 +1016,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 +1047,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 +1092,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 +1150,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 +1168,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 +1195,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: @@ -1183,19 +1230,12 @@ 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.lora import MoEExpertsLoRA - from xorl.models.layers.moe.experts import MoEExperts - 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(): if isinstance(mod, FSDPModule) and mname != "": @@ -1260,25 +1300,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: @@ -1289,19 +1329,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)) @@ -1319,8 +1353,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 @@ -1332,15 +1365,16 @@ 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. - - 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'). - - This splits fused weights back to individual projections before sending. + """Convert training parameter names to HF/SGLang inference names. + + 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.) """ + config = model.config num_heads = config.num_attention_heads num_kv_heads = getattr(config, "num_key_value_heads", num_heads) @@ -1351,16 +1385,27 @@ 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() - 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)) + 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) half = tensor.shape[0] // 2 gate = tensor[:half].clone() @@ -1369,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( @@ -1376,6 +1425,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. @@ -1390,12 +1440,13 @@ 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) + backend.transfer_bucket( + buffer, + flush_cache=flush_cache, + weight_version=weight_version, + ) return bucket_bytes, len(buffer) @staticmethod @@ -1409,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 = [] @@ -1423,4 +1473,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..39644534 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -11,8 +11,26 @@ 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 +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 + logger = helper.create_logger(__name__) @@ -33,6 +51,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 --- @@ -42,6 +105,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, @@ -52,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, @@ -75,7 +138,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, @@ -97,8 +160,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 @@ -111,12 +172,13 @@ 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, 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, @@ -125,7 +187,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: @@ -167,18 +228,18 @@ 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, ) # ------------------------------------------------------------------ # 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) @@ -203,7 +264,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) @@ -236,7 +300,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: @@ -279,6 +342,7 @@ def build_training_model( # Internal helpers # ====================================================================== + def _inject_qlora( model: nn.Module, *, @@ -295,11 +359,6 @@ 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, - ) is_prequantized = False checkpoint_quant_format = None @@ -318,12 +377,10 @@ 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( - 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. @@ -375,28 +432,23 @@ 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): + 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 " - f"(shared={moe_shared_lora}, hybrid_shared={moe_hybrid_shared_lora})" - ) + + 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: - from xorl.lora.utils import inject_lora_into_model inject_lora_into_model( model, r=lora_rank, @@ -419,53 +471,44 @@ 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( - 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 4360cf5e..dc58df5c 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -10,15 +10,18 @@ import json import os +import socket import time from dataclasses import asdict, dataclass, field from typing import Any, Dict, List, Optional, Tuple 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 @@ -26,34 +29,64 @@ 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 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.models.transformers.qwen3_moe.modeling_qwen3_moe import load_balancing_loss_func 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, + should_skip_generic_param_upcast, +) 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__) +_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 # --------------------------------------------------------------------------- # TrainState — checkpointable training state # --------------------------------------------------------------------------- + @dataclass class TrainState: """Minimal checkpointable state (follows torchtitan's Stateful pattern).""" @@ -82,6 +115,7 @@ def load_state_dict(self, d: Dict[str, Any]) -> None: # Trainer # --------------------------------------------------------------------------- + class Trainer: """Offline training loop. Analogous to torchtitan's Trainer. @@ -94,6 +128,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, @@ -102,7 +139,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) @@ -112,6 +152,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 @@ -120,7 +161,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)) @@ -131,8 +174,36 @@ 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 # noqa: PLC0415 + + 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, @@ -155,12 +226,122 @@ 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 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 # noqa: PLC0415 + + 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(), + } + 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, group=_get_trainer_cpu_group()) + 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 # noqa: PLC0415 + + 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 @@ -185,32 +366,36 @@ 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.""" 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, + 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, + rmsnorm_mode=args.model.rmsnorm_mode, init_device=args.train.init_device, ) self.model_config = self.model.config @@ -232,13 +417,19 @@ 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) 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 if detect_prequantized_nvfp4(args.model.model_path): self.is_prequantized = True @@ -253,7 +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( @@ -285,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 - 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, - ) + is_moe_model = getattr(self.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( + 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: @@ -301,8 +506,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, @@ -314,16 +518,18 @@ 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, 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") + 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 @@ -355,14 +561,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_quantize_qlora, maybe_load_and_quantize_moe_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 +575,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: @@ -417,17 +621,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 @@ -447,11 +640,17 @@ 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( + { + "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.""" @@ -493,7 +692,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, @@ -601,8 +799,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() @@ -643,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) @@ -711,17 +917,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) @@ -825,16 +1035,15 @@ 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} tflops={tflops_per_gpu:.1f}" - ) + 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 "?" @@ -842,19 +1051,31 @@ 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 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, - }) + if ( + args.train.global_rank == 0 + and args.train.use_wandb + and self.state.global_step % args.train.wandb_log_interval == 0 + ): + 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": epoch_fraction, + "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 +1098,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 +1120,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 +1145,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 get_model_state_dict, StateDictOptions 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,26 +1155,25 @@ 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, 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: 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 +1185,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() diff --git a/src/xorl/trainers/training_utils.py b/src/xorl/trainers/training_utils.py index 1fa41ba8..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,14 @@ 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" - ) + 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: @@ -203,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 @@ -262,12 +259,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/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) diff --git a/src/xorl/utils/count_flops.py b/src/xorl/utils/count_flops.py index 67b06660..214779c0 100644 --- a/src/xorl/utils/count_flops.py +++ b/src/xorl/utils/count_flops.py @@ -9,64 +9,42 @@ logger = logging.get_logger(__name__) -def _gc_multipliers( - gc_enabled: bool, - recompute_modules: Optional[List[str]], - moe_checkpoint_method: Optional[str], +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 _grad_ckpt_multipliers( + gradient_checkpointing_enabled: bool, + gradient_checkpointing_method: Optional[str] = None, ) -> 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 QK^T/SV training FLOPs multiplier per attended element 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 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) def get_device_flops(unit="T"): @@ -81,7 +59,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: @@ -113,26 +91,23 @@ 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 # 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. @@ -175,30 +150,27 @@ 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 * 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 - + 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 ) - 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 @@ -220,24 +192,24 @@ 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 * 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 with gradient-checkpoint-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 # 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 @@ -269,31 +241,29 @@ 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 - + 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 @@ -301,8 +271,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 @@ -346,17 +316,25 @@ 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 - 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 @@ -364,8 +342,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 @@ -383,18 +361,15 @@ 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 * 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 - + 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) - 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 @@ -419,17 +394,15 @@ 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 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) @@ -476,9 +449,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 @@ -491,10 +462,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/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 37e803e6..af098e95 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 @@ -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,12 +34,14 @@ 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 - if TYPE_CHECKING: from torch.utils.data import DataLoader from transformers import PretrainedConfig @@ -42,28 +50,31 @@ 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]: - """Return the global (pre-SP-slice) token count for FLOPs accounting. +def _get_global_token_count(micro_batch: Dict[str, "torch.Tensor"]) -> List[int]: + """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: @@ -82,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 @@ -98,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 @@ -292,7 +301,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 @@ -336,10 +345,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) @@ -372,10 +381,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: @@ -539,12 +548,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..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"] @@ -72,5 +71,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/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/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 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..4066f69c 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] @@ -170,12 +173,16 @@ def test_packing_dataset(): args.data.datasets = [] 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], - }) + args.train.ringattn_parallel_size = 1 + args.train.ulysses_parallel_size = 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], + } + ) 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..1421ecd2 100644 --- a/tests/data/prepare/test_shared.py +++ b/tests/data/prepare/test_shared.py @@ -1,19 +1,23 @@ """Tests for xorl.data.prepare.shared module.""" from unittest.mock import Mock, patch + 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 ( - 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, + load_preprocessed_dataset, merge_datasets, + save_preprocessed_dataset, ) -from xorl.arguments import DatasetConfig pytestmark = pytest.mark.cpu @@ -89,7 +93,7 @@ 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 mock_response.headers = {} @@ -116,10 +120,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 @@ -161,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" @@ -178,17 +189,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 +245,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 save_preprocessed_dataset, load_preprocessed_dataset args = Mock() args.data.dataset_prepared_path = str(tmp_path) @@ -236,10 +254,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..6a6d74d1 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 CollatePipeline, DataCollator, FlattenCollator, 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,12 @@ 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 +137,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 +147,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 +160,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 +204,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 +218,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 +229,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 +243,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 +266,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 +304,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 +328,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)) @@ -280,17 +343,37 @@ 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) 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 +384,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..cfa67ef3 100644 --- a/tests/data/test_data_loader_distributed.py +++ b/tests/data/test_data_loader_distributed.py @@ -4,17 +4,20 @@ 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 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 ( - 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 +26,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 --- @@ -84,7 +89,7 @@ def test_partitioning_micro_batches_sp_batch_size_and_drop_last(self, mock_sampl 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) assert any(isinstance(c, TextSequenceShardCollator) for c in internal.data_collators) @@ -111,7 +116,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,30 +131,34 @@ 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 - from xorl.data.constants import IGNORE_INDEX cp_size = 4 seq_len = 128 @@ -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) @@ -250,15 +271,25 @@ def test_pipeline_epoch_consistency_and_packed_sequences(self, mock_ps_collator, assert len(epoch1) == len(epoch2) # --- Packed sequence tests --- - from torch.utils.data import Dataset 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 +300,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 +318,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 +351,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 +390,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..91c88fe7 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 # noqa: PLC0415 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) @@ -102,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) @@ -117,6 +120,7 @@ def run_one_pass( # ── Single test case ────────────────────────────────────────────────────────── + def run_test( num_tokens: int, hidden_dim: int, @@ -131,7 +135,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 +145,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 +164,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 + 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( - "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 +211,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 +219,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 +242,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 +252,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_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_ep_lora_weight_slicing.py b/tests/distributed/test_ep_lora_weight_slicing.py index d6161bdd..8dcef965 100644 --- a/tests/distributed/test_ep_lora_weight_slicing.py +++ b/tests/distributed/test_ep_lora_weight_slicing.py @@ -8,13 +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 +from xorl.models.transformers.qwen3_moe.parallelize import get_ep_plan + pytestmark = [pytest.mark.distributed] @@ -74,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() @@ -119,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 @@ -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_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/distributed/test_parallel_state.py b/tests/distributed/test_parallel_state.py index 5473e140..c4a34aae 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] @@ -20,23 +22,24 @@ 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: 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) @@ -89,30 +92,32 @@ 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 - @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..ca310b7c 100644 --- a/tests/distributed/test_tensor_parallel.py +++ b/tests/distributed/test_tensor_parallel.py @@ -8,6 +8,10 @@ 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] @@ -16,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, @@ -50,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 Qwen3MLP, Qwen3ForCausalLM config = Qwen3Config( hidden_size=256, @@ -87,6 +87,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..d45f7289 100644 --- a/tests/e2e/e2e_utils.py +++ b/tests/e2e/e2e_utils.py @@ -15,14 +15,22 @@ 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 +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from transformers import AutoConfig, AutoModelForCausalLM -import pytest + +try: + import xorl_client +except ModuleNotFoundError: + xorl_client = None # --------------------------------------------------------------------------- @@ -36,6 +44,7 @@ # GPU helpers # --------------------------------------------------------------------------- + def gpu_count() -> int: if not torch.cuda.is_available(): return 0 @@ -65,19 +74,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 +138,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 +155,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}" # --------------------------------------------------------------------------- @@ -251,17 +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 @@ -284,6 +278,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 +390,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 +422,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 +441,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 +462,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 +505,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, @@ -530,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_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..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 @@ -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,8 @@ 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] @@ -38,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 @@ -58,6 +67,7 @@ # Helpers # --------------------------------------------------------------------------- + def _skip_if_no_qwen3_8b(): return pytest.mark.skipif( not os.path.isdir(QWEN3_8B_DIR), @@ -67,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() @@ -89,13 +97,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 +132,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,35 +157,30 @@ 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.""" @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["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) @_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["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) @_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["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..7d57ebb6 100644 --- a/tests/e2e/server_utils.py +++ b/tests/e2e/server_utils.py @@ -14,6 +14,19 @@ import requests 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, @@ -60,8 +73,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 +153,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 +218,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 +296,7 @@ def get_log(self) -> str: # Random data generation # --------------------------------------------------------------------------- + def generate_random_sft_data( num_samples: int, seq_len: int = 64, @@ -282,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 = [] @@ -340,9 +366,10 @@ 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 + xorl_client = _require_xorl_client() adam_params = xorl_client.AdamParams(learning_rate=lr, beta1=0.9, beta2=0.95, eps=1e-8) losses = [] @@ -376,15 +403,12 @@ 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): """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( @@ -401,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( @@ -415,4 +439,3 @@ def _create_full_weight_client(base_url, model_dir): # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- - diff --git a/tests/models/test_module_utils_broadcast.py b/tests/models/test_module_utils_broadcast.py new file mode 100644 index 00000000..bb4681c5 --- /dev/null +++ b/tests/models/test_module_utils_broadcast.py @@ -0,0 +1,173 @@ +from types import SimpleNamespace + +import pytest +import torch + +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, 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, 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", + 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_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, device=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_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) + + 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"] diff --git a/tests/models/test_moe_experts_lora.py b/tests/models/test_moe_experts_lora.py index 0355242f..b88a9c76 100644 --- a/tests/models/test_moe_experts_lora.py +++ b/tests/models/test_moe_experts_lora.py @@ -2,20 +2,24 @@ import pytest +from xorl.lora import LoraLinear, inject_lora_into_model + + 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 +200,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) @@ -214,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) @@ -261,8 +287,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 +317,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 +378,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 +406,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 +421,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 +442,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 +473,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 +491,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 +510,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 +526,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 +549,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}" # --------------------------------------------------------------------------- @@ -501,7 +581,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): @@ -549,13 +628,13 @@ 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 # 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 +643,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 +651,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,11 +661,176 @@ 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) 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"]) 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..c996d57e --- /dev/null +++ b/tests/models/test_moe_fused_gate_up_proj.py @@ -0,0 +1,178 @@ +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]) + + +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)) == [] diff --git a/tests/models/test_moe_routing_cache.py b/tests/models/test_moe_routing_cache.py deleted file mode 100644 index 7420663f..00000000 --- a/tests/models/test_moe_routing_cache.py +++ /dev/null @@ -1,423 +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 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, - ) - 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 - 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 - 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.""" - import xorl.models.layers.moe.moe_block as mb - - class _FakeConfig: - pass - - class _MoEModel(nn.Module): - def __init__(self): - super().__init__() - 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) - 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/models/test_moe_routing_replay.py b/tests/models/test_moe_routing_replay.py index 21bb9445..a5fa9c74 100644 --- a/tests/models/test_moe_routing_replay.py +++ b/tests/models/test_moe_routing_replay.py @@ -10,18 +10,19 @@ import pytest import torch import torch.nn as nn -from functools import partial from torch.utils.checkpoint import checkpoint +from xorl.models.base import XorlPreTrainedModel + + 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 +213,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 +253,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 +270,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 +283,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 +299,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 +312,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 +332,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 +363,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: @@ -442,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 @@ -498,8 +512,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_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 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..6b00a675 100644 --- a/tests/models/test_moe_weight_loading_integration.py +++ b/tests/models/test_moe_weight_loading_integration.py @@ -5,33 +5,29 @@ into a model that expects fused (stacked) expert format. """ -import os -import tempfile +import random 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 +58,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 +74,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 +98,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 +116,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 +161,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 +240,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 +264,12 @@ 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) - - import random + 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) + items = list(state_dict2.items()) random.seed(42) random.shuffle(items) @@ -293,8 +277,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 +288,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 +302,17 @@ 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 +329,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_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) diff --git a/tests/models/test_qwen3_moe_fused_lora.py b/tests/models/test_qwen3_moe_fused_lora.py index 1becdb78..f8e109f4 100644 --- a/tests/models/test_qwen3_moe_fused_lora.py +++ b/tests/models/test_qwen3_moe_fused_lora.py @@ -4,16 +4,18 @@ 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 ( - 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 +180,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 +247,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 +276,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 +332,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 +360,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 +374,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 +393,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 +439,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 +456,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 +473,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})", ) @@ -461,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): @@ -506,12 +561,12 @@ 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 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 +574,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 +582,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 +591,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/__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..e8d0ca98 --- /dev/null +++ b/tests/ops/loss/test_drgrpo_loss.py @@ -0,0 +1,339 @@ +# 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 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). + """ + 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" diff --git a/tests/ops/test_attention.py b/tests/ops/test_attention.py index 47614ca5..0ab07c5d 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,38 +159,23 @@ 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 - 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 - 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.""" @@ -177,8 +193,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 +208,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..758aea12 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.experts import MoEExperts # noqa: PLC0415 + from xorl.models.layers.moe.moe_block import MoEBlock # noqa: PLC0415 + 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.""" @@ -136,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) @@ -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_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()" + ) diff --git a/tests/ops/test_ep_routing_scores.py b/tests/ops/test_ep_routing_scores.py new file mode 100644 index 00000000..a4a1b6ad --- /dev/null +++ b/tests/ops/test_ep_routing_scores.py @@ -0,0 +1,191 @@ +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.quack", "QuackEPGroupGemm", id="quack"), + ], +) +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) + + # 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() + + 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) diff --git a/tests/ops/test_group_gemm.py b/tests/ops/test_group_gemm.py index 2d632d85..f8b3fcde 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] @@ -84,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") @@ -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)) @@ -150,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") @@ -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() @@ -212,14 +220,14 @@ 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") 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..4616740d 100644 --- a/tests/ops/test_moe_act.py +++ b/tests/ops/test_moe_act.py @@ -17,6 +17,10 @@ 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 @@ -25,10 +29,17 @@ # 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()) + try: + from xorl.models.layers.moe.backend import MOE_EXPERT_BACKENDS_MOE_ACT # noqa: PLC0415 + + 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 [] @@ -39,7 +50,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) @@ -67,11 +78,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 +105,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 +118,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 +143,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 +155,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 +174,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 +213,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 +229,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 +280,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 +308,12 @@ 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.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 + from xorl.models.layers.moe.experts import MoEExperts # noqa: PLC0415 config = Qwen3MoeConfig( vocab_size=1000, @@ -319,11 +336,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 +363,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 +388,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..8e4e547d 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) @@ -126,7 +138,8 @@ 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): 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) @@ -192,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") @@ -212,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) @@ -223,27 +241,46 @@ 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) # --- 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, 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 # noqa: PLC0415 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..aa5986db 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 @@ -53,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") @@ -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 # noqa: PLC0415 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, ) @@ -187,8 +190,11 @@ 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 ( - expert_histogram, moe_scatter, moe_gather, moe_index_compute, + from xorl.ops.group_gemm.kernel.moe import ( # noqa: PLC0415 + 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..df0f88e8 100644 --- a/tests/ops/test_moe_torch_compile.py +++ b/tests/ops/test_moe_torch_compile.py @@ -13,17 +13,23 @@ """ import pytest -import time 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 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,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) @@ -59,7 +65,8 @@ 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, num_experts=num_experts, @@ -78,13 +85,13 @@ 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") 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 +105,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,11 +156,11 @@ 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): """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"]: @@ -181,20 +189,21 @@ 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, - ) config = _tiny_moe_config(_moe_implementation=moe_backend) model = Qwen3MoeForCausalLM(config).to(DEVICE, DTYPE) @@ -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() @@ -486,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 @@ -497,18 +530,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 +571,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 +606,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/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/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/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_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/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..2808571e 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) @@ -75,7 +82,7 @@ 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: w = torch.randn(m.out_features, m.in_features, device="cuda", dtype=torch.bfloat16) @@ -93,13 +100,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 +136,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 +145,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 +174,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 +196,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 +221,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 +240,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 +257,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 +275,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 +304,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 +325,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 +353,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 +375,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 +392,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 +499,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 +541,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 +574,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 +606,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 +630,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 +647,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 +663,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 +686,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 +707,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 +740,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..a3d46897 100644 --- a/tests/server/api_server/test_api_server.py +++ b/tests/server/api_server/test_api_server.py @@ -5,20 +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.server import APIServer from xorl.server.api_server.api_types import ( + AdamParams, + CreateSessionRequest, Datum, DatumInput, ForwardBackwardRequest, + LoadWeightsRequest, OptimStepRequest, SaveWeightsRequest, - LoadWeightsRequest, - AdamParams, + SessionHeartbeatRequest, + UntypedAPIFuture, ) +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: @@ -69,10 +77,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" @@ -102,5 +112,72 @@ def test_request_creation_and_serialization(self): 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 451aca10..098b4ab2 100644 --- a/tests/server/api_server/test_api_types.py +++ b/tests/server/api_server/test_api_types.py @@ -5,31 +5,32 @@ """ 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, - LossFnOutput, - AdamParams, + ErrorResponse, ForwardBackwardRequest, ForwardBackwardResponse, - OptimStepRequest, - OptimStepResponse, - SaveWeightsRequest, - SaveWeightsResponse, + HealthCheckResponse, LoadWeightsRequest, LoadWeightsResponse, + LossFnOutput, + OptimStepRequest, + OptimStepResponse, SaveWeightsForSamplerRequest, SaveWeightsForSamplerResponse, - HealthCheckResponse, - ErrorResponse, + SaveWeightsRequest, + SaveWeightsResponse, ) +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + class TestDatumAndForwardBackward: """Test Datum, DatumInput, ForwardBackwardRequest and ForwardBackwardResponse.""" @@ -89,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=[ @@ -112,7 +121,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 +148,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 +162,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={}) @@ -162,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 @@ -175,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" @@ -201,12 +227,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 +274,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..d64dc5d2 100644 --- a/tests/server/api_server/test_checkpoint_paths.py +++ b/tests/server/api_server/test_checkpoint_paths.py @@ -12,14 +12,17 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException +from fastapi.exceptions import HTTPException -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 +from xorl.server.api_server.utils import validate_model_id class TestTomiUriAndPathConstruction: @@ -43,9 +46,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" @@ -88,16 +89,19 @@ 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() # --- 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 +119,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 +169,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 +182,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 +190,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" @@ -196,7 +212,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" @@ -217,15 +232,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: @@ -256,7 +277,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"]: @@ -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 @@ -295,7 +324,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) @@ -315,13 +343,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..6a066c76 100644 --- a/tests/server/api_server/test_future_store.py +++ b/tests/server/api_server/test_future_store.py @@ -3,20 +3,43 @@ """ import asyncio -import pytest import time +import pytest + + pytestmark = [pytest.mark.cpu, pytest.mark.server] +from typing import Any, Dict, Optional + from xorl.server.api_server.future_store import ( - FutureStore, - FutureStatus, FutureEntry, - make_try_again_response, - make_failed_response, + FutureStatus, + FutureStore, ) +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.""" @@ -47,22 +70,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 +102,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 +131,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 +145,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 +154,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 +172,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 +206,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 +219,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 +246,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 +263,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 +282,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/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/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..92769e4f 100644 --- a/tests/server/orchestrator/test_cu_seqlens_alignment.py +++ b/tests/server/orchestrator/test_cu_seqlens_alignment.py @@ -15,17 +15,21 @@ """ 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.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 + pytestmark = [pytest.mark.cpu] @@ -33,6 +37,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 +94,7 @@ def _make_feature(input_ids, position_ids): # Tests # --------------------------------------------------------------------------- + class TestCuSeqlensAlignmentAndDtype: """Verify server/CLI cu_seq_lens match, dtype, and edge cases.""" @@ -103,10 +109,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"] @@ -155,11 +163,12 @@ 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 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 +186,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 +204,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 +218,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 +227,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,34 +271,39 @@ 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 + 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 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) 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 @@ -291,7 +315,7 @@ 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) datum_list = [_make_datum([10, 20, 30]), _make_datum([40, 50, 60, 70])] @@ -302,7 +326,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..f2f5ddb3 100644 --- a/tests/server/orchestrator/test_orchestrator.py +++ b/tests/server/orchestrator/test_orchestrator.py @@ -13,41 +13,41 @@ - Tests focus on Orchestrator's scheduling, routing, and output formatting """ +import msgpack 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 +105,10 @@ 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..29b71413 100644 --- a/tests/server/orchestrator/test_orchestrator_client_communication.py +++ b/tests/server/orchestrator/test_orchestrator_client_communication.py @@ -9,21 +9,21 @@ import asyncio import logging +import socket as sock + 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 +36,7 @@ # Mock Engine # ============================================================================ + class MockEngine: """ Mock Engine that simulates the INPUT/OUTPUT socket behavior of Orchestrator. @@ -116,33 +117,51 @@ 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 +185,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 +202,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 +228,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 +294,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) @@ -285,10 +314,10 @@ 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: - 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_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 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/test_server_arguments.py b/tests/server/test_server_arguments.py new file mode 100644 index 00000000..1f63a119 --- /dev/null +++ b/tests/server/test_server_arguments.py @@ -0,0 +1,71 @@ +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" + + +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/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" 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 diff --git a/tests/server/weight_sync/test_pp_nccl_transfer.py b/tests/server/weight_sync/test_pp_nccl_transfer.py index db7854d3..b757592a 100644 --- a/tests/server/weight_sync/test_pp_nccl_transfer.py +++ b/tests/server/weight_sync/test_pp_nccl_transfer.py @@ -5,11 +5,11 @@ 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 +from xorl.server.weight_sync.handler import WeightSyncHandler, _prod class TestProd: @@ -37,7 +37,7 @@ 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 @@ -80,6 +80,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") 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", + ) diff --git a/tests/test_arguments.py b/tests/test_arguments.py new file mode 100644 index 00000000..07d32bc6 --- /dev/null +++ b/tests/test_arguments.py @@ -0,0 +1,93 @@ +import sys + +import pytest +import torch +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 == {} + + +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 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 + ) 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 + )